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

Author: Ian Atol

base/boot.jl

@@ -279,9 +279,14 @@ 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, copy to avoid escaping a
+    # Eventually, we want to figure out if the copy is needed to save the performance of copying
+    #BoundsError(@nospecialize(a)) = (@_noinline_meta; new(a))
+    BoundsError(@nospecialize(a)) = (@_noinline_meta; new(Array(a)))
+    #BoundsError(@nospecialize(a), i) = (@_noinline_meta; new(a, i))
+    BoundsError(@nospecialize(a), i) = (@_noinline_meta; new(Array(a), i))
 end
+
 struct DivideError         <: Exception end
 struct OutOfMemoryError    <: Exception end
 struct ReadOnlyMemoryError <: Exception end

base/compiler/ssair/passes.jl

@@ -1256,26 +1256,47 @@ function cfg_simplify!(ir::IRCode)
     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)
+    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 mutating_arrayfreeze
+    relevant = IdSet{Int}() # allocations and their sizes
+    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
@@ -1284,6 +1305,26 @@ 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
@@ -1294,54 +1335,101 @@ function memory_opt!(ir::IRCode)
             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
+        if is_known_call(stmt, arrayset, compact) && length(stmt.args) >= 5
             # The value being set escapes, everything else doesn't
-            (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, arrayref, compact) && length(stmt.args) == 4
+            arr = stmt.args[3]
 
-        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(arr, SSAValue) && arr.id in relevant
+                (haskey(uses, arr.id)) || (uses[arr.id] = Int[])
+                push!(uses[arr.id], idx)
+            end
 
-            while !isa(typ, Type)
-                typ = typeof(typ)
+        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
 
-            if isa(typ, Core.Const)
-                typ = typ.val
+        # 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
 
-            # 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
+        # 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_escape(ur[])
             end
@@ -1357,8 +1445,18 @@ function memory_opt!(ir::IRCode)
         # Make sure that the value we reference didn't escape
         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
         if haskey(uses, id)
@@ -1370,6 +1468,7 @@ function memory_opt!(ir::IRCode)
             end
         end
         ok || continue
+        #println("Optimization of ", stmt)
         stmt.args[1] = Core.mutating_arrayfreeze
     end
     return ir

src/codegen.cpp

@@ -663,12 +663,13 @@ static const auto jl_newbits_func = new JuliaFunction{
 // `julia.typeof` does read memory, but it is effectively readnone before we lower
 // the allocation function. This is OK as long as we lower `julia.typeof` no later than
 // `julia.gc_alloc_obj`.
+// Updated to argmemonly due to C++ deconstructor style usage in jl_f_arrayfreeze / mutating_arrayfreeze
 static const auto jl_typeof_func = new JuliaFunction{
     "julia.typeof",
     [](LLVMContext &C) { return FunctionType::get(T_prjlvalue,
                 {T_prjlvalue}, false); },
     [](LLVMContext &C) { return AttributeList::get(C,
-            Attributes(C, {Attribute::ReadNone, Attribute::NoUnwind, Attribute::NoRecurse}),
+            Attributes(C, {Attribute::ArgMemOnly, Attribute::NoUnwind, Attribute::NoRecurse}),
             Attributes(C, {Attribute::NonNull}),
             None); },
 };

test/compiler/immutablearray.jl

@@ -4,9 +4,9 @@ using Test
 function test_allocate1()
     a = Vector{Float64}(undef, 5)
     for i = 1:5
-        a[i] === i
+        a[i] = i
     end
-    ImmutableArray(a)
+    Core.ImmutableArray(a)
 end
 
 function test_allocate2()
@@ -33,15 +33,29 @@ function test_broadcast1()
     typeof(a .+ a) <: Core.ImmutableArray
 end
 
+function test_allocate5()
+    a = [1,2,3]
+    try 
+        getindex(a, 4)
+    catch end 
+    Core.ImmutableArray(a)
+end
+
 let
+    # warmup for @allocated
+    a,b,c,d,e = test_allocate1(), test_allocate2(), test_allocate3(), test_allocate4(), test_allocate5()
+
+    # these magic values are ~ what the mutable array version would allocate
     @test @allocated(test_allocate1()) < 100
     @test @allocated(test_allocate2()) < 100
     @test @allocated(test_allocate3()) < 150
     @test @allocated(test_allocate4()) < 100
+    @test @allocated(test_allocate5()) < 170
     @test test_broadcast1() == true
 end
 
-#DiffEq Performance Tests
+
+# DiffEq Performance Tests
 
 # using DifferentialEquations
 # using StaticArrays