Very WIP: Refactor core inference loops to use less memory

Currently inference uses `O(<number of statements>*<number of slots>)` state
in the core inference loop. This is usually fine, because users don't tend
to write functions that are particularly long. However, MTK does generate
functions that are excessively long and we've observed MTK models that spend
99% of their inference time just allocating and copying this state.
It is possible to get away with significantly smaller state, and this PR is
a first step in that direction, reducing the state to `O(<number of basic blocks>*<number of slots>)`.
Further improvements are possible by making use of slot liveness information
and only storing those slots that are live across a particular basic block.

The core change here is to keep a full set of `slottypes` only at
basic block boundaries rather than at each statement. For statements
in between, the full variable state can be fully recovered by
linearly scanning throughout the basic block, taking note of
slot assignments (together with the SSA type) and NewVarNodes.

The current status of this branch is that the changes appear correct
(no known functional regressions) and significantly improve the MTK
test cases in question (no exact benchmarks here for now, since
the branch still needs a number of fixes before final numbers make
sense), but somewhat regress optimizer quality (which is expected
and just a missing TODO) and bootstrap time (which is not expected
and something I need to dig into).

Author: Keno

base/compiler/abstractinterpretation.jl

@@ -128,6 +128,14 @@ include("compiler/utilities.jl")
 include("compiler/validation.jl")
 include("compiler/methodtable.jl")
 
+function argextype end # imported by EscapeAnalysis
+function stmt_effect_free end # imported by EscapeAnalysis
+function alloc_array_ndims end # imported by EscapeAnalysis
+function try_compute_field end # imported by EscapeAnalysis
+include("compiler/ssair/basicblock.jl")
+include("compiler/ssair/domtree.jl")
+include("compiler/ssair/ir.jl")
+
 include("compiler/inferenceresult.jl")
 include("compiler/inferencestate.jl")
 

base/compiler/compiler.jl

@@ -88,17 +88,22 @@ mutable struct InferenceState
     sptypes::Vector{Any}
     slottypes::Vector{Any}
     src::CodeInfo
+    cfg::CFG
 
     #= intermediate states for local abstract interpretation =#
+    currbb::Int
     currpc::Int
-    ip::BitSetBoundedMinPrioritySet # current active instruction pointers
+    ip::BitSet#=TODO BoundedMinPrioritySet=# # current active instruction pointers
+    was_reached::BitSet
     handler_at::Vector{Int} # current exception handler info
     ssavalue_uses::Vector{BitSet} # ssavalue sparsity and restart info
-    stmt_types::Vector{Union{Nothing, VarTable}}
+    # TODO: Could keep this sparsely by doing structural liveness analysis ahead of time.
+    bb_vartables::Vector{VarTable}
+    pc_vartable::VarTable
     stmt_edges::Vector{Union{Nothing, Vector{Any}}}
     stmt_info::Vector{Any}
 
-    #= interprocedural intermediate states for abstract interpretation =#
+    #= intermediate states for interprocedural abstract interpretation =#
     pclimitations::IdSet{InferenceState} # causes of precision restrictions (LimitedAccuracy) on currpc ssavalue
     limitations::IdSet{InferenceState} # causes of precision restrictions (LimitedAccuracy) on return
     cycle_backedges::Vector{Tuple{InferenceState, Int}} # call-graph backedges connecting from callee to caller
@@ -125,36 +130,40 @@ mutable struct InferenceState
     interp::AbstractInterpreter
 
     # src is assumed to be a newly-allocated CodeInfo, that can be modified in-place to contain intermediate results
-    function InferenceState(result::InferenceResult,
-        src::CodeInfo, cache::Symbol, interp::AbstractInterpreter)
+    function InferenceState(result::InferenceResult, src::CodeInfo, cache::Symbol,
+        interp::AbstractInterpreter)
         linfo = result.linfo
         world = get_world_counter(interp)
         def = linfo.def
         mod = isa(def, Method) ? def.module : def
         sptypes = sptypes_from_meth_instance(linfo)
-
         code = src.code::Vector{Any}
-        nstmts = length(code)
-        currpc = 1
-        ip = BitSetBoundedMinPrioritySet(nstmts)
-        handler_at = compute_trycatch(code, ip.elems)
-        push!(ip, 1)
+        cfg = compute_basic_blocks(code)
+
+        currbb = currpc = 1
+        ip = BitSet(1) # TODO BitSetBoundedMinPrioritySet(1)
+        was_reached = BitSet()
+        handler_at = compute_trycatch(code, BitSet())
         nssavalues = src.ssavaluetypes::Int
         ssavalue_uses = find_ssavalue_uses(code, nssavalues)
-        stmt_types = Union{Nothing, VarTable}[ nothing for i = 1:nstmts ]
+        nstmts = length(code)
         stmt_edges = Union{Nothing, Vector{Any}}[ nothing for i = 1:nstmts ]
         stmt_info = Any[ nothing for i = 1:nstmts ]
 
         nslots = length(src.slotflags)
         slottypes = Vector{Any}(undef, nslots)
+        pc_vartable = VarTable(undef, nslots)
+        bb_vartable_proto = VarTable(undef, nslots)
         argtypes = result.argtypes
-        nargs = length(argtypes)
-        stmt_types[1] = stmt_type1 = VarTable(undef, nslots)
+        nargtypes = length(argtypes)
         for i in 1:nslots
-            argtyp = (i > nargs) ? Bottom : argtypes[i]
-            stmt_type1[i] = VarState(argtyp, i > nargs)
+            argtyp = (i > nargtypes) ? Bottom : argtypes[i]
+            pc_vartable[i] = VarState(argtyp, i > nargtypes)
+            bb_vartable_proto[i] = VarState(Bottom, i > nargtypes)
             slottypes[i] = argtyp
         end
+        bb_vartables = VarTable[i == 1 ? copy(pc_vartable) : copy(bb_vartable_proto)
+            for i = 1:length(cfg.blocks)]
 
         pclimitations = IdSet{InferenceState}()
         limitations = IdSet{InferenceState}()
@@ -183,8 +192,8 @@ mutable struct InferenceState
         cached = cache === :global
 
         frame = new(
-            linfo, world, mod, sptypes, slottypes, src,
-            currpc, ip, handler_at, ssavalue_uses, stmt_types, stmt_edges, stmt_info,
+            linfo, world, mod, sptypes, slottypes, src, cfg,
+            currbb, currpc, ip, was_reached, handler_at, ssavalue_uses, bb_vartables, pc_vartable, stmt_edges, stmt_info,
             pclimitations, limitations, cycle_backedges, callers_in_cycle, dont_work_on_me, parent, inferred,
             result, valid_worlds, bestguess, ipo_effects,
             params, restrict_abstract_call_sites, cached,
@@ -226,6 +235,8 @@ function any_inbounds(code::Vector{Any})
     return false
 end
 
+was_reached((; was_reached)::InferenceState, pc::Int) = pc in was_reached
+
 function compute_trycatch(code::Vector{Any}, ip::BitSet)
     # The goal initially is to record the frame like this for the state at exit:
     # 1: (enter 3) # == 0
@@ -422,7 +433,7 @@ end
 
 update_valid_age!(edge::InferenceState, sv::InferenceState) = update_valid_age!(sv, edge.valid_worlds)
 
-function record_ssa_assign(ssa_id::Int, @nospecialize(new), frame::InferenceState)
+function record_ssa_assign!(ssa_id::Int, @nospecialize(new), frame::InferenceState)
     ssavaluetypes = frame.src.ssavaluetypes::Vector{Any}
     old = ssavaluetypes[ssa_id]
     if old === NOT_FOUND || !(new ⊑ old)
@@ -431,14 +442,19 @@ function record_ssa_assign(ssa_id::Int, @nospecialize(new), frame::InferenceStat
         # guarantee convergence we need to use tmerge here to ensure that is true
         ssavaluetypes[ssa_id] = old === NOT_FOUND ? new : tmerge(old, new)
         W = frame.ip
-        s = frame.stmt_types
         for r in frame.ssavalue_uses[ssa_id]
-            if s[r] !== nothing # s[r] === nothing => unreached statement
-                push!(W, r)
+            if was_reached(frame, r)
+                usebb = block_for_inst(frame.cfg, r)
+                # We're guaranteed to visit the statement if it's in the current
+                # basic block, since SSA values can only ever appear after their
+                # def.
+                if usebb != frame.currbb
+                    push!(W, usebb)
+                end
             end
         end
     end
-    nothing
+    return nothing
 end
 
 function add_cycle_backedge!(frame::InferenceState, caller::InferenceState, currpc::Int)

base/compiler/inferencestate.jl

@@ -88,6 +88,7 @@ mutable struct OptimizationState
     linfo::MethodInstance
     src::CodeInfo
     ir::Union{Nothing, IRCode}
+    was_reached::Union{Nothing, BitSet}
     stmt_info::Vector{Any}
     mod::Module
     sptypes::Vector{Any} # static parameters
@@ -100,7 +101,7 @@ mutable struct OptimizationState
             WorldView(code_cache(interp), frame.world),
             interp)
         return new(frame.linfo,
-                   frame.src, nothing, frame.stmt_info, frame.mod,
+                   frame.src, nothing, frame.was_reached, frame.stmt_info, frame.mod,
                    frame.sptypes, frame.slottypes, inlining)
     end
     function OptimizationState(linfo::MethodInstance, src::CodeInfo, params::OptimizationParams, interp::AbstractInterpreter)
@@ -128,11 +129,13 @@ mutable struct OptimizationState
             WorldView(code_cache(interp), get_world_counter()),
             interp)
         return new(linfo,
-                   src, nothing, stmt_info, mod,
+                   src, nothing, nothing, stmt_info, mod,
                    sptypes_from_meth_instance(linfo), slottypes, inlining)
     end
 end
 
+was_reached((; was_reached)::OptimizationState, pc::Int) = was_reached === nothing || pc in was_reached
+
 function OptimizationState(linfo::MethodInstance, params::OptimizationParams, interp::AbstractInterpreter)
     src = retrieve_code_info(linfo)
     src === nothing && return nothing
@@ -572,9 +575,29 @@ function convert_to_ircode(ci::CodeInfo, sv::OptimizationState)
             end
         end
     end
+    meta = Any[]
     labelmap = coverage ? fill(0, length(code)) : changemap
     while idx <= length(code)
         codeloc = codelocs[idx]
+        stmt = code[idx]
+        if process_meta!(meta, stmt) || !was_reached(sv, oldidx)
+            if oldidx < length(labelmap)
+                changemap[oldidx] != 0 && (changemap[oldidx+1] = changemap[oldidx])
+                if coverage && labelmap[oldidx] != 0
+                    labelmap[oldidx + 1] = labelmap[oldidx]
+                end
+                changemap[oldidx] = -1
+                coverage && (labelmap[oldidx] = -1)
+            end
+            # TODO: It would be more efficient to do this in bulk
+            deleteat!(code, idx)
+            deleteat!(codelocs, idx)
+            deleteat!(ssavaluetypes, idx)
+            deleteat!(stmtinfo, idx)
+            deleteat!(ssaflags, idx)
+            oldidx += 1
+            continue
+        end
         if coverage && codeloc != prevloc && codeloc != 0
             # insert a side-effect instruction before the current instruction in the same basic block
             insert!(code, idx, Expr(:code_coverage_effect))
@@ -589,7 +612,13 @@ function convert_to_ircode(ci::CodeInfo, sv::OptimizationState)
             idx += 1
             prevloc = codeloc
         end
-        if code[idx] isa Expr && ssavaluetypes[idx] === Union{}
+        if isa(stmt, GotoIfNot)
+            if !was_reached(sv, oldidx + 1)
+                code[idx] = GotoNode(stmt.dest)
+            elseif !was_reached(sv, stmt.dest)
+                code[idx] = nothing
+            end
+        elseif stmt isa Expr && ssavaluetypes[idx] === Union{}
             if !(idx < length(code) && isa(code[idx + 1], ReturnNode) && !isdefined((code[idx + 1]::ReturnNode), :val))
                 # insert unreachable in the same basic block after the current instruction (splitting it)
                 insert!(code, idx + 1, ReturnNode())
@@ -607,12 +636,9 @@ function convert_to_ircode(ci::CodeInfo, sv::OptimizationState)
         idx += 1
         oldidx += 1
     end
+
     renumber_ir_elements!(code, changemap, labelmap)
 
-    meta = Any[]
-    for i = 1:length(code)
-        code[i] = remove_meta!(code[i], meta)
-    end
     strip_trailing_junk!(ci, code, stmtinfo)
     cfg = compute_basic_blocks(code)
     types = Any[]
@@ -623,18 +649,13 @@ function convert_to_ircode(ci::CodeInfo, sv::OptimizationState)
     return ir
 end
 
-function remove_meta!(@nospecialize(stmt), meta::Vector{Any})
-    if isa(stmt, Expr)
-        head = stmt.head
-        if head === :meta
-            args = stmt.args
-            if length(args) > 0
-                push!(meta, stmt)
-            end
-            return nothing
-        end
+function process_meta!(meta::Vector{Any}, @nospecialize stmt)
+    isa(stmt, Expr) || return false
+    stmt.head === :meta || return false
+    if length(stmt.args) > 0
+        push!(meta, stmt)
     end
-    return stmt
+    return true
 end
 
 function slot2reg(ir::IRCode, ci::CodeInfo, sv::OptimizationState)
@@ -796,7 +817,9 @@ end
 
 function cumsum_ssamap!(ssamap::Vector{Int})
     rel_change = 0
+    any_change = false
     for i = 1:length(ssamap)
+        any_change = any_change || ssamap[i] != 0
         rel_change += ssamap[i]
         if ssamap[i] == -1
             # Keep a marker that this statement was deleted
@@ -805,16 +828,15 @@ function cumsum_ssamap!(ssamap::Vector{Int})
             ssamap[i] = rel_change
         end
     end
+    return any_change
 end
 
 function renumber_ir_elements!(body::Vector{Any}, ssachangemap::Vector{Int}, labelchangemap::Vector{Int})
-    cumsum_ssamap!(labelchangemap)
+    any_change = cumsum_ssamap!(labelchangemap)
     if ssachangemap !== labelchangemap
-        cumsum_ssamap!(ssachangemap)
-    end
-    if labelchangemap[end] == 0 && ssachangemap[end] == 0
-        return
+        any_change = cumsum_ssamap!(ssachangemap)
     end
+    any_change || return
     for i = 1:length(body)
         el = body[i]
         if isa(el, GotoNode)
@@ -824,7 +846,8 @@ function renumber_ir_elements!(body::Vector{Any}, ssachangemap::Vector{Int}, lab
             if isa(cond, SSAValue)
                 cond = SSAValue(cond.id + ssachangemap[cond.id])
             end
-            body[i] = GotoIfNot(cond, el.dest + labelchangemap[el.dest])
+            was_deleted = labelchangemap[el.dest] == typemin(Int)
+            body[i] = was_deleted ? cond : GotoIfNot(cond, el.dest + labelchangemap[el.dest])
         elseif isa(el, ReturnNode)
             if isdefined(el, :val)
                 val = el.val

base/compiler/optimize.jl

@@ -8,14 +8,6 @@ else
     end
 end
 
-function argextype end # imported by EscapeAnalysis
-function stmt_effect_free end # imported by EscapeAnalysis
-function alloc_array_ndims end # imported by EscapeAnalysis
-function try_compute_field end # imported by EscapeAnalysis
-
-include("compiler/ssair/basicblock.jl")
-include("compiler/ssair/domtree.jl")
-include("compiler/ssair/ir.jl")
 include("compiler/ssair/slot2ssa.jl")
 include("compiler/ssair/inlining.jl")
 include("compiler/ssair/verify.jl")

base/compiler/ssair/driver.jl

@@ -736,7 +736,7 @@ function dominates_ssa(compact::IncrementalCompact, domtree::DomTree, x::AnySSAV
             elseif xinfo !== nothing
                 return !xinfo.attach_after
             else
-                return yinfo.attach_after
+                return (yinfo::NewNodeInfo).attach_after
             end
         end
         return x′.id < y′.id