Reorganize code to keep code as IRCode slightly longer (#39885)

Keno · web-flow · commit 22f84e7d472e · 2021-03-05T22:19:37.000-05:00
This moved the primary place where IRCode gets converted into CodeInfo
into the transform_result_for_cache call, which is a sensible place for
it to be since the primary reason we need to convert back to IRCode is
to make it acceptable for storing in the global cache. The reason we
might want to not perform the conversion, is that the conversion is
slightly lossy, because it drops stmtinfo. However, in Cthulhu, I would
like to keep the statement info around such that Cthulhu can present
it to the user, even in optimized code (i.e. where inlining decided
not to inline, I would still like Cthulhu to be able to introspect
inference's original annotations). This change makes that possible.

In an ideal world, we wouldn't have to do this at all for uncached
code, but of course both code_typed and typeinf_ext do look at the
code, even if it is uncached. Unfortunately, at the moment we don't
really have a good way to indicate whether or not the code will
be looked at, so there is a fallback path that always does the
conversion if we decided not to do the caching. Some future
refactoring can save some additional time here.
diff --git a/base/compiler/abstractinterpretation.jl b/base/compiler/abstractinterpretation.jl
@@ -32,6 +32,7 @@ function abstract_call_gf_by_type(interp::AbstractInterpreter, @nospecialize(f),
                                   fargs::Union{Nothing,Vector{Any}}, argtypes::Vector{Any}, @nospecialize(atype),
                                   sv::InferenceState, max_methods::Int = InferenceParams(interp).MAX_METHODS)
     if sv.params.unoptimize_throw_blocks && sv.currpc in sv.throw_blocks
+        add_remark!(interp, sv, "Skipped call in throw block")
         return CallMeta(Any, false)
     end
     valid_worlds = WorldRange()
@@ -365,7 +366,7 @@ function const_prop_heuristic(interp::AbstractInterpreter, method::Method, mi::M
     if isdefined(code, :inferred) && !cache_inlineable
         cache_inf = code.inferred
         if !(cache_inf === nothing)
-            cache_inlineable = inlining_policy(interp)(cache_inf)
+            cache_inlineable = inlining_policy(interp)(cache_inf) !== nothing
         end
     end
     if !cache_inlineable
diff --git a/base/compiler/optimize.jl b/base/compiler/optimize.jl
@@ -32,14 +32,18 @@ function default_inlining_policy(@nospecialize(src))
     if isa(src, CodeInfo) || isa(src, Vector{UInt8})
         src_inferred = ccall(:jl_ir_flag_inferred, Bool, (Any,), src)
         src_inlineable = ccall(:jl_ir_flag_inlineable, Bool, (Any,), src)
-        return src_inferred && src_inlineable
+        return src_inferred && src_inlineable ? src : nothing
     end
-    return false
+    if isa(src, OptimizationState) && isdefined(src, :ir)
+        return src.src.inlineable ? src.ir : nothing
+    end
+    return nothing
 end
 
 mutable struct OptimizationState
     linfo::MethodInstance
     src::CodeInfo
+    ir::Any # Union{Nothing, IRCode}
     stmt_info::Vector{Any}
     mod::Module
     nargs::Int
@@ -54,7 +58,7 @@ mutable struct OptimizationState
             WorldView(code_cache(interp), frame.world),
             inlining_policy(interp))
         return new(frame.linfo,
-                   frame.src, frame.stmt_info, frame.mod, frame.nargs,
+                   frame.src, nothing, frame.stmt_info, frame.mod, frame.nargs,
                    frame.sptypes, frame.slottypes, false,
                    inlining)
     end
@@ -88,7 +92,7 @@ mutable struct OptimizationState
             WorldView(code_cache(interp), get_world_counter()),
             inlining_policy(interp))
         return new(linfo,
-                   src, stmt_info, inmodule, nargs,
+                   src, nothing, stmt_info, inmodule, nargs,
                    sptypes_from_meth_instance(linfo), slottypes, false,
                    inlining)
         end
@@ -100,6 +104,20 @@ function OptimizationState(linfo::MethodInstance, params::OptimizationParams, in
     return OptimizationState(linfo, src, params, interp)
 end
 
+function ir_to_codeinf!(opt::OptimizationState)
+    replace_code_newstyle!(opt.src, opt.ir, opt.nargs - 1)
+    opt.ir = nothing
+    let src = opt.src::CodeInfo
+        widen_all_consts!(src)
+        src.inferred = true
+        # finish updating the result struct
+        validate_code_in_debug_mode(opt.linfo, src, "optimized")
+        return src
+    end
+end
+
+include("compiler/ssair/driver.jl")
+
 
 #############
 # constants #
@@ -152,7 +170,7 @@ function isinlineable(m::Method, me::OptimizationState, params::OptimizationPara
         end
     end
     if !inlineable
-        inlineable = inline_worthy(me.src.code, me.src, me.sptypes, me.slottypes, params, union_penalties, cost_threshold + bonus)
+        inlineable = inline_worthy(me.ir, params, union_penalties, cost_threshold + bonus)
     end
     return inlineable
 end
@@ -175,7 +193,7 @@ function stmt_affects_purity(@nospecialize(stmt), ir)
 end
 
 # Convert IRCode back to CodeInfo and compute inlining cost and sideeffects
-function finish(opt::OptimizationState, params::OptimizationParams, ir, @nospecialize(result))
+function finish(interp::AbstractInterpreter, opt::OptimizationState, params::OptimizationParams, ir, @nospecialize(result))
     def = opt.linfo.def
     nargs = Int(opt.nargs) - 1
 
@@ -225,7 +243,7 @@ function finish(opt::OptimizationState, params::OptimizationParams, ir, @nospeci
         end
     end
 
-    replace_code_newstyle!(opt.src, ir, nargs)
+    opt.ir = ir
 
     # determine and cache inlineability
     union_penalties = false
@@ -263,14 +281,15 @@ function finish(opt::OptimizationState, params::OptimizationParams, ir, @nospeci
             opt.src.inlineable = isinlineable(def, opt, params, union_penalties, bonus)
         end
     end
+
     nothing
 end
 
 # run the optimization work
 function optimize(interp::AbstractInterpreter, opt::OptimizationState, params::OptimizationParams, @nospecialize(result))
     nargs = Int(opt.nargs) - 1
     @timeit "optimizer" ir = run_passes(opt.src, nargs, opt)
-    finish(opt, params, ir, result)
+    finish(interp, opt, params, ir, result)
 end
 
 
@@ -296,7 +315,7 @@ plus_saturate(x::Int, y::Int) = max(x, y, x+y)
 # known return type
 isknowntype(@nospecialize T) = (T === Union{}) || isa(T, Const) || isconcretetype(widenconst(T))
 
-function statement_cost(ex::Expr, line::Int, src::CodeInfo, sptypes::Vector{Any},
+function statement_cost(ex::Expr, line::Int, src::Union{CodeInfo, IRCode}, sptypes::Vector{Any},
                         slottypes::Vector{Any}, union_penalties::Bool,
                         params::OptimizationParams, error_path::Bool = false)
     head = ex.head
@@ -308,7 +327,7 @@ function statement_cost(ex::Expr, line::Int, src::CodeInfo, sptypes::Vector{Any}
         if ftyp === IntrinsicFunction && farg isa SSAValue
             # if this comes from code that was already inlined into another function,
             # Consts have been widened. try to recover in simple cases.
-            farg = src.code[farg.id]
+            farg = isa(src, CodeInfo) ? src.code[farg.id] : src.stmts[farg.id][:inst]
             if isa(farg, GlobalRef) || isa(farg, QuoteNode) || isa(farg, IntrinsicFunction) || isexpr(farg, :static_parameter)
                 ftyp = argextype(farg, src, sptypes, slottypes)
             end
@@ -350,7 +369,7 @@ function statement_cost(ex::Expr, line::Int, src::CodeInfo, sptypes::Vector{Any}
             end
             return T_FFUNC_COST[fidx]
         end
-        extyp = line == -1 ? Any : src.ssavaluetypes[line]
+        extyp = line == -1 ? Any : argextype(SSAValue(line), src, sptypes, slottypes)
         if extyp === Union{}
             return 0
         end
@@ -361,7 +380,7 @@ function statement_cost(ex::Expr, line::Int, src::CodeInfo, sptypes::Vector{Any}
         # run-time of the function, we omit them from
         # consideration. This way, non-inlined error branches do not
         # prevent inlining.
-        extyp = line == -1 ? Any : src.ssavaluetypes[line]
+        extyp = line == -1 ? Any : argextype(SSAValue(line), src, sptypes, slottypes)
         return extyp === Union{} ? 0 : 20
     elseif head === :(=)
         if ex.args[1] isa GlobalRef
@@ -386,31 +405,32 @@ function statement_cost(ex::Expr, line::Int, src::CodeInfo, sptypes::Vector{Any}
     return 0
 end
 
-function statement_or_branch_cost(@nospecialize(stmt), line::Int, src::CodeInfo, sptypes::Vector{Any},
+function statement_or_branch_cost(@nospecialize(stmt), line::Int, src::Union{CodeInfo, IRCode}, sptypes::Vector{Any},
                                   slottypes::Vector{Any}, union_penalties::Bool, params::OptimizationParams,
                                   throw_blocks::Union{Nothing,BitSet})
     thiscost = 0
+    dst(tgt) = isa(src, IRCode) ? first(src.cfg.blocks[tgt].stmts) : tgt
     if stmt isa Expr
         thiscost = statement_cost(stmt, line, src, sptypes, slottypes, union_penalties, params,
                                   throw_blocks !== nothing && line in throw_blocks)::Int
     elseif stmt isa GotoNode
         # loops are generally always expensive
         # but assume that forward jumps are already counted for from
         # summing the cost of the not-taken branch
-        thiscost = stmt.label < line ? 40 : 0
+        thiscost = dst(stmt.label) < line ? 40 : 0
     elseif stmt isa GotoIfNot
-        thiscost = stmt.dest < line ? 40 : 0
+        thiscost = dst(stmt.dest) < line ? 40 : 0
     end
     return thiscost
 end
 
-function inline_worthy(body::Array{Any,1}, src::CodeInfo, sptypes::Vector{Any}, slottypes::Vector{Any},
+function inline_worthy(ir::IRCode,
                        params::OptimizationParams, union_penalties::Bool=false, cost_threshold::Integer=params.inline_cost_threshold)
     bodycost::Int = 0
-    throw_blocks = params.unoptimize_throw_blocks ? find_throw_blocks(body) : nothing
-    for line = 1:length(body)
-        stmt = body[line]
-        thiscost = statement_or_branch_cost(stmt, line, src, sptypes, slottypes, union_penalties, params, throw_blocks)
+    throw_blocks = params.unoptimize_throw_blocks ? find_throw_blocks(ir.stmts.inst, RefValue(ir)) : nothing
+    for line = 1:length(ir.stmts)
+        stmt = ir.stmts[line][:inst]
+        thiscost = statement_or_branch_cost(stmt, line, ir, ir.sptypes, ir.argtypes, union_penalties, params, throw_blocks)
         bodycost = plus_saturate(bodycost, thiscost)
         bodycost > cost_threshold && return false
     end
@@ -490,5 +510,3 @@ function renumber_ir_elements!(body::Vector{Any}, ssachangemap::Vector{Int}, lab
         end
     end
 end
-
-include("compiler/ssair/driver.jl")
diff --git a/base/compiler/ssair/inlining.jl b/base/compiler/ssair/inlining.jl
@@ -696,13 +696,13 @@ function resolve_todo(todo::InliningTodo, state::InliningState)
     isconst, src = false, nothing
     if isa(spec.match, InferenceResult)
         let inferred_src = spec.match.src
-            if isa(inferred_src, CodeInfo)
-                isconst, src = false, inferred_src
-            elseif isa(inferred_src, Const)
+            if isa(inferred_src, Const)
                 if !is_inlineable_constant(inferred_src.val)
                     return compileable_specialization(state.et, spec.match)
                 end
                 isconst, src = true, quoted(inferred_src.val)
+            else
+                isconst, src = false, inferred_src
             end
         end
     else
@@ -724,17 +724,15 @@ function resolve_todo(todo::InliningTodo, state::InliningState)
         return ConstantCase(src)
     end
 
-    if src !== nothing && !state.policy(src)
-        src = nothing
+    if src !== nothing
+        src = state.policy(src)
     end
 
     if src === nothing
         return compileable_specialization(state.et, spec.match)
     end
 
-    if isa(src, CodeInfo) || isa(src, Vector{UInt8})
-
-    elseif isa(src, IRCode)
+    if isa(src, IRCode)
         src = copy(src)
     end
 
@@ -1028,7 +1026,7 @@ function inline_invoke!(ir::IRCode, idx::Int, sig::Signature, info::InvokeCallIn
     calltype = ir.stmts[idx][:type]
 
     if !info.match.fully_covers
-        # XXX: We could union split this
+        # TODO: We could union split out the signature check and continue on
         return nothing
     end
 
diff --git a/base/compiler/typeinfer.jl b/base/compiler/typeinfer.jl
@@ -210,6 +210,18 @@ function typeinf(interp::AbstractInterpreter, frame::InferenceState)
     end
 end
 
+function finish!(interp::AbstractInterpreter, caller::InferenceResult)
+    # If we didn't transform the src for caching, we may have to transform
+    # it anyway for users like typeinf_ext. Do that here.
+    opt = caller.src
+    if may_optimize(interp) && opt isa OptimizationState
+        if opt.ir !== nothing
+            caller.src = ir_to_codeinf!(opt)
+        end
+    end
+    return caller.src
+end
+
 function _typeinf(interp::AbstractInterpreter, frame::InferenceState)
     typeinf_nocycle(interp, frame) || return false # frame is now part of a higher cycle
     # with no active ip's, frame is done
@@ -235,30 +247,26 @@ function _typeinf(interp::AbstractInterpreter, frame::InferenceState)
               frames[i].cached )
         for i in 1:length(frames) ]
     empty!(frames)
-    if may_optimize(interp)
-        for (caller, _, _) in results
-            opt = caller.src
-            if opt isa OptimizationState
-                result_type = caller.result
-                @assert !(result_type isa LimitedAccuracy)
-                optimize(interp, opt, OptimizationParams(interp), result_type)
-                finish(opt.src, interp)
-                # finish updating the result struct
-                validate_code_in_debug_mode(opt.linfo, opt.src, "optimized")
-                if opt.const_api
-                    if result_type isa Const
-                        caller.src = result_type
-                    else
-                        @assert isconstType(result_type)
-                        caller.src = Const(result_type.parameters[1])
-                    end
-                elseif opt.src.inferred
-                    caller.src = opt.src::CodeInfo # stash a copy of the code (for inlining)
+    for (caller, _, _) in results
+        opt = caller.src
+        if may_optimize(interp) && opt isa OptimizationState
+            result_type = caller.result
+            @assert !(result_type isa LimitedAccuracy)
+            optimize(interp, opt, OptimizationParams(interp), result_type)
+            if opt.const_api
+                # XXX: The work in ir_to_codeinf! is essentially wasted. The only reason
+                # we're doing it is so that code_llvm can return the code
+                # for the `return ...::Const` (which never runs anyway). We should do this
+                # as a post processing step instead.
+                ir_to_codeinf!(opt)
+                if result_type isa Const
+                    caller.src = result_type
                 else
-                    caller.src = nothing
+                    @assert isconstType(result_type)
+                    caller.src = Const(result_type.parameters[1])
                 end
-                caller.valid_worlds = opt.inlining.et.valid_worlds[]
             end
+            caller.valid_worlds = opt.inlining.et.valid_worlds[]
         end
     end
     for (caller, edges, cached) in results
@@ -271,6 +279,7 @@ function _typeinf(interp::AbstractInterpreter, frame::InferenceState)
         if cached
             cache_result!(interp, caller)
         end
+        finish!(interp, caller)
     end
     return true
 end
@@ -349,7 +358,12 @@ function transform_result_for_cache(interp::AbstractInterpreter, linfo::MethodIn
     # If we decided not to optimize, drop the OptimizationState now.
     # External interpreters can override as necessary to cache additional information
     if inferred_result isa OptimizationState
-        inferred_result = inferred_result.src
+        opt = inferred_result
+        if isa(opt.src, CodeInfo)
+            inferred_result = ir_to_codeinf!(opt)
+        else
+            inferred_result = opt.src
+        end
     end
     if inferred_result isa CodeInfo
         inferred_result.min_world = first(valid_worlds)
@@ -470,13 +484,6 @@ function finish(me::InferenceState, interp::AbstractInterpreter)
     nothing
 end
 
-function finish(src::CodeInfo, interp::AbstractInterpreter)
-    # convert all type information into the form consumed by the cache for inlining and code-generation
-    widen_all_consts!(src)
-    src.inferred = true
-    nothing
-end
-
 # record the backedges
 function store_backedges(frame::InferenceResult, edges::Vector{Any})
     toplevel = !isa(frame.linfo.def, Method)
@@ -839,9 +846,9 @@ function typeinf_code(interp::AbstractInterpreter, method::Method, @nospecialize
     frame === nothing && return (nothing, Any)
     if typeinf(interp, frame) && run_optimizer
         opt_params = OptimizationParams(interp)
-        opt = OptimizationState(frame, opt_params, interp)
-        optimize(interp, opt, opt_params, ignorelimited(result.result))
-        opt.src.inferred = true
+        result.src = OptimizationState(frame, opt_params, interp)
+        optimize(interp, result.src, opt_params, ignorelimited(result.result))
+        frame.src = finish!(interp, result)
     end
     ccall(:jl_typeinf_end, Cvoid, ())
     frame.inferred || return (nothing, Any)
