Implement maybecopy in BoundsError, start optimization, refactor memo…

…ry_opt!

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,18 @@ 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))
+
+"""
+    maybecopy(x)
+
+`maybecopy` provides access to `x` while ensuring it does not escape.
+To do so, the optimizer decides whether to create a copy of `x` or not based on the implementation
+That is, `maybecopy` will either be a call to [`copy`](@ref) or just a reference to x.
+"""
+maybecopy = Core.maybecopy
+
 # T[x...] constructs Array{T,1}
 """
     getindex(type[, elements...])
@@ -626,8 +666,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 +677,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,8 +279,17 @@ 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 arrays to avoid escaping them
+    # Eventually, we want to figure out if the copy is needed to save the performance of copying
+    # (i.e., if a escapes elsewhere, don't bother to make a copy)
+
+    BoundsError(@nospecialize(a)) = a isa Array ?
+                                    (@_noinline_meta; new(Core.maybecopy(a))) :
+                                    (@_noinline_meta; new(a))
+
+    BoundsError(@nospecialize(a), i) = a isa Array ?
+                                       (@_noinline_meta; new(Core.maybecopy(a), i)) :
+                                       (@_noinline_meta; new(a, i))
 end
 struct DivideError         <: Exception end
 struct OutOfMemoryError    <: 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,76 +1256,184 @@ 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
+    maybecopies = Int[] # calls to maybecopy
+
+    function mark_escape(@nospecialize val)
         isa(val, SSAValue) || return
+        #println(val.id, " escaped.")
         val.id in relevant && pop!(relevant, val.id)
     end
+
+    function mark_use(val, idx)
+        isa(val, SSAValue) || return
+        id = val.id
+        id in relevant || return
+        (haskey(uses, id)) || (uses[id] = Int[])
+        push!(uses[id], idx)
+    end
+
     for ((_, idx), stmt) in compact
+
+        #println("idx: ", idx, " = ", stmt)
+
         if isa(stmt, ReturnNode)
             isdefined(stmt, :val) || continue
             val = stmt.val
-            if isa(val, SSAValue) && val.id in relevant
-                (haskey(uses, val.id)) || (uses[val.id] = Int[])
-                push!(uses[val.id], idx)
-            end
+            mark_use(val, idx)
             continue
+
+        # check for phinodes that are possibly allocations
+        elseif isa(stmt, PhiNode)
+
+            # ensure all of the phinode values are defined
+            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
+                    push!(relevant, idx)
+                end
+            end
         end
+
         (isexpr(stmt, :call) || isexpr(stmt, :foreigncall)) || continue
+
+        if is_known_call(stmt, Core.maybecopy, compact)
+            push!(maybecopies, idx)
+            continue
+        end
+
         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
+            mark_use(arr, idx)
+
+        elseif is_known_call(stmt, arrayref, compact) && length(stmt.args) == 4
+            arr = stmt.args[3]
+            mark_use(arr, idx)
+
+        elseif is_known_call(stmt, setindex!, compact) && length(stmt.args) == 4
+            # handle similarly to arrayset
+            val = stmt.args[3]
+            mark_escape(val)
+
+            arr = stmt.args[2]
+            mark_use(arr, idx)
+
+        elseif is_known_call(stmt, (===), compact) && length(stmt.args) == 3
+            arr1 = stmt.args[2]
+            arr2 = stmt.args[3]
+
+            mark_use(arr1, idx)
+            mark_use(arr2, idx)
+
+        # 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
+            arr = stmt.args[6]
+            mark_use(arr, idx)
+
+        elseif is_known_call(stmt, arraysize, compact) && isa(stmt.args[2], SSAValue)
+            arr = stmt.args[2]
+            mark_use(arr, idx)
+
         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
+    isempty(revisit) && isempty(maybecopies) && 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
         (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
+        stmt.args[1] = Core.mutating_arrayfreeze
     end
+
+    # TODO: Use escape analysis info to determine if maybecopy should copy
+
+    # for idx in maybecopies
+    #     stmt = ir.stmts[idx][:inst]::Expr
+    #     #println(stmt.args)
+    #     arr = stmt.args[2]
+    #     id = isa(arr, SSAValue) ? arr.id : arr.n # SSAValue or Core.Argument
+
+    #     if (id in relevant) # didn't escape elsewhere, so make a copy to keep it un-escaped
+    #         #println("didn't escape maybecopy")
+    #         stmt.args[1] = Main.Base.copy
+    #     else # already escaped, so save the cost of copying and just pass the actual object
+    #         #println("escaped maybecopy")
+    #         ir.stmts[idx][:inst] = arr
+    #     end
+    # 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))")
 

src/builtin_proto.h

@@ -54,6 +54,7 @@ DECLARE_BUILTIN(_typevar);
 DECLARE_BUILTIN(arrayfreeze);
 DECLARE_BUILTIN(arraythaw);
 DECLARE_BUILTIN(mutating_arrayfreeze);
+DECLARE_BUILTIN(maybecopy);
 
 JL_CALLABLE(jl_f_invoke_kwsorter);
 JL_CALLABLE(jl_f__structtype);