ImmutableArrays

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,12 +176,19 @@ function vect(X...)
     return copyto!(Vector{T}(undef, length(X)), X)
 end
 
-size(a::Array, d::Integer) = arraysize(a, convert(Int, d))
-size(a::Vector) = (arraysize(a,1),)
-size(a::Matrix) = (arraysize(a,1), arraysize(a,2))
-size(a::Array{<:Any,N}) where {N} = (@_inline_meta; ntuple(M -> size(a, M), Val(N))::Dims)
+# 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)
 
-asize_from(a::Array, n) = n > ndims(a) ? () : (arraysize(a,n), asize_from(a, n+1)...)
+# 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))
+size(a::IMArray{<:Any,N}) where {N} = (@_inline_meta; ntuple(M -> size(a, M), Val(N))::Dims)
+
+asize_from(a::IMArray, n) = n > ndims(a) ? () : (arraysize(a,n), asize_from(a, n+1)...)
 
 allocatedinline(T::Type) = (@_pure_meta; ccall(:jl_stored_inline, Cint, (Any,), T) != Cint(0))
 
@@ -223,6 +259,13 @@ function isassigned(a::Array, i::Int...)
     ccall(:jl_array_isassigned, Cint, (Any, UInt), a, ii) == 1
 end
 
+function isassigned(a::ImmutableArray, i::Int...)
+    @_inline_meta
+    ii = (_sub2ind(size(a), i...) % UInt) - 1
+    @boundscheck ii < length(a) % UInt || return false
+    ccall(:jl_array_isassigned, Cint, (Any, UInt), a, ii) == 1
+end
+
 ## copy ##
 
 """
@@ -378,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...])
@@ -611,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)
 
@@ -622,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
@@ -895,6 +951,9 @@ function getindex end
 @eval getindex(A::Array, i1::Int) = arrayref($(Expr(:boundscheck)), A, i1)
 @eval getindex(A::Array, i1::Int, i2::Int, I::Int...) = (@_inline_meta; arrayref($(Expr(:boundscheck)), A, i1, i2, I...))
 
+@eval getindex(A::ImmutableArray, i1::Int) = arrayref($(Expr(:boundscheck)), A, i1)
+@eval getindex(A::ImmutableArray, i1::Int, i2::Int, I::Int...) = (@_inline_meta; arrayref($(Expr(:boundscheck)), A, i1, i2, I...))
+
 # Faster contiguous indexing using copyto! for UnitRange and Colon
 function getindex(A::Array, I::AbstractUnitRange{<:Integer})
     @_inline_meta

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,7 +307,8 @@ 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)
+    #@timeit "compact 3" ir = compact!(ir)
+    ir = memory_opt!(ir)
     #@Base.show ir
     if JLOptions().debug_level == 2
         @timeit "verify 3" (verify_ir(ir); verify_linetable(ir.linetable))

base/compiler/ssair/ir.jl

@@ -319,6 +319,13 @@ function setindex!(x::IRCode, @nospecialize(repl), s::SSAValue)
     return x
 end
 
+function ssadominates(ir::IRCode, domtree::DomTree, ssa1::Int, ssa2::Int)
+    bb1 = block_for_inst(ir.cfg, ssa1)
+    bb2 = block_for_inst(ir.cfg, ssa2)
+    bb1 == bb2 && return ssa1 < ssa2
+    return dominates(domtree, bb1, bb2)
+end
+
 # SSA values that need renaming
 struct OldSSAValue
     id::Int

base/compiler/ssair/passes.jl

@@ -1255,3 +1255,185 @@ function cfg_simplify!(ir::IRCode)
     compact.active_result_bb = length(bb_starts)
     return finish(compact)
 end
+
+# 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 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}() # 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
+            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
+
+        if is_known_call(stmt, arrayset, compact) && length(stmt.args) >= 5
+            # The value being set escapes, everything else doesn't
+            mark_escape(stmt.args[4])
+            arr = stmt.args[3]
+            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
+            # Assume everything else escapes
+            for ur in userefs(stmt)
+                mark_escape(ur[])
+            end
+        end
+    end
+
+    ir = finish(compact)
+    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
+        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
+        if haskey(uses, id)
+            for use in uses[id]
+                if ssadominates(ir, domtree, idx, use)
+                    ok = false
+                    break
+                end
+            end
+        end
+        ok || continue
+        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