Add simple immutable array memory optimization

Author: Ian Atol

base/compiler/optimize.jl

@@ -551,12 +551,14 @@ function run_passes(ci::CodeInfo, sv::OptimizationState, caller::InferenceResult
     @timeit "Inlining"  ir = ssa_inlining_pass!(ir, ir.linetable, sv.inlining, ci.propagate_inbounds)
     # @timeit "verify 2" verify_ir(ir)
     @timeit "compact 2" ir = compact!(ir)
-    @timeit "SROA" ir, memory_opt = linear_pass!(ir)
-    if memory_opt
-        @timeit "memory_opt_pass!" begin
-            @timeit "Local EA" estate = analyze_escapes(ir,
-                nargs, #=call_resolved=#true, null_escape_cache)
-            @timeit "memory_opt_pass!" ir = memory_opt_pass!(ir, estate)
+    @timeit "SROA" ir, memory_opt, imarray_memory_opt = linear_pass!(ir)
+    if memory_opt || imarray_memory_opt
+        @timeit "Local EA" estate = analyze_escapes(ir, nargs, #=call_resolved=#true, null_escape_cache)
+        if memory_opt
+            @timeit "Memory opt" ir = memory_opt_pass!(ir, estate)
+        end
+        if imarray_memory_opt
+            @timeit "imarray opt" ir = imarray_memoryopt_pass!(ir, estate)
         end
     end
     @timeit "ADCE"      ir = adce_pass!(ir)

base/compiler/ssair/EscapeAnalysis/EscapeAnalysis.jl

@@ -27,7 +27,7 @@ import ._TOP_MOD:     # Base definitions
     pop!, push!, pushfirst!, empty!, delete!, max, min, enumerate, unwrap_unionall,
     ismutabletype
 import Core.Compiler: # Core.Compiler specific definitions
-    Bottom, InferenceResult, IRCode, IR_FLAG_EFFECT_FREE,
+    Arrayish, Bottom, InferenceResult, IRCode, IR_FLAG_EFFECT_FREE,
     isbitstype, isexpr, is_meta_expr_head, println, widenconst, argextype, singleton_type,
     fieldcount_noerror, try_compute_field, try_compute_fieldidx, hasintersect, ⊑,
     intrinsic_nothrow, array_builtin_common_typecheck, arrayset_typecheck,
@@ -1608,7 +1608,7 @@ function escape_builtin!(::typeof(arrayref), astate::AnalysisState, pc::Int, arg
     argtypes = Any[argextype(args[i], astate.ir) for i in 2:length(args)]
     boundcheckt = argtypes[1]
     aryt = argtypes[2]
-    if !array_builtin_common_typecheck(Array, boundcheckt, aryt, argtypes, 3)
+    if !array_builtin_common_typecheck(Arrayish, boundcheckt, aryt, argtypes, 3)
         add_thrown_escapes!(astate, pc, args, 2)
     end
     ary = args[3]

base/compiler/ssair/passes.jl

@@ -547,7 +547,11 @@ function linear_pass!(ir::IRCode)
     compact = IncrementalCompact(ir)
     lifting_cache = IdDict{Pair{AnySSAValue, Any}, AnySSAValue}()
     local memory_opt = false # whether or not to run the memory_opt_pass! pass later
+    local imarray_memory_opt = false
     for ((_, idx), stmt) in compact
+        # presence of arrayfreeze means possible copy elision opportunity, so run imarray_memoryopt_pass!
+        isa(stmt, GlobalRef) && stmt.name === :arrayfreeze && (imarray_memory_opt = true)
+
         isa(stmt, Expr) || continue
         field_ordering = :unspecified
         if isexpr(stmt, :new)
@@ -663,7 +667,7 @@ function linear_pass!(ir::IRCode)
     non_dce_finish!(compact)
     simple_dce!(compact)
     ir = complete(compact)
-    return ir, memory_opt
+    return ir, memory_opt, imarray_memory_opt
 end
 
 function form_new_preserves(origex::Expr, intermediates::Vector{Int}, new_preserves::Vector{Any})
@@ -1761,76 +1765,21 @@ function cfg_simplify!(ir::IRCode)
     return finish(compact)
 end
 
-function is_allocation(stmt)
-    isexpr(stmt, :foreigncall) || return false
-    s = stmt.args[1]
-    isa(s, QuoteNode) && (s = s.value)
-    return s === :jl_alloc_array_1d
-end
-
-function memory_opt!(ir::IRCode)
-    compact = IncrementalCompact(ir, false)
-    uses = IdDict{Int, Vector{Int}}()
-    relevant = IdSet{Int}()
-    revisit = Int[]
-    function mark_val(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
-            if isa(val, SSAValue) && val.id in relevant
-                (haskey(uses, val.id)) || (uses[val.id] = Int[])
-                push!(uses[val.id], idx)
-            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)
-            # The value being set escapes, everything else doesn't
-            mark_val(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)
-        else
-            # For now we assume everything escapes
-            # TODO: We could handle PhiNodes specially and improve this
-            for ur in userefs(stmt)
-                mark_val(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
-        (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
+function imarray_memoryopt_pass!(ir::IRCode, estate::EscapeState)
+    # mark statements that possibly can be optimized
+    for idx in 1:length(ir.stmts)
+        stmt = ir.stmts[idx][:inst]
+        isexpr(stmt, :call) || continue
+        if is_known_call(stmt, Core.arrayfreeze, ir)
+            # array as SSA value might have been initialized within this frame
+            # (thus potentially doesn't escape to anywhere)
+            ary = stmt.args[2]
+            if isa(ary, SSAValue)
+                # if array doesn't escape, we can just change the tag and avoid allocation
+                has_no_escape(estate[ary]) || continue
+                stmt.args[1] = GlobalRef(Core, :mutating_arrayfreeze)
             end
         end
-        ok || continue
-
-        ir.stmts[idx][:inst].args[1] = Core.mutating_arrayfreeze
     end
     return ir
-end
+end