ImmutableArray works as uType for DiffEq problems

memory_opt! performance and safety improvements

Quick fixes to build

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/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,26 @@ function cfg_simplify!(ir::IRCode)
     return finish(compact)
 end
 
-function is_allocation(stmt)
+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
+    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 mutating_arrayfreeze
+
+    function mark_escape(@nospecialize val)
         isa(val, SSAValue) || return
         val.id in relevant && pop!(relevant, val.id)
     end
+
     for ((_, idx), stmt) in compact
+
         if isa(stmt, ReturnNode)
             isdefined(stmt, :val) || continue
             val = stmt.val
@@ -1282,50 +1285,92 @@ function memory_opt!(ir::IRCode)
             end
             continue
         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)
+            # arrayset expr.args:
+            # :(Base.arrayset)
+            # false
+            # :(%2) array
+            # :(%8) value
+            # :(%7) index
             # The value being set escapes, everything else doesn't
-            mark_val(stmt.args[4])
+            (length(stmt.args) == 5) || continue # fix boundserror during precompile --- but how do we have arrayset with < 5 args?
+            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, Core.arrayfreeze, compact) && isa(stmt.args[2], SSAValue)
             push!(revisit, idx)
+
+        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]
+
+            while !isa(typ, Type)
+                typ = typeof(typ)
+            end
+
+            if isa(typ, Core.Const)
+                typ = typ.val
+            end
+
+            # make sure this call isn't going to throw
+            if 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, the array definitely escapes
+                for ur in userefs(stmt)
+                    mark_escape(ur[])
+                end
+            end
         else
             # For now we assume everything escapes
             # TODO: We could handle PhiNodes specially and improve this
             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
         (id in relevant) || continue
 
         # 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
     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))")