opaque_closure: Allow opting out of PartialOpaque (#54734)

Keno · web-flow · commit 1c59231b2dae · 2024-06-09T01:06:03.000-04:00
`PartialOpaque` is a powerful mechanism to recover some amount of
identity from opaque closure for inlining and other optimizations.
However, there are two downsides:
1. It causes additional inference work, which is unnecessary if you
already know that you don't want the identity-based optimization.
2. We (currently) disallow :new_opaque_closure in code returned from
generated functions, because we cannot ensure that the identity of any
resulting PartialOpaque will be stable. This somewhat defeats the
purpose of having the opaque closure be, well, opaque.

This PR adds an additional argument to `new_opaque_closure` that decides
whether or not inference is allowed to form `PartialOpaque` for this
`:new_opaque_closure`. If not, it is also permitted to run during
precompile.
diff --git a/base/compiler/abstractinterpretation.jl b/base/compiler/abstractinterpretation.jl
@@ -2547,7 +2547,7 @@ function abstract_eval_new_opaque_closure(interp::AbstractInterpreter, e::Expr,
     𝕃ᵢ = typeinf_lattice(interp)
     rt = Union{}
     effects = Effects() # TODO
-    if length(e.args) >= 4
+    if length(e.args) >= 5
         ea = e.args
         argtypes = collect_argtypes(interp, ea, vtypes, sv)
         if argtypes === nothing
@@ -2556,7 +2556,11 @@ function abstract_eval_new_opaque_closure(interp::AbstractInterpreter, e::Expr,
         else
             mi = frame_instance(sv)
             rt = opaque_closure_tfunc(𝕃ᵢ, argtypes[1], argtypes[2], argtypes[3],
-                argtypes[4], argtypes[5:end], mi)
+                argtypes[5], argtypes[6:end], mi)
+            if ea[4] !== true && isa(rt, PartialOpaque)
+                rt = widenconst(rt)
+                # Propagation of PartialOpaque disabled
+            end
             if isa(rt, PartialOpaque) && isa(sv, InferenceState) && !call_result_unused(sv, sv.currpc)
                 # Infer this now so that the specialization is available to
                 # optimization.
diff --git a/base/compiler/optimize.jl b/base/compiler/optimize.jl
@@ -347,7 +347,7 @@ function stmt_effect_flags(𝕃ₒ::AbstractLattice, @nospecialize(stmt), @nospe
             ⊑(𝕃ₒ, typ, Tuple) || return (false, false, false)
             rt_lb = argextype(args[2], src)
             rt_ub = argextype(args[3], src)
-            source = argextype(args[4], src)
+            source = argextype(args[5], src)
             if !(⊑(𝕃ₒ, rt_lb, Type) && ⊑(𝕃ₒ, rt_ub, Type) && ⊑(𝕃ₒ, source, Method))
                 return (false, false, false)
             end
diff --git a/base/compiler/ssair/passes.jl b/base/compiler/ssair/passes.jl
@@ -443,8 +443,8 @@ function lift_leaves(compact::IncrementalCompact, field::Int,
                     ocleaf = simple_walk(compact, ocleaf)
                 end
                 ocdef, _ = walk_to_def(compact, ocleaf)
-                if isexpr(ocdef, :new_opaque_closure) && isa(field, Int) && 1 ≤ field ≤ length(ocdef.args)-4
-                    lift_arg!(compact, leaf, cache_key, ocdef, 4+field, lifted_leaves)
+                if isexpr(ocdef, :new_opaque_closure) && isa(field, Int) && 1 ≤ field ≤ length(ocdef.args)-5
+                    lift_arg!(compact, leaf, cache_key, ocdef, 5+field, lifted_leaves)
                     continue
                 end
                 return nothing
diff --git a/base/compiler/validation.jl b/base/compiler/validation.jl
@@ -34,7 +34,7 @@ const VALID_EXPR_HEADS = IdDict{Symbol,UnitRange{Int}}(
     :throw_undef_if_not => 2:2,
     :aliasscope => 0:0,
     :popaliasscope => 0:0,
-    :new_opaque_closure => 4:typemax(Int),
+    :new_opaque_closure => 5:typemax(Int),
     :import => 1:typemax(Int),
     :using => 1:typemax(Int),
     :export => 1:typemax(Int),
diff --git a/base/deprecated.jl b/base/deprecated.jl
@@ -24,7 +24,8 @@ const __internal_changes_list = (
     :invertedlinetables,
     :codeinforefactor,
     :miuninferredrm,
-    :codeinfonargs  # #54341
+    :codeinfonargs,  # #54341
+    :ocnopartial,
     # Add new change names above this line
 )
 
diff --git a/doc/src/devdocs/ast.md b/doc/src/devdocs/ast.md
@@ -519,18 +519,22 @@ These symbols appear in the `head` field of [`Expr`](@ref)s in lowered form.
 
         The function signature of the opaque closure. Opaque closures don't participate in dispatch, but the input types can be restricted.
 
-      * `args[2]` : isva
-
-        Indicates whether the closure accepts varargs.
-
-      * `args[3]` : lb
+      * `args[2]` : lb
 
         Lower bound on the output type. (Defaults to `Union{}`)
 
-      * `args[4]` : ub
+      * `args[3]` : ub
 
         Upper bound on the output type. (Defaults to `Any`)
 
+      * `args[4]` : constprop
+
+        Indicates whether the opaque closure's identity may be used for constant
+        propagation. The `@opaque` macro enables this by default, but this will
+        cause additional inference which may be undesirable and prevents the
+        code from running during precompile.
+        If `args[4]` is a method, the argument is considered skipped.
+
       * `args[5]` : method
 
         The actual method as an `opaque_closure_method` expression.
diff --git a/src/ast.c b/src/ast.c
@@ -464,7 +464,7 @@ static jl_value_t *scm_to_julia(fl_context_t *fl_ctx, value_t e, jl_module_t *mo
     }
     JL_CATCH {
         // if expression cannot be converted, replace with error expr
-        //jl_(jl_current_exception(ct));
+        //jl_(jl_current_exception(jl_current_task));
         //jlbacktrace();
         jl_expr_t *ex = jl_exprn(jl_error_sym, 1);
         v = (jl_value_t*)ex;
diff --git a/src/codegen.cpp b/src/codegen.cpp
@@ -6513,15 +6513,16 @@ static jl_cgval_t emit_expr(jl_codectx_t &ctx, jl_value_t *expr, ssize_t ssaidx_
         return mark_julia_type(ctx, val, true, (jl_value_t*)jl_any_type);
     }
     else if (head == jl_new_opaque_closure_sym) {
-        assert(nargs >= 4 && "Not enough arguments in new_opaque_closure");
-        SmallVector<jl_cgval_t, 4> argv(nargs, jl_cgval_t());
+        assert(nargs >= 5 && "Not enough arguments in new_opaque_closure");
+        SmallVector<jl_cgval_t, 5> argv(nargs, jl_cgval_t());
         for (size_t i = 0; i < nargs; ++i) {
             argv[i] = emit_expr(ctx, args[i]);
         }
         const jl_cgval_t &argt = argv[0];
         const jl_cgval_t &lb = argv[1];
         const jl_cgval_t &ub = argv[2];
-        const jl_cgval_t &source = argv[3];
+        // argv[3] - constprop marker not used here
+        const jl_cgval_t &source = argv[4];
         if (source.constant == NULL) {
             // For now, we require non-constant source to be handled by using
             // eval. This should probably be a verifier error and an abort here.
@@ -6539,17 +6540,18 @@ static jl_cgval_t emit_expr(jl_codectx_t &ctx, jl_value_t *expr, ssize_t ssaidx_
             jl_value_t *env_t = NULL;
             JL_GC_PUSH2(&closure_t, &env_t);
 
-            SmallVector<jl_value_t *, 0> env_component_ts(nargs-4);
-            for (size_t i = 0; i < nargs - 4; ++i) {
-                jl_value_t *typ = argv[4+i].typ;
+            size_t ncapture_args = nargs-5;
+            SmallVector<jl_value_t *, 0> env_component_ts(ncapture_args);
+            for (size_t i = 0; i < ncapture_args; ++i) {
+                jl_value_t *typ = argv[nargs-ncapture_args+i].typ;
                 if (typ == jl_bottom_type) {
                     JL_GC_POP();
                     return jl_cgval_t();
                 }
                 env_component_ts[i] = typ;
             }
 
-            env_t = jl_apply_tuple_type_v(env_component_ts.data(), nargs-4);
+            env_t = jl_apply_tuple_type_v(env_component_ts.data(), ncapture_args);
             // we need to know the full env type to look up the right specialization
             if (jl_is_concrete_type(env_t)) {
                 jl_tupletype_t *argt_typ = (jl_tupletype_t*)argt.constant;
@@ -6567,7 +6569,7 @@ static jl_cgval_t emit_expr(jl_codectx_t &ctx, jl_value_t *expr, ssize_t ssaidx_
                         fptr = mark_julia_type(ctx, Constant::getNullValue(ctx.types().T_size), false, jl_voidpointer_type);
 
                     // TODO: Inline the env at the end of the opaque closure and generate a descriptor for GC
-                    jl_cgval_t env = emit_new_struct(ctx, env_t, nargs-4, ArrayRef<jl_cgval_t>(argv).drop_front(4));
+                    jl_cgval_t env = emit_new_struct(ctx, env_t, ncapture_args, ArrayRef<jl_cgval_t>(argv).drop_front(nargs-ncapture_args));
 
                     jl_cgval_t closure_fields[5] = {
                         env,
diff --git a/src/interpreter.c b/src/interpreter.c
@@ -298,7 +298,7 @@ static jl_value_t *eval_value(jl_value_t *e, interpreter_state *s)
             argv[i] = eval_value(args[i], s);
         JL_NARGSV(new_opaque_closure, 4);
         jl_value_t *ret = (jl_value_t*)jl_new_opaque_closure((jl_tupletype_t*)argv[0], argv[1], argv[2],
-            argv[3], argv+4, nargs-4, 1);
+            argv[4], argv+5, nargs-5, 1);
         JL_GC_POP();
         return ret;
     }
diff --git a/src/julia-syntax.scm b/src/julia-syntax.scm
@@ -3998,7 +3998,7 @@ f(x) = yt(x)
                                            v)))
                                    cvs)))
                `(new_opaque_closure
-                 ,(cadr e) (call (core apply_type) (core Union)) (core Any)
+                 ,(cadr e) (call (core apply_type) (core Union)) (core Any) (true)
                  (opaque_closure_method (null) ,nargs ,isva ,functionloc ,(convert-lambda lam2 (car (lam:args lam2)) #f '() (symbol-to-idx-map cvs)))
                  ,@var-exprs))))
           ((method)
@@ -4564,13 +4564,13 @@ f(x) = yt(x)
                               (cons (cadr e) (cons fptr (cdddr e)))))
                            ;; Leave a literal lambda in place for later global expansion
                            ((eq? (car e) 'new_opaque_closure)
-                            (let* ((oc_method (car (list-tail (cdr e) 3))) ;; opaque_closure_method
+                            (let* ((oc_method (car (list-tail (cdr e) 4))) ;; opaque_closure_method
                                    (lambda (list-ref oc_method 5))
                                    (lambda (linearize lambda)))
                               (append
-                               (compile-args (list-head (cdr e) 3) break-labels)
+                               (compile-args (list-head (cdr e) 4) break-labels)
                                (list (append (butlast oc_method) (list lambda)))
-                               (compile-args (list-tail (cdr e) 4) break-labels))))
+                               (compile-args (list-tail (cdr e) 5) break-labels))))
                            ;; NOTE: 1st argument to cglobal treated same as for ccall
                            ((and (length> e 2)
                                  (or (eq? (cadr e) 'cglobal)
diff --git a/src/method.c b/src/method.c
@@ -783,11 +783,20 @@ JL_DLLEXPORT jl_code_info_t *jl_code_for_staged(jl_method_instance_t *mi, size_t
         }
 
         // If this generated function has an opaque closure, cache it for
-        // correctness of method identity
+        // correctness of method identity. In particular, other methods that call
+        // this method may end up referencing it in a PartialOpaque lattice element
+        // type. If the method identity were to change (for the same world age)
+        // in between invocations of this method, that return type inference would
+        // no longer be correct.
         int needs_cache_for_correctness = 0;
         for (int i = 0; i < jl_array_nrows(func->code); ++i) {
             jl_value_t *stmt = jl_array_ptr_ref(func->code, i);
             if (jl_is_expr(stmt) && ((jl_expr_t*)stmt)->head == jl_new_opaque_closure_sym) {
+                if (jl_expr_nargs(stmt) >= 4 && jl_is_bool(jl_exprarg(stmt, 3)) && !jl_unbox_bool(jl_exprarg(stmt, 3))) {
+                    // If this new_opaque_closure is prohibited from sourcing PartialOpaque,
+                    // there is no problem
+                    continue;
+                }
                 if (jl_options.incremental && jl_generating_output())
                     jl_error("Impossible to correctly handle OpaqueClosure inside @generated returned during precompile process.");
                 needs_cache_for_correctness = 1;
diff --git a/src/opaque_closure.c b/src/opaque_closure.c
@@ -170,10 +170,10 @@ JL_DLLEXPORT jl_opaque_closure_t *jl_new_opaque_closure_from_code_info(jl_tuplet
 
 JL_CALLABLE(jl_new_opaque_closure_jlcall)
 {
-    if (nargs < 4)
+    if (nargs < 5)
         jl_error("new_opaque_closure: Not enough arguments");
     return (jl_value_t*)jl_new_opaque_closure((jl_tupletype_t*)args[0],
-        args[1], args[2], args[3], &args[4], nargs-4, 1);
+        args[1], args[2], args[4], &args[5], nargs-5, 1);
 }
 
 // check whether the specified number of arguments is compatible with the
diff --git a/test/compiler/inference.jl b/test/compiler/inference.jl
@@ -4262,10 +4262,10 @@ let # Test the presence of PhiNodes in lowered IR by taking the above function,
     Core.Compiler.replace_code_newstyle!(ci, ir)
     ci.ssavaluetypes = length(ci.ssavaluetypes)
     @test any(x->isa(x, Core.PhiNode), ci.code)
-    oc = @eval b->$(Expr(:new_opaque_closure, Tuple{Bool, Float64}, Any, Any,
+    oc = @eval b->$(Expr(:new_opaque_closure, Tuple{Bool, Float64}, Any, Any, true,
         Expr(:opaque_closure_method, nothing, 2, false, LineNumberNode(0, nothing), ci)))(b, 1.0)
     @test Base.return_types(oc, Tuple{Bool}) == Any[Float64]
-    oc = @eval ()->$(Expr(:new_opaque_closure, Tuple{Bool, Float64}, Any, Any,
+    oc = @eval ()->$(Expr(:new_opaque_closure, Tuple{Bool, Float64}, Any, Any, true,
         Expr(:opaque_closure_method, nothing, 2, false, LineNumberNode(0, nothing), ci)))(true, 1.0)
     @test Base.return_types(oc, Tuple{}) == Any[Float64]
 end
diff --git a/test/opaque_closure.jl b/test/opaque_closure.jl
@@ -10,7 +10,7 @@ const lno = LineNumberNode(1, :none)
 
 let ci = @code_lowered const_int()
     @eval function oc_trivial()
-        $(Expr(:new_opaque_closure, Tuple{}, Any, Any,
+        $(Expr(:new_opaque_closure, Tuple{}, Any, Any, true,
             Expr(:opaque_closure_method, nothing, 0, false, lno, ci)))
     end
 end
@@ -19,7 +19,7 @@ end
 
 let ci = @code_lowered const_int()
     @eval function oc_simple_inf()
-        $(Expr(:new_opaque_closure, Tuple{}, Union{}, Any,
+        $(Expr(:new_opaque_closure, Tuple{}, Union{}, Any, true,
             Expr(:opaque_closure_method, nothing, 0, false, lno, ci)))
     end
 end
@@ -33,7 +33,7 @@ end
 (a::OcClos2Int)() = getfield(a, 1) + getfield(a, 2)
 let ci = @code_lowered OcClos2Int(1, 2)();
     @eval function oc_trivial_clos()
-        $(Expr(:new_opaque_closure, Tuple{}, Int, Int,
+        $(Expr(:new_opaque_closure, Tuple{}, Int, Int, true,
             Expr(:opaque_closure_method, nothing, 0, false, lno, ci),
             1, 2))
     end
@@ -42,7 +42,7 @@ end
 
 let ci = @code_lowered OcClos2Int(1, 2)();
     @eval function oc_self_call_clos()
-        $(Expr(:new_opaque_closure, Tuple{}, Int, Int,
+        $(Expr(:new_opaque_closure, Tuple{}, Int, Int, true,
             Expr(:opaque_closure_method, nothing, 0, false, lno, ci),
             1, 2))()
     end
@@ -59,7 +59,7 @@ end
 (a::OcClos1Any)() = getfield(a, 1)
 let ci = @code_lowered OcClos1Any(1)()
     @eval function oc_pass_clos(x)
-        $(Expr(:new_opaque_closure, Tuple{}, Any, Any,
+        $(Expr(:new_opaque_closure, Tuple{}, Any, Any, true,
             Expr(:opaque_closure_method, nothing, 0, false, lno, ci),
             :x))
     end
@@ -69,7 +69,7 @@ end
 
 let ci = @code_lowered OcClos1Any(1)()
     @eval function oc_infer_pass_clos(x)
-        $(Expr(:new_opaque_closure, Tuple{}, Union{}, Any,
+        $(Expr(:new_opaque_closure, Tuple{}, Union{}, Any, true,
             Expr(:opaque_closure_method, nothing, 0, false, lno, ci),
             :x))
     end
@@ -81,7 +81,7 @@ end
 
 let ci = @code_lowered identity(1)
     @eval function oc_infer_pass_id()
-        $(Expr(:new_opaque_closure, Tuple{Any}, Any, Any,
+        $(Expr(:new_opaque_closure, Tuple{Any}, Any, Any, true,
             Expr(:opaque_closure_method, nothing, 1, false, lno, ci)))
     end
 end
@@ -103,7 +103,7 @@ end
 
 let ci = @code_lowered OcOpt([1 2])()
     @eval function oc_opt_ndims(A)
-        $(Expr(:new_opaque_closure, Tuple{}, Union{}, Any,
+        $(Expr(:new_opaque_closure, Tuple{}, Union{}, Any, true,
             Expr(:opaque_closure_method, nothing, 0, false, lno, ci),
             :A))
     end
@@ -184,7 +184,7 @@ mk_va_opaque() = @opaque (x...)->x
 let ci = @code_lowered const_int()
     global function mk_ocg(world::UInt, source, args...)
         @nospecialize
-        cig = Meta.lower(@__MODULE__, Expr(:new_opaque_closure, Tuple{}, Any, Any,
+        cig = Meta.lower(@__MODULE__, Expr(:new_opaque_closure, Tuple{}, Any, Any, true,
             Expr(:opaque_closure_method, nothing, 0, false, lno, ci))).args[1]
         cig.slotnames = Symbol[Symbol("#self#")]
         cig.slottypes = Any[Any]
@@ -298,7 +298,7 @@ eval_oc_spec(oc) = oc()
 for f in (const_int, const_int_barrier)
     ci = code_lowered(f, Tuple{})[1]
     for compiled in (true, false)
-        oc_expr = Expr(:new_opaque_closure, Tuple{}, Union{}, Float64,
+        oc_expr = Expr(:new_opaque_closure, Tuple{}, Union{}, Float64, true,
             Expr(:opaque_closure_method, nothing, 0, false, lno, ci))
         oc_mismatch = let ci = code_lowered(f, Tuple{})[1]
             if compiled
diff --git a/test/precompile.jl b/test/precompile.jl
@@ -1969,4 +1969,31 @@ precompile_test_harness("No backedge precompile") do load_path
     @test first(methods(NoBackEdges.f)).specializations.cache.max_world === typemax(UInt)
 end
 
+# Test precompilation of generated functions that return opaque closures
+# (with constprop marker set to false).
+precompile_test_harness("Generated Opaque") do load_path
+    write(joinpath(load_path, "GeneratedOpaque.jl"),
+        """
+        module GeneratedOpaque
+        using Base.Experimental: @opaque
+        using InteractiveUtils
+        const_int_barrier() = Base.inferencebarrier(1)::typeof(1)
+        const lno = LineNumberNode(1, :none)
+
+        const ci = @code_lowered const_int_barrier()
+        @generated function oc_re_generated_no_partial()
+            Expr(:new_opaque_closure, Tuple{}, Any, Any, false,
+                Expr(:opaque_closure_method, nothing, 0, false, lno, ci))
+        end
+        @assert oc_re_generated_no_partial()() === 1
+        end
+        """)
+    Base.compilecache(Base.PkgId("GeneratedOpaque"))
+    @eval using GeneratedOpaque
+    let oc = invokelatest(GeneratedOpaque.oc_re_generated_no_partial)
+        @test oc.source.specializations.cache.max_world === typemax(UInt)
+        @test oc() === 1
+    end
+end
+
 finish_precompile_test!()

Original file line number	Diff line number	Diff line change
`@@ -24,7 +24,8 @@ const __internal_changes_list = (`
`24`	`24`	`:invertedlinetables,`
`25`	`25`	`:codeinforefactor,`
`26`	`26`	`:miuninferredrm,`
`27`		`- :codeinfonargs # #54341`
	`27`	`+ :codeinfonargs, # #54341`
	`28`	`+ :ocnopartial,`
`28`	`29`	`# Add new change names above this line`
`29`	`30`	`)`
`30`	`31`
Original file line number	Diff line number	Diff line change
`@@ -464,7 +464,7 @@ static jl_value_t scm_to_julia(fl_context_t fl_ctx, value_t e, jl_module_t *mo`
`464`	`464`	`}`
`465`	`465`	`JL_CATCH {`
`466`	`466`	`// if expression cannot be converted, replace with error expr`
`467`		`- //jl_(jl_current_exception(ct));`
	`467`	`+ //jl_(jl_current_exception(jl_current_task));`
`468`	`468`	`//jlbacktrace();`
`469`	`469`	`jl_expr_t *ex = jl_exprn(jl_error_sym, 1);`
`470`	`470`	`v = (jl_value_t*)ex;`
Original file line number	Diff line number	Diff line change
`@@ -298,7 +298,7 @@ static jl_value_t eval_value(jl_value_t e, interpreter_state *s)`
`298`	`298`	`argv[i] = eval_value(args[i], s);`
`299`	`299`	`JL_NARGSV(new_opaque_closure, 4);`
`300`	`300`	`jl_value_t ret = (jl_value_t)jl_new_opaque_closure((jl_tupletype_t*)argv[0], argv[1], argv[2],`
`301`		`- argv[3], argv+4, nargs-4, 1);`
	`301`	`+ argv[4], argv+5, nargs-5, 1);`
`302`	`302`	`JL_GC_POP();`
`303`	`303`	`return ret;`
`304`	`304`	`}`