perform inference using optimizer-derived type information

In certain cases, the optimizer can introduce new type information.
This is particularly evident in SROA, where load forwarding can reveal
type information that was not visible during abstract interpretation.
In such cases, re-running abstract interpretation using this new type
information can be highly valuable, however, currently, this only occurs
when semi-concrete interpretation happens to be triggered.

This commit introduces a new "post-optimization inference" phase at the
end of the optimizer pipeline. When the optimizer derives new type
information, this phase performs IR abstract interpretation to further
optimize the IR.

Author: aviatesk

Compiler/src/inferencestate.jl

@@ -825,6 +825,7 @@ mutable struct IRInterpretationState
     callstack #::Vector{AbsIntState}
     frameid::Int
     parentid::Int
+    new_call_inferred::Bool
 
     function IRInterpretationState(interp::AbstractInterpreter,
         spec_info::SpecInfo, ir::IRCode, mi::MethodInstance, argtypes::Vector{Any},
@@ -850,7 +851,7 @@ mutable struct IRInterpretationState
         edges = Any[]
         callstack = AbsIntState[]
         return new(spec_info, ir, mi, WorldWithRange(world, valid_worlds), curridx, argtypes_refined, ir.sptypes, tpdum,
-                ssa_refined, lazyreachability, tasks, edges, callstack, 0, 0)
+                ssa_refined, lazyreachability, tasks, edges, callstack, 0, 0, #=new_call_inferred=#false)
     end
 end
 

Compiler/src/optimize.jl

@@ -999,7 +999,7 @@ end
 
 # run the optimization work
 function optimize(interp::AbstractInterpreter, opt::OptimizationState, caller::InferenceResult)
-    @timeit "optimizer" ir = run_passes_ipo_safe(opt.src, opt)
+    @timeit "optimizer" ir = run_passes_ipo_safe(interp, opt, caller)
     ipo_dataflow_analysis!(interp, opt, ir, caller)
     return finish(interp, opt, ir, caller)
 end
@@ -1019,27 +1019,25 @@ matchpass(optimize_until::Int, stage, _) = optimize_until == stage
 matchpass(optimize_until::String, _, name) = optimize_until == name
 matchpass(::Nothing, _, _) = false
 
-function run_passes_ipo_safe(
-    ci::CodeInfo,
-    sv::OptimizationState,
-    optimize_until = nothing,  # run all passes by default
-)
+function run_passes_ipo_safe(interp::AbstractInterpreter, sv::OptimizationState, result::InferenceResult;
+                             optimize_until = nothing)  # run all passes by default
+    ci = sv.src
     __stage__ = 0  # used by @pass
     # NOTE: The pass name MUST be unique for `optimize_until::String` to work
     @pass "convert"   ir = convert_to_ircode(ci, sv)
     @pass "slot2reg"  ir = slot2reg(ir, ci, sv)
     # TODO: Domsorting can produce an updated domtree - no need to recompute here
     @pass "compact 1" ir = compact!(ir)
     @pass "Inlining"  ir = ssa_inlining_pass!(ir, sv.inlining, ci.propagate_inbounds)
-    # @timeit "verify 2" verify_ir(ir)
     @pass "compact 2" ir = compact!(ir)
     @pass "SROA"      ir = sroa_pass!(ir, sv.inlining)
-    @pass "ADCE"      (ir, made_changes) = adce_pass!(ir, sv.inlining)
-    if made_changes
-        @pass "compact 3" ir = compact!(ir, true)
-    end
+    @pass "ADCE"      ir, changed = adce_pass!(ir, sv.inlining)
+    @pass "compact 3" changed && (
+                      ir = compact!(ir, true))
+    @pass "optinf"    optinf_worthwhile(ir) && (
+                      ir = optinf!(ir, interp, sv, result))
     if is_asserts()
-        @timeit "verify 3" begin
+        @timeit "verify" begin
             verify_ir(ir, true, false, optimizer_lattice(sv.inlining.interp), sv.linfo)
             verify_linetable(ir.debuginfo, length(ir.stmts))
         end
@@ -1048,6 +1046,48 @@ function run_passes_ipo_safe(
     return ir
 end
 
+# If the optimizer derives new type information (as implied by `IR_FLAG_REFINED`),
+# and this new type information is available for the arguments of a call expression,
+# further optimizations may be possible by performing irinterp on the optimized IR.
+function optinf_worthwhile(ir::IRCode)
+    @assert isempty(ir.new_nodes) "expected compacted IRCode"
+    for i = 1:length(ir.stmts)
+        inst = ir[SSAValue(i)]
+        if has_flag(inst, IR_FLAG_REFINED)
+            if isexpr(inst[:stmt], :call)
+                return true
+            end
+        end
+    end
+    return false
+end
+
+function optinf!(ir::IRCode, interp::AbstractInterpreter, sv::OptimizationState, result::InferenceResult)
+    ci = sv.src
+    spec_info = SpecInfo(ci)
+    world = get_inference_world(interp)
+    min_world, max_world = first(result.valid_worlds), last(result.valid_worlds)
+    irsv = IRInterpretationState(interp, spec_info, ir, result.linfo, ir.argtypes,
+                                 world, min_world, max_world)
+    rt, (nothrow, noub) = ir_abstract_constant_propagation(interp, irsv)
+    if irsv.new_call_inferred
+        ir = ssa_inlining_pass!(ir, sv.inlining, ci.propagate_inbounds)
+        ir = compact!(ir)
+        effects = result.effects
+        if nothrow
+            effects = Effects(effects; nothrow=true)
+        end
+        if noub
+            effects = Effects(effects; noub=ALWAYS_TRUE)
+        end
+        result.effects = effects
+        result.exc_result = refine_exception_type(result.exc_result, effects)
+        ⋤ = strictneqpartialorder(ipo_lattice(interp))
+        result.result = rt ⋤ result.result ? rt : result.result
+    end
+    return ir
+end
+
 function strip_trailing_junk!(code::Vector{Any}, ssavaluetypes::Vector{Any}, ssaflags::Vector, debuginfo::DebugInfoStream, cfg::CFG, info::Vector{CallInfo})
     # Remove `nothing`s at the end, we don't handle them well
     # (we expect the last instruction to be a terminator)

Compiler/src/ssair/EscapeAnalysis.jl

@@ -24,10 +24,10 @@ using Base:       # Base definitions
     isempty, length, max, min, missing, println, push!, pushfirst!,
     !, !==, &, *, +, -, :, <, <<, >, |, ∈, ∉, ∩, ∪, ≠, ≤, ≥, ⊆
 using ..Compiler: # Compiler specific definitions
-    AbstractLattice, Compiler, IRCode, IR_FLAG_NOTHROW,
+    @show, AbstractLattice, Compiler, IRCode, IR_FLAG_NOTHROW,
     argextype, fieldcount_noerror, has_flag, intrinsic_nothrow, is_meta_expr_head,
     is_identity_free_argtype, isexpr, setfield!_nothrow, singleton_type, try_compute_field,
-    try_compute_fieldidx, widenconst
+    try_compute_fieldidx, widenconst, ⊑
 
 function include(x::String)
     if !isdefined(Base, :end_base_include)

Compiler/src/ssair/ir.jl

@@ -1709,7 +1709,8 @@ function reprocess_phi_node!(𝕃ₒ::AbstractLattice, compact::IncrementalCompa
 
     # There's only one predecessor left - just replace it
     v = phi.values[1]
-    if !⊑(𝕃ₒ, compact[compact.ssa_rename[old_idx]][:type], argextype(v, compact))
+    ⋤ = strictneqpartialorder(𝕃ₒ)
+    if argextype(v, compact) ⋤ compact[compact.ssa_rename[old_idx]][:type]
         v = Refined(v)
     end
     compact.ssa_rename[old_idx] = v

Compiler/src/ssair/irinterp.jl

@@ -58,6 +58,7 @@ function abstract_call(interp::AbstractInterpreter, arginfo::ArgInfo, sstate::St
     call = abstract_call(interp, arginfo, si, irsv)::Future
     Future{Any}(call, interp, irsv) do call, interp, irsv
         irsv.ir.stmts[irsv.curridx][:info] = call.info
+        irsv.new_call_inferred |= true
         nothing
     end
     return call
@@ -204,7 +205,8 @@ function reprocess_instruction!(interp::AbstractInterpreter, inst::Instruction,
         # Handled at the very end
         return false
     elseif isa(stmt, PiNode)
-        rt = tmeet(typeinf_lattice(interp), argextype(stmt.val, ir), widenconst(stmt.typ))
+        ⊓ = join(typeinf_lattice(interp))
+        rt = argextype(stmt.val, ir) ⊓ widenconst(stmt.typ)
     elseif stmt === nothing
         return false
     elseif isa(stmt, GlobalRef)
@@ -226,7 +228,9 @@ function reprocess_instruction!(interp::AbstractInterpreter, inst::Instruction,
                 inst[:stmt] = quoted(rt.val)
             end
             return true
-        elseif !⊑(typeinf_lattice(interp), inst[:type], rt)
+        end
+        ⋤ = strictneqpartialorder(typeinf_lattice(interp))
+        if rt ⋤ inst[:type]
             inst[:type] = rt
             return true
         end

Compiler/src/ssair/passes.jl

@@ -989,9 +989,10 @@ function lift_keyvalue_get!(compact::IncrementalCompact, idx::Int, stmt::Expr, 
     lifted_leaves === nothing && return
 
     result_t = Union{}
+    ⊔ = join(𝕃ₒ)
     for v in values(lifted_leaves)
         v === nothing && return
-        result_t = tmerge(𝕃ₒ, result_t, argextype(v.val, compact))
+        result_t = result_t ⊔ argextype(v.val, compact)
     end
 
     (lifted_val, nest) = perform_lifting!(compact,
@@ -1001,8 +1002,12 @@ function lift_keyvalue_get!(compact::IncrementalCompact, idx::Int, stmt::Expr, 
     compact[idx] = lifted_val === nothing ? nothing : Expr(:call, GlobalRef(Core, :tuple), lifted_val.val)
     finish_phi_nest!(compact, nest)
     if lifted_val !== nothing
-        if !⊑(𝕃ₒ, compact[SSAValue(idx)][:type], tuple_tfunc(𝕃ₒ, Any[result_t]))
-            add_flag!(compact[SSAValue(idx)], IR_FLAG_REFINED)
+        stmttype = tuple_tfunc(𝕃ₒ, Any[result_t])
+        inst = compact[SSAValue(idx)]
+        ⋤ = strictneqpartialorder(𝕃ₒ)
+        if stmttype ⋤ inst[:type]
+            inst[:type] = stmttype
+            add_flag!(inst, IR_FLAG_REFINED)
         end
     end
 
@@ -1440,19 +1445,23 @@ function sroa_pass!(ir::IRCode, inlining::Union{Nothing,InliningState}=nothing)
         lifted_leaves, any_undef = lifted_result
 
         result_t = Union{}
+        ⊔ = join(𝕃ₒ)
         for v in values(lifted_leaves)
             v === nothing && continue
-            result_t = tmerge(𝕃ₒ, result_t, argextype(v.val, compact))
+            result_t = result_t ⊔ argextype(v.val, compact)
         end
 
         (lifted_val, nest) = perform_lifting!(compact,
             visited_philikes, field, result_t, lifted_leaves, val, lazydomtree)
 
         should_delete_node = false
-        line = compact[SSAValue(idx)][:line]
-        if lifted_val !== nothing && !⊑(𝕃ₒ, compact[SSAValue(idx)][:type], result_t)
+        inst = compact[SSAValue(idx)]
+        line = inst[:line]
+        ⋤ = strictneqpartialorder(𝕃ₒ)
+        if lifted_val !== nothing && result_t ⋤ inst[:type]
             compact[idx] = lifted_val === nothing ? nothing : lifted_val.val
-            add_flag!(compact[SSAValue(idx)], IR_FLAG_REFINED)
+            inst[:type] = result_t
+            add_flag!(inst, IR_FLAG_REFINED)
         elseif lifted_val === nothing || isa(lifted_val.val, AnySSAValue)
             # Save some work in a later compaction, by inserting this into the renamer now,
             # but only do this if we didn't set the REFINED flag, to save work for irinterp
@@ -1855,9 +1864,15 @@ function sroa_mutables!(ir::IRCode, defuses::IdDict{Int,Tuple{SPCSet,SSADefUse}}
                 for use in du.uses
                     if use.kind === :getfield
                         inst = ir[SSAValue(use.idx)]
-                        inst[:stmt] = compute_value_for_use(ir, domtree, allblocks,
+                        newvalue = compute_value_for_use(ir, domtree, allblocks,
                             du, phinodes, fidx, use.idx)
-                        add_flag!(inst, IR_FLAG_REFINED)
+                        inst[:stmt] = newvalue
+                        newvaluetyp = argextype(newvalue, ir)
+                        ⋤ = strictneqpartialorder(𝕃ₒ)
+                        if newvaluetyp ⋤ inst[:type]
+                            inst[:type] = newvaluetyp
+                            add_flag!(inst, IR_FLAG_REFINED)
+                        end
                     elseif use.kind === :isdefined
                         continue # already rewritten if possible
                     elseif use.kind === :nopreserve
@@ -1878,11 +1893,12 @@ function sroa_mutables!(ir::IRCode, defuses::IdDict{Int,Tuple{SPCSet,SSADefUse}}
                 for b in phiblocks
                     n = ir[phinodes[b]][:stmt]::PhiNode
                     result_t = Bottom
+                    ⊔ = join(𝕃ₒ)
                     for p in ir.cfg.blocks[b].preds
                         push!(n.edges, p)
                         v = compute_value_for_block(ir, domtree, allblocks, du, phinodes, fidx, p)
                         push!(n.values, v)
-                        result_t = tmerge(𝕃ₒ, result_t, argextype(v, ir))
+                        result_t = result_t ⊔ argextype(v, ir)
                     end
                     ir[phinodes[b]][:type] = result_t
                 end

Compiler/src/typeinfer.jl

@@ -118,10 +118,9 @@ function finish!(interp::AbstractInterpreter, caller::InferenceState)
             # we can now widen our applicability in the global cache too
             store_backedges(ci, edges)
         end
-        inferred_result = nothing
-        uncompressed = inferred_result
-        const_flag = is_result_constabi_eligible(result)
-        discard_src = caller.cache_mode === CACHE_MODE_NULL || const_flag
+        uncompressed = inferred_result = nothing
+        (; rettype, exctype, rettype_const, const_flags) = ResultForCache(result)
+        discard_src = caller.cache_mode === CACHE_MODE_NULL || is_result_constabi_eligible(result)
         if !discard_src
             inferred_result = transform_result_for_cache(interp, result)
             # TODO: do we want to augment edges here with any :invoke targets that we got from inlining (such that we didn't have a direct edge to it already)?
@@ -143,9 +142,13 @@ function finish!(interp::AbstractInterpreter, caller::InferenceState)
         if !@isdefined di
             di = DebugInfo(result.linfo)
         end
-        ccall(:jl_update_codeinst, Cvoid, (Any, Any, Int32, UInt, UInt, UInt32, Any, Any, Any),
-            ci, inferred_result, const_flag, first(result.valid_worlds), last(result.valid_worlds), encode_effects(result.ipo_effects),
-            result.analysis_results, di, edges)
+        min_world, max_world = first(result.valid_worlds), last(result.valid_worlds)
+        ipo_effects = encode_effects(result.ipo_effects)
+        ccall(:jl_update_codeinst, Cvoid, (
+            Any, Any, Any, Any, Any, Int32, UInt, UInt,
+            UInt32, Any, Any, Any),
+            ci, inferred_result, rettype, exctype, rettype_const, const_flags, min_world, max_world,
+            ipo_effects, result.analysis_results, di, edges)
         engine_reject(interp, ci)
         if !discard_src && isdefined(interp, :codegen) && uncompressed isa CodeInfo
             # record that the caller could use this result to generate code when required, if desired, to avoid repeating n^2 work
@@ -488,37 +491,15 @@ function finishinfer!(me::InferenceState, interp::AbstractInterpreter)
 
     # finish populating inference results into the CodeInstance if possible, and maybe cache that globally for use elsewhere
     if isdefined(result, :ci)
-        result_type = result.result
-        result_type isa LimitedAccuracy && (result_type = result_type.typ)
-        @assert !(result_type === nothing)
-        if isa(result_type, Const)
-            rettype_const = result_type.val
-            const_flags = is_result_constabi_eligible(result) ? 0x3 : 0x2
-        elseif isa(result_type, PartialOpaque)
-            rettype_const = result_type
-            const_flags = 0x2
-        elseif isconstType(result_type)
-            rettype_const = result_type.parameters[1]
-            const_flags = 0x2
-        elseif isa(result_type, PartialStruct)
-            rettype_const = result_type.fields
-            const_flags = 0x2
-        elseif isa(result_type, InterConditional)
-            rettype_const = result_type
-            const_flags = 0x2
-        elseif isa(result_type, InterMustAlias)
-            rettype_const = result_type
-            const_flags = 0x2
-        else
-            rettype_const = nothing
-            const_flags = 0x0
-        end
+        (; rettype, exctype, rettype_const, const_flags) = ResultForCache(result)
         di = nothing
         edges = empty_edges # `edges` will be updated within `finish!`
         ci = result.ci
-        ccall(:jl_fill_codeinst, Cvoid, (Any, Any, Any, Any, Int32, UInt, UInt, UInt32, Any, Any, Any),
-            ci, widenconst(result_type), widenconst(result.exc_result), rettype_const, const_flags,
-            first(result.valid_worlds), last(result.valid_worlds),
+        min_world, max_world = first(result.valid_worlds), last(result.valid_worlds)
+        ccall(:jl_fill_codeinst, Cvoid, (
+            Any, Any, Any, Any, Int32, UInt, UInt,
+            UInt32, Any, Any, Any),
+            ci, rettype, exctype, rettype_const, const_flags, min_world, max_world,
             encode_effects(result.ipo_effects), result.analysis_results, di, edges)
         if is_cached(me) # CACHE_MODE_GLOBAL
             cached_result = cache_result!(me.interp, result, ci)
@@ -530,6 +511,46 @@ function finishinfer!(me::InferenceState, interp::AbstractInterpreter)
     nothing
 end
 
+struct ResultForCache
+    rettype
+    exctype
+    rettype_const
+    const_flags::UInt8
+    ResultForCache(rettype, exctype, rettype_const, const_flags::UInt8) = (
+        @nospecialize rettype exctype rettype_const;
+        new(rettype, exctype, rettype_const, const_flags))
+end
+@inline function ResultForCache(result::InferenceResult)
+    result_type = result.result
+    result_type isa LimitedAccuracy && (result_type = result_type.typ)
+    @assert !(result_type === nothing)
+    rettype = widenconst(result_type)
+    exctype = widenconst(result.exc_result)
+    if isa(result_type, Const)
+        rettype_const = result_type.val
+        const_flags = is_result_constabi_eligible(result) ? 0x3 : 0x2
+    elseif isa(result_type, PartialOpaque)
+        rettype_const = result_type
+        const_flags = 0x2
+    elseif isconstType(result_type)
+        rettype_const = result_type.parameters[1]
+        const_flags = 0x2
+    elseif isa(result_type, PartialStruct)
+        rettype_const = result_type.fields
+        const_flags = 0x2
+    elseif isa(result_type, InterConditional)
+        rettype_const = result_type
+        const_flags = 0x2
+    elseif isa(result_type, InterMustAlias)
+        rettype_const = result_type
+        const_flags = 0x2
+    else
+        rettype_const = nothing
+        const_flags = 0x0
+    end
+    return ResultForCache(rettype, exctype, rettype_const, const_flags)
+end
+
 # record the backedges
 function store_backedges(caller::CodeInstance, edges::SimpleVector)
     isa(caller.def.def, Method) || return # don't add backedges to toplevel method instance
@@ -1025,7 +1046,7 @@ function typeinf_ircode(interp::AbstractInterpreter, mi::MethodInstance,
     end
     (; result) = frame
     opt = OptimizationState(frame, interp)
-    ir = run_passes_ipo_safe(opt.src, opt, optimize_until)
+    ir = run_passes_ipo_safe(interp, opt, result; optimize_until)
     rt = widenconst(ignorelimited(result.result))
     return ir, rt
 end
@@ -1050,6 +1071,7 @@ function typeinf_frame(interp::AbstractInterpreter, mi::MethodInstance, run_opti
             opt = OptimizationState(frame, interp)
             optimize(interp, opt, frame.result)
             src = ir_to_codeinf!(opt)
+            src.rettype = widenconst(result.result)
         end
         result.src = frame.src = src
     end