Implement broadcast, memory_opt! performance improvements, additional tests

Add a couple allocation test cases

Improvements to memory_opt, change jl_typeof attribute, always-copy on BoundsError

Whitespace, fix copy in BoundsError contstructor

Author: Ian Atol

base/abstractarray.jl

@@ -1073,6 +1073,10 @@ function copy(a::AbstractArray)
     copymutable(a)
 end
 
+function copy(a::Core.ImmutableArray)
+    a
+end
+
 function copyto!(B::AbstractVecOrMat{R}, ir_dest::AbstractRange{Int}, jr_dest::AbstractRange{Int},
                A::AbstractVecOrMat{S}, ir_src::AbstractRange{Int}, jr_src::AbstractRange{Int}) where {R,S}
     if length(ir_dest) != length(ir_src)

base/array.jl

@@ -118,6 +118,35 @@ Union type of [`DenseVector{T}`](@ref) and [`DenseMatrix{T}`](@ref).
 """
 const DenseVecOrMat{T} = Union{DenseVector{T}, DenseMatrix{T}}
 
+"""
+    ImmutableArray
+
+Dynamically allocated, immutable array.
+
+"""
+const ImmutableArray = Core.ImmutableArray
+
+"""
+    IMArray{T,N}
+
+Union type of [`Array{T,N}`](@ref) and [`ImmutableArray{T,N}`](@ref)
+"""
+const IMArray{T,N} = Union{Array{T, N}, ImmutableArray{T,N}}
+
+"""
+    IMVector{T}
+
+One-dimensional [`ImmutableArray`](@ref) or [`Array`](@ref) with elements of type `T`. Alias for `IMArray{T, 1}`.
+"""
+const IMVector{T} = IMArray{T, 1}
+
+"""
+    IMMatrix{T}
+
+Two-dimensional [`ImmutableArray`](@ref) or [`Array`](@ref) with elements of type `T`. Alias for `IMArray{T,2}`.
+"""
+const IMMatrix{T} = IMArray{T, 2}
+
 ## Basic functions ##
 
 import Core: arraysize, arrayset, arrayref, const_arrayref
@@ -147,14 +176,13 @@ function vect(X...)
     return copyto!(Vector{T}(undef, length(X)), X)
 end
 
-const ImmutableArray = Core.ImmutableArray
-const IMArray{T,N} = Union{Array{T, N}, ImmutableArray{T,N}}
-const IMVector{T} = IMArray{T, 1}
-const IMMatrix{T} = IMArray{T, 2}
-
+# Freeze and thaw constructors
 ImmutableArray(a::Array) = Core.arrayfreeze(a)
 Array(a::ImmutableArray) = Core.arraythaw(a)
 
+ImmutableArray(a::AbstractArray{T,N}) where {T,N} = ImmutableArray{T,N}(a)
+
+# Size functions for arrays, both mutable and immutable
 size(a::IMArray, d::Integer) = arraysize(a, convert(Int, d))
 size(a::IMVector) = (arraysize(a,1),)
 size(a::IMMatrix) = (arraysize(a,1), arraysize(a,2))
@@ -393,6 +421,9 @@ similar(a::Array{T}, m::Int) where {T}              = Vector{T}(undef, m)
 similar(a::Array, T::Type, dims::Dims{N}) where {N} = Array{T,N}(undef, dims)
 similar(a::Array{T}, dims::Dims{N}) where {T,N}     = Array{T,N}(undef, dims)
 
+ImmutableArray{T}(undef::UndefInitializer, m::Int) where T = ImmutableArray(Array{T}(undef, m))
+ImmutableArray{T}(undef::UndefInitializer, dims::Dims) where T = ImmutableArray(Array{T}(undef, dims))
+
 # T[x...] constructs Array{T,1}
 """
     getindex(type[, elements...])
@@ -626,8 +657,8 @@ oneunit(x::AbstractMatrix{T}) where {T} = _one(oneunit(T), x)
 
 ## Conversions ##
 
-convert(::Type{T}, a::AbstractArray) where {T<:Array} = a isa T ? a : T(a)
 convert(::Type{Union{}}, a::AbstractArray) = throw(MethodError(convert, (Union{}, a)))
+convert(T::Union{Type{<:Array},Type{<:Core.ImmutableArray}}, a::AbstractArray) = a isa T ? a : T(a)
 
 promote_rule(a::Type{Array{T,n}}, b::Type{Array{S,n}}) where {T,n,S} = el_same(promote_type(T,S), a, b)
 
@@ -637,6 +668,7 @@ if nameof(@__MODULE__) === :Base  # avoid method overwrite
 # constructors should make copies
 Array{T,N}(x::AbstractArray{S,N})         where {T,N,S} = copyto_axcheck!(Array{T,N}(undef, size(x)), x)
 AbstractArray{T,N}(A::AbstractArray{S,N}) where {T,N,S} = copyto_axcheck!(similar(A,T), A)
+ImmutableArray{T,N}(Ar::AbstractArray{S,N}) where {T,N,S} = Core.arrayfreeze(copyto_axcheck!(Array{T,N}(undef, size(Ar)), Ar))
 end
 
 ## copying iterators to containers

base/boot.jl

@@ -279,9 +279,21 @@ struct BoundsError <: Exception
     a::Any
     i::Any
     BoundsError() = new()
-    BoundsError(@nospecialize(a)) = (@_noinline_meta; new(a))
-    BoundsError(@nospecialize(a), i) = (@_noinline_meta; new(a,i))
+    # For now, always* copy to avoid escaping a
+    # Eventually, we want to figure out if the copy is needed to save the performance of copying
+
+    # * to avoid throwing a MethodError and breaking many test suites, first check if copy is defined for a
+
+    #BoundsError(@nospecialize(a)) = (@_noinline_meta; new(a))
+    BoundsError(@nospecialize(a)) = Main.Base.hasmethod(Main.Base.copy, Tuple{typeof(a)}) ?
+                                    (@_noinline_meta; new(Main.Base.copy(a))) :
+                                    (@_noinline_meta; new(a))
+    #BoundsError(@nospecialize(a), i) = (@_noinline_meta; new(a, i))
+    BoundsError(@nospecialize(a), i) = Main.Base.hasmethod(Main.Base.copy, Tuple{typeof(a)}) ?
+                                       (@_noinline_meta; new(Main.Base.copy(a), i)) :
+                                       (@_noinline_meta; new(a, i))
 end
+
 struct DivideError         <: Exception end
 struct OutOfMemoryError    <: Exception end
 struct ReadOnlyMemoryError <: Exception end

base/broadcast.jl

@@ -1385,4 +1385,17 @@ function Base.show(io::IO, op::BroadcastFunction)
 end
 Base.show(io::IO, ::MIME"text/plain", op::BroadcastFunction) = show(io, op)
 
+struct IMArrayStyle <: Broadcast.AbstractArrayStyle{Any} end
+BroadcastStyle(::Type{<:Core.ImmutableArray}) = IMArrayStyle()
+
+#similar has to return mutable array
+function Base.similar(bc::Broadcasted{IMArrayStyle}, ::Type{ElType}) where ElType
+    similar(Array{ElType}, axes(bc))
+end
+
+@inline function copy(bc::Broadcasted{IMArrayStyle})
+    ElType = combine_eltypes(bc.f, bc.args)
+    return Core.ImmutableArray(copyto!(similar(bc, ElType), bc))
+end
+
 end # module

base/compiler/optimize.jl

@@ -307,6 +307,7 @@ function run_passes(ci::CodeInfo, sv::OptimizationState)
     ir = adce_pass!(ir)
     #@Base.show ("after_adce", ir)
     @timeit "type lift" ir = type_lift_pass!(ir)
+    #@timeit "compact 3" ir = compact!(ir)
     ir = memory_opt!(ir)
     #@Base.show ir
     if JLOptions().debug_level == 2

base/compiler/ssair/passes.jl

@@ -1256,23 +1256,47 @@ function cfg_simplify!(ir::IRCode)
     return finish(compact)
 end
 
-function is_allocation(stmt)
+# function is_known_fcall(stmt::Expr, @nospecialize(func))
+#     isexpr(stmt, :foreigncall) || return false
+#     s = stmt.args[1]
+#     isa(s, QuoteNode) && (s = s.value)
+#     return s === func
+# end
+
+function is_known_fcall(stmt::Expr, funcs::Vector{Symbol})
     isexpr(stmt, :foreigncall) || return false
     s = stmt.args[1]
     isa(s, QuoteNode) && (s = s.value)
-    return s === :jl_alloc_array_1d
+    # return any(e -> s === e, funcs)
+    return true in map(e -> s === e, funcs)
+end
+
+function is_allocation(stmt::Expr)
+    isexpr(stmt, :foreigncall) || return false
+    s = stmt.args[1]
+    isa(s, QuoteNode) && (s = s.value)
+    return (s === :jl_alloc_array_1d
+         || s === :jl_alloc_array_2d
+         || s ===  :jl_alloc_array_3d
+         || s === :jl_new_array)
 end
 
 function memory_opt!(ir::IRCode)
     compact = IncrementalCompact(ir, false)
     uses = IdDict{Int, Vector{Int}}()
-    relevant = IdSet{Int}()
-    revisit = Int[]
-    function mark_val(val)
+    relevant = IdSet{Int}() # allocations
+    revisit = Int[] # potential targets for a mutating_arrayfreeze drop-in
+
+    function mark_escape(@nospecialize val)
         isa(val, SSAValue) || return
+        #println(val.id, " escaped.")
         val.id in relevant && pop!(relevant, val.id)
     end
+
     for ((_, idx), stmt) in compact
+
+        #println("idx: ", idx, " = ", stmt)
+
         if isa(stmt, ReturnNode)
             isdefined(stmt, :val) || continue
             val = stmt.val
@@ -1281,51 +1305,171 @@ function memory_opt!(ir::IRCode)
                 push!(uses[val.id], idx)
             end
             continue
+
+        # check for phinodes that are possibly allocations
+        elseif isa(stmt, PhiNode)
+
+            # this loop seems like a waste, but using map here didn't go well
+            defined = true
+            for i = 1:length(stmt.values)
+                if !isassigned(stmt.values, i)
+                    defined = false
+                end
+            end
+
+            defined || continue
+
+            for val in stmt.values
+                if isa(val, SSAValue) && val.id in relevant
+                    #println("Adding ", idx ," to relevant from PhiNode: " , stmt)
+                    push!(relevant, idx)
+                end
+            end
         end
+
         (isexpr(stmt, :call) || isexpr(stmt, :foreigncall)) || continue
+
         if is_allocation(stmt)
             push!(relevant, idx)
             # TODO: Mark everything else here
             continue
         end
-        # TODO: Replace this by interprocedural escape analysis
-        if is_known_call(stmt, arrayset, compact)
+
+        if is_known_call(stmt, arrayset, compact) && length(stmt.args) >= 5
             # The value being set escapes, everything else doesn't
-            mark_val(stmt.args[4])
+            mark_escape(stmt.args[4])
+
+            arr = stmt.args[3]
+
+            if isa(arr, SSAValue) && arr.id in relevant
+                (haskey(uses, arr.id)) || (uses[arr.id] = Int[])
+                push!(uses[arr.id], idx)
+            end
+
+        elseif is_known_call(stmt, arrayref, compact) && length(stmt.args) == 4
             arr = stmt.args[3]
+
+            if isa(arr, SSAValue) && arr.id in relevant
+                (haskey(uses, arr.id)) || (uses[arr.id] = Int[])
+                push!(uses[arr.id], idx)
+            end
+
+        elseif is_known_call(stmt, setindex!, compact) && length(stmt.args) == 4
+            #println("setindex!: ", stmt.args)
+            # handle similarly to arrayset
+            # escape the value being set
+            val = stmt.args[3]
+            mark_escape(val)
+            # track usage of arr for dominance analysis
+            arr = stmt.args[2]
+            if isa(arr, SSAValue) && arr.id in relevant
+                (haskey(uses, arr.id)) || (uses[arr.id] = Int[])
+                push!(uses[arr.id], idx)
+            end
+
+        # these foreigncalls have similar structure and don't escape our array, so handle them all at once
+        elseif is_known_fcall(stmt, [:jl_array_ptr, :jl_array_copy]) && length(stmt.args) == 6
+            #println("is_known_fcall: ", stmt)
+            arr = stmt.args[6]
+
+            # just record usage info
             if isa(arr, SSAValue) && arr.id in relevant
                 (haskey(uses, arr.id)) || (uses[arr.id] = Int[])
                 push!(uses[arr.id], idx)
             end
+
+        # elseif is_known_fcall(stmt, :jl_array_ptr) && length(stmt.args) == 6
+        #     arr = stmt.args[6]
+
+        #     if isa(arr, SSAValue) && arr.id in relevant
+        #         (haskey(uses, arr.id)) || (uses[arr.id] = Int[])
+        #         push!(uses[arr.id], idx)
+        #     end
+
+        # elseif is_known_fcall(stmt, :jl_array_copy) && length(stmt.args) == 6
+        #     arr = stmt.args[6]
+
+        #     if isa(arr, SSAValue) && arr.id in relevant
+        #         (haskey(uses, arr.id)) || (uses[arr.id] = Int[])
+        #         push!(uses[arr.id], idx)
+        #     end
+
+        elseif is_known_call(stmt, arraysize, compact) && isa(stmt.args[2], SSAValue) #&& isa(stmt.args[3], Number)
+            arr = stmt.args[2]
+            # dim = stmt.args[3]
+            # typ = types(compact)[arr]
+
+            # if isa(typ, Core.Const)
+            #     typ = typ.val
+            # end
+
+            # NEW: since exceptions no longer escape arrays, we can just assume no escape
+
+            if arr.id in relevant
+                (haskey(uses, arr.id)) || (uses[arr.id] = Int[])
+                push!(uses[arr.id], idx)
+            end
+
+            # # make sure this call isn't going to throw
+            # if isa(typ, Type) && typ <: AbstractArray && dim >= 1
+            #     # don't escape the array, but mark usage for dom analysis
+            #     if arr.id in relevant
+            #         (haskey(uses, arr.id)) || (uses[arr.id] = Int[])
+            #         push!(uses[arr.id], idx)
+            #     end
+
+            # else # if this call throws or we can't tell, assume all uses escape
+            #     for ur in userefs(stmt)
+            #         mark_escape(ur[])
+            #     end
+            # end
+
         elseif is_known_call(stmt, Core.arrayfreeze, compact) && isa(stmt.args[2], SSAValue)
+            # mark these for potential replacement with mutating_arrayfreeze
             push!(revisit, idx)
+
         else
-            # For now we assume everything escapes
-            # TODO: We could handle PhiNodes specially and improve this
+            # Assume everything else escapes
             for ur in userefs(stmt)
-                mark_val(ur[])
+                mark_escape(ur[])
             end
         end
     end
+
     ir = finish(compact)
     isempty(revisit) && return ir
+
     domtree = construct_domtree(ir.cfg.blocks)
+
     for idx in revisit
         # Make sure that the value we reference didn't escape
-        id = ir.stmts[idx][:inst].args[2].id
+        stmt = ir.stmts[idx][:inst]::Expr
+        id = (stmt.args[2]::SSAValue).id
+        # print("Relevant: ")
+        # for id in relevant
+        #     print(id, " ")
+        # end
+        # println(" ")
+        # print("Revisit: ")
+        # Main.Base.show(revisit)
+        # println()
         (id in relevant) || continue
 
+        #println("Revisiting ", stmt)
+
         # We're ok to steal the memory if we don't dominate any uses
         ok = true
-        for use in uses[id]
-            if ssadominates(ir, domtree, idx, use)
-                ok = false
-                break
+        if haskey(uses, id)
+            for use in uses[id]
+                if ssadominates(ir, domtree, idx, use)
+                    ok = false
+                    break
+                end
             end
         end
         ok || continue
-
-        ir.stmts[idx][:inst].args[1] = Core.mutating_arrayfreeze
+        #println("Optimization of ", stmt)
+        stmt.args[1] = Core.mutating_arrayfreeze
     end
     return ir
 end

base/pointer.jl

@@ -63,6 +63,7 @@ cconvert(::Type{Ptr{UInt8}}, s::AbstractString) = String(s)
 cconvert(::Type{Ptr{Int8}}, s::AbstractString) = String(s)
 
 unsafe_convert(::Type{Ptr{T}}, a::Array{T}) where {T} = ccall(:jl_array_ptr, Ptr{T}, (Any,), a)
+unsafe_convert(::Type{Ptr{T}}, a::Core.ImmutableArray{T}) where {T} = ccall(:jl_array_ptr, Ptr{T}, (Any,), a)
 unsafe_convert(::Type{Ptr{S}}, a::AbstractArray{T}) where {S,T} = convert(Ptr{S}, unsafe_convert(Ptr{T}, a))
 unsafe_convert(::Type{Ptr{T}}, a::AbstractArray{T}) where {T} = error("conversion to pointer not defined for $(typeof(a))")