Don't inline generated functions if we can't invoke their generator (JuliaLang#59018)

This PR disables inlining (via an outer `src_inlining_policy` method)
for method instances obtained from `@generated` methods if we cannot
invoke their generator during compilation, deferring their execution to
runtime via dynamic dispatch.

Generated functions may be called with arguments that are not fully
inferrable, yet we might inline them. The lack of concretely inferred
argument types can cause the generator function to not be invoked
`may_invoke_generator(mi) === false`, which causes its fallback
definition to be executed when one is available to expand its body (the
`else` branch of a `if @generated` construct, if present). However,
fallback definitions were intended to enable more accurate type
inference, but not to be actually used for execution; they will
typically be slower, and in this case, it is best to defer the
generation of the method body to runtime, where the generator can be
invoked with concrete argument types.

Implements
JuliaLang#59013 (comment),
fixing JuliaLang#58915.

Author: serenity4

Compiler/src/abstractinterpretation.jl

@@ -1245,7 +1245,7 @@ function const_prop_methodinstance_heuristic(interp::AbstractInterpreter,
         if isa(code, CodeInstance)
             inferred = @atomic :monotonic code.inferred
             # TODO propagate a specific `CallInfo` that conveys information about this call
-            if src_inlining_policy(interp, inferred, NoCallInfo(), IR_FLAG_NULL)
+            if src_inlining_policy(interp, mi, inferred, NoCallInfo(), IR_FLAG_NULL)
                 return true
             end
         end

Compiler/src/optimize.jl

@@ -132,6 +132,17 @@ is_declared_noinline(@nospecialize src::MaybeCompressed) =
 #####################
 
 # return whether this src should be inlined. If so, retrieve_ir_for_inlining must return an IRCode from it
+
+function src_inlining_policy(interp::AbstractInterpreter, mi::MethodInstance,
+    @nospecialize(src), @nospecialize(info::CallInfo), stmt_flag::UInt32)
+    # If we have a generator, but we can't invoke it (because argument type information is lacking),
+    # don't inline so we defer its invocation to runtime where we'll have precise type information.
+    if isa(mi.def, Method) && hasgenerator(mi)
+        may_invoke_generator(mi) || return false
+    end
+    return src_inlining_policy(interp, src, info, stmt_flag)
+end
+
 function src_inlining_policy(interp::AbstractInterpreter,
     @nospecialize(src), @nospecialize(info::CallInfo), stmt_flag::UInt32)
     isa(src, OptimizationState) && (src = src.src)

Compiler/src/ssair/inlining.jl

@@ -865,7 +865,7 @@ function resolve_todo(mi::MethodInstance, result::Union{Nothing,InferenceResult,
         return compileable_specialization(edge, effects, et, info, state)
     end
 
-    src_inlining_policy(state.interp, src, info, flag) ||
+    src_inlining_policy(state.interp, mi, src, info, flag) ||
         return compileable_specialization(edge, effects, et, info, state)
 
     add_inlining_edge!(et, edge)
@@ -897,7 +897,7 @@ function resolve_todo(mi::MethodInstance, @nospecialize(info::CallInfo), flag::U
         return nothing
     end
 
-    src_inlining_policy(state.interp, src, info, flag) || return nothing
+    src_inlining_policy(state.interp, mi, src, info, flag) || return nothing
     ir, spec_info, debuginfo = retrieve_ir_for_inlining(cached_result, src)
     add_inlining_edge!(et, cached_result)
     return InliningTodo(mi, ir, spec_info, debuginfo, effects)
@@ -1448,7 +1448,7 @@ end
 function semiconcrete_result_item(result::SemiConcreteResult,
         @nospecialize(info::CallInfo), flag::UInt32, state::InliningState)
     code = result.edge
-    mi = code.def
+    mi = get_ci_mi(code)
     et = InliningEdgeTracker(state)
 
     if (!OptimizationParams(state.interp).inlining || is_stmt_noinline(flag) ||
@@ -1458,7 +1458,7 @@ function semiconcrete_result_item(result::SemiConcreteResult,
         (is_declared_noinline(mi.def::Method) && !is_stmt_inline(flag)))
         return compileable_specialization(code, result.effects, et, info, state)
     end
-    src_inlining_policy(state.interp, result.ir, info, flag) ||
+    src_inlining_policy(state.interp, mi, result.ir, info, flag) ||
         return compileable_specialization(code, result.effects, et, info, state)
 
     add_inlining_edge!(et, result.edge)
@@ -1602,10 +1602,8 @@ end
 
 function handle_invoke_expr!(todo::Vector{Pair{Int,Any}}, ir::IRCode,
     idx::Int, stmt::Expr, @nospecialize(info::CallInfo), flag::UInt32, sig::Signature, state::InliningState)
-    mi = stmt.args[1]
-    if !(mi isa MethodInstance)
-        mi = (mi::CodeInstance).def
-    end
+    edge = stmt.args[1]
+    mi = isa(edge, MethodInstance) ? edge : get_ci_mi(edge::CodeInstance)
     case = resolve_todo(mi, info, flag, state)
     handle_single_case!(todo, ir, idx, stmt, case, false)
     return nothing
@@ -1625,13 +1623,13 @@ function assemble_inline_todo!(ir::IRCode, state::InliningState)
     todo = Pair{Int, Any}[]
 
     for idx in 1:length(ir.stmts)
-        flag = ir.stmts[idx][:flag]
+        inst = ir.stmts[idx]
+        flag = inst[:flag]
 
         simpleres = process_simple!(todo, ir, idx, flag, state)
         simpleres === nothing && continue
         stmt, sig = simpleres
-
-        info = ir.stmts[idx][:info]
+        info = inst[:info]
 
         # `NativeInterpreter` won't need this, but provide a support for `:invoke` exprs here
         # for external `AbstractInterpreter`s that may run the inlining pass multiple times

Compiler/src/ssair/passes.jl

@@ -1532,7 +1532,7 @@ end
 function try_inline_finalizer!(ir::IRCode, argexprs::Vector{Any}, idx::Int,
     code::CodeInstance, @nospecialize(info::CallInfo), inlining::InliningState,
     attach_after::Bool)
-    mi = code.def
+    mi = get_ci_mi(code)
     et = InliningEdgeTracker(inlining)
     if code isa CodeInstance
         if use_const_api(code)
@@ -1545,7 +1545,7 @@ function try_inline_finalizer!(ir::IRCode, argexprs::Vector{Any}, idx::Int,
         return false
     end
 
-    src_inlining_policy(inlining.interp, src, info, IR_FLAG_NULL) || return false
+    src_inlining_policy(inlining.interp, mi, src, info, IR_FLAG_NULL) || return false
     src, spec_info, di = retrieve_ir_for_inlining(code, src)
 
     # For now: Require finalizer to only have one basic block

Compiler/test/inline.jl

@@ -2314,4 +2314,43 @@ let src = code_typed1(g_noinline_invoke, (Union{Symbol,Nothing},))
     @test !any(@nospecialize(x)->isa(x,GlobalRef), src.code)
 end
 
+path = Ref{Symbol}(:unknown)
+function f59018_generator(x)
+    if @generated
+        if x isa DataType && x.name === Type.body.name
+            path[] = :generator
+            return Core.sizeof(x.parameters[1])
+        end
+    else
+        path[] = :fallback
+        return Core.sizeof(x.parameters[1])
+    end
+end
+f59018() = f59018_generator(Base.inferencebarrier(Int64))
+let src = code_typed1(f59018, ())
+    # We should hit a dynamic dispatch, because not enough information
+    # is available to expand the generator during compilation.
+    @test iscall((src, f59018_generator), src.code[end - 1])
+    @test path[] === :unknown
+    @test f59018() === 8
+    @test path[] === :generator
+end
+
+# https://github.com/JuliaLang/julia/issues/58915
+f58915(nt) = @inline Base.setindex(nt, 2, :next)
+# This function should fully-inline, i.e. it should have only built-in / intrinsic calls
+# and no invokes or dynamic calls of user code
+let src = code_typed1(f58915, Tuple{@NamedTuple{next::UInt32,prev::UInt32}})
+    # Any calls should be built-in calls
+    @test count(iscall(f->!isa(singleton_type(argextype(f, src)), Core.Builtin)), src.code) == 0
+    # There should be no invoke at all
+    @test count(isinvoke(Returns(true)), src.code) == 0
+end
+
+# https://github.com/JuliaLang/julia/issues/58915#issuecomment-3061421895
+let src = code_typed1(Base.setindex, (@NamedTuple{next::UInt32,prev::UInt32}, Int, Symbol))
+    @test count(isinvoke(:merge_fallback), src.code) == 0
+    @test count(iscall((src, Base.merge_fallback)), src.code) == 0
+end
+
 end # module inline_tests