Add cfunction support for --trim (JuliaLang#58812)

Author: topolarity

Compiler/src/typeinfer.jl

@@ -1480,16 +1480,30 @@ function collectinvokes!(workqueue::CompilationQueue, ci::CodeInfo, sptypes::Vec
                 # No dynamic dispatch to resolve / enqueue
                 continue
             end
+        elseif isexpr(stmt, :cfunction) && length(stmt.args) == 5
+            (pointer_type, f, rt, at, call_type) = stmt.args
+            linfo = ci.parent
 
-            let workqueue = invokelatest_queue
-                # make a best-effort attempt to enqueue the relevant code for the finalizer
-                mi = compileable_specialization_for_call(workqueue.interp, atype)
-                mi === nothing && continue
+            linfo isa MethodInstance || continue
+            at isa SimpleVector || continue
 
-                push!(workqueue, mi)
+            ft = argextype(f, ci, sptypes)
+            argtypes = Any[ft]
+            for i = 1:length(at)
+                push!(argtypes, sp_type_rewrap(at[i], linfo, #= isreturn =# false))
             end
+            atype = argtypes_to_type(argtypes)
+        else
+            # TODO: handle other StmtInfo like OpaqueClosure?
+            continue
+        end
+        let workqueue = invokelatest_queue
+            # make a best-effort attempt to enqueue the relevant code for the dynamic invokelatest call
+            mi = compileable_specialization_for_call(workqueue.interp, atype)
+            mi === nothing && continue
+
+            push!(workqueue, mi)
         end
-        # TODO: handle other StmtInfo like @cfunction and OpaqueClosure?
     end
 end
 

Compiler/src/verifytrim.jl

@@ -14,7 +14,7 @@ using ..Compiler:
      argextype, empty!, error, get, get_ci_mi, get_world_counter, getindex, getproperty,
      hasintersect, haskey, in, isdispatchelem, isempty, isexpr, iterate, length, map!, max,
      pop!, popfirst!, push!, pushfirst!, reinterpret, reverse!, reverse, setindex!,
-     setproperty!, similar, singleton_type, sptypes_from_meth_instance,
+     setproperty!, similar, singleton_type, sptypes_from_meth_instance, sp_type_rewrap,
      unsafe_pointer_to_objref, widenconst, isconcretetype,
      # misc
      @nospecialize, @assert, C_NULL
@@ -256,9 +256,27 @@ function verify_codeinstance!(interp::NativeInterpreter, codeinst::CodeInstance,
                 warn = true # downgrade must-throw calls to be only a warning
             end
         elseif isexpr(stmt, :cfunction)
+            length(stmt.args) != 5 && continue # required by IR legality
+            (pointer_type, f, rt, at, call_type) = stmt.args
+
+            at isa SimpleVector || continue  # required by IR legality
+            ft = argextype(f, codeinfo, sptypes)
+            argtypes = Any[ft]
+            for i = 1:length(at)
+                push!(argtypes, sp_type_rewrap(at[i], get_ci_mi(codeinst), #= isreturn =# false))
+            end
+            atype = argtypes_to_type(argtypes)
+
+            mi = compileable_specialization_for_call(interp, atype)
+            if mi !== nothing
+                # n.b.: Codegen may choose unpredictably to emit this `@cfunction` as a dynamic invoke or a full
+                # dynamic call, but in either case it guarantees that the required adapter(s) are emitted. All
+                # that we are required to verify here is that the callee CodeInstance is covered.
+                ci = get(caches, mi, nothing)
+                ci isa CodeInstance && continue
+            end
+
             error = "unresolved cfunction"
-            #TODO: parse the cfunction expression to check the target is defined
-            warn = true
         elseif isexpr(stmt, :foreigncall)
             foreigncall = stmt.args[1]
             if foreigncall isa QuoteNode

Compiler/test/verifytrim.jl

@@ -36,36 +36,39 @@ let infos = typeinf_ext_toplevel(Any[Core.svec(Nothing, Tuple{typeof(finalizer),
              \[1\] finalizer\(f::Any, o::Any\)""", repr)
 end
 
+# test that basic `cfunction` generation is allowed, when the dispatch target can be resolved
 make_cfunction() = @cfunction(+, Float64, (Int64,Int64))
+let infos = typeinf_ext_toplevel(Any[Core.svec(Ptr{Cvoid}, Tuple{typeof(make_cfunction)})], [Base.get_world_counter()], TRIM_UNSAFE)
+    errors, parents = get_verify_typeinf_trim(infos)
+    @test isempty(errors)
+end
 
 # use TRIM_UNSAFE to bypass verifier inside typeinf_ext_toplevel
-let infos = typeinf_ext_toplevel(Any[Core.svec(Ptr{Cvoid}, Tuple{typeof(make_cfunction)})], [Base.get_world_counter()], TRIM_UNSAFE)
+make_cfunction_bad(@nospecialize(f::Any)) = @cfunction($f, Float64, (Int64,Int64))::Base.CFunction
+let infos = typeinf_ext_toplevel(Any[Core.svec(Base.CFunction, Tuple{typeof(make_cfunction_bad), Any})], [Base.get_world_counter()], TRIM_UNSAFE)
     errors, parents = get_verify_typeinf_trim(infos)
-    @test_broken isempty(errors) # missing cfunction
+    @test !isempty(errors) # missing cfunction
 
-    desc = only(errors)
-    @test desc.first
-    desc = desc.second
+    (is_warning, desc) = only(errors)
+    @test !is_warning
     @test desc isa CallMissing
     @test occursin("cfunction", desc.desc)
     repr = sprint(verify_print_error, desc, parents)
-    @test occursin(
-        r"""^unresolved cfunction from statement \$\(Expr\(:cfunction, Ptr{Nothing}, :\(\$\(QuoteNode\(\+\)\)\), Float64, :\(svec\(Int64, Int64\)::Core.SimpleVector\), :\(:ccall\)\)\)::Ptr{Nothing}
+    @test occursin(r"""^unresolved cfunction from statement \$\(Expr\(:cfunction, Base.CFunction, :\(f::Any\), Float64, :\(svec\(Int64, Int64\)::Core.SimpleVector\), :\(:ccall\)\)\)::Base.CFunction
             Stacktrace:
-             \[1\] make_cfunction\(\)""", repr)
-
+             \[1\] make_cfunction_bad\(f::Any\)""", repr)
     resize!(infos, 1)
     @test infos[1] isa Core.SimpleVector && infos[1][1] isa Type && infos[1][2] isa Type
     errors, parents = get_verify_typeinf_trim(infos)
     desc = only(errors)
     @test !desc.first
     desc = desc.second
     @test desc isa CCallableMissing
-    @test desc.rt == Ptr{Cvoid}
-    @test desc.sig == Tuple{typeof(make_cfunction)}
+    @test desc.rt == Base.CFunction
+    @test desc.sig == Tuple{typeof(make_cfunction_bad), Any}
     @test occursin("unresolved ccallable", desc.desc)
     repr = sprint(verify_print_error, desc, parents)
-    @test repr == "unresolved ccallable for Tuple{$(typeof(make_cfunction))} => Ptr{Nothing}\n\n"
+    @test repr == "unresolved ccallable for Tuple{$(typeof(make_cfunction_bad)), Any} => Base.CFunction\n\n"
 end
 
 let infos = typeinf_ext_toplevel(Any[Core.svec(Base.SecretBuffer, Tuple{Type{Base.SecretBuffer}})], [Base.get_world_counter()], TRIM_UNSAFE)

src/codegen.cpp

@@ -6356,8 +6356,8 @@ static jl_cgval_t emit_expr(jl_codectx_t &ctx, jl_value_t *expr, ssize_t ssaidx_
     }
     else if (head == jl_cfunction_sym) {
         assert(nargs == 5);
-        jl_cgval_t fexpr_rt = emit_expr(ctx, args[1]);
-        return emit_cfunction(ctx, args[0], fexpr_rt, args[2], (jl_svec_t*)args[3]);
+        jl_cgval_t fexpr_val = emit_expr(ctx, args[1]);
+        return emit_cfunction(ctx, args[0], fexpr_val, args[2], (jl_svec_t*)args[3]);
     }
     else if (head == jl_assign_sym) {
         assert(nargs == 2);
@@ -7576,7 +7576,7 @@ static const char *derive_sigt_name(jl_value_t *jargty)
 // Get the LLVM Function* for the C-callable entry point for a certain function
 // and argument types.
 // here argt does not include the leading function type argument
-static jl_cgval_t emit_cfunction(jl_codectx_t &ctx, jl_value_t *output_type, const jl_cgval_t &fexpr_rt, jl_value_t *declrt, jl_svec_t *argt)
+static jl_cgval_t emit_cfunction(jl_codectx_t &ctx, jl_value_t *output_type, const jl_cgval_t &fexpr_val, jl_value_t *declrt, jl_svec_t *argt)
 {
     jl_unionall_t *unionall_env = (jl_is_method(ctx.linfo->def.method) && jl_is_unionall(ctx.linfo->def.method->sig))
         ? (jl_unionall_t*)ctx.linfo->def.method->sig
@@ -7634,8 +7634,8 @@ static jl_cgval_t emit_cfunction(jl_codectx_t &ctx, jl_value_t *output_type, con
     // compute+verify the dispatch signature, and see if it depends on the environment sparams
     bool approx = false;
     sigt = (jl_value_t*)jl_alloc_svec(nargt + 1);
-    jl_svecset(sigt, 0, fexpr_rt.typ);
-    if (!fexpr_rt.constant && (!jl_is_concrete_type(fexpr_rt.typ) || jl_is_kind(fexpr_rt.typ)))
+    jl_svecset(sigt, 0, fexpr_val.typ);
+    if (!fexpr_val.constant && (!jl_is_concrete_type(fexpr_val.typ) || jl_is_kind(fexpr_val.typ)))
         approx = true;
     for (size_t i = 0; i < nargt; i++) {
         jl_value_t *jargty = jl_svecref(argt, i);
@@ -7664,25 +7664,25 @@ static jl_cgval_t emit_cfunction(jl_codectx_t &ctx, jl_value_t *output_type, con
         unionall_env = NULL;
     }
 
-    bool nest = (!fexpr_rt.constant || unionall_env);
+    bool nest = (!fexpr_val.constant || unionall_env);
     if (ctx.emission_context.TargetTriple.isAArch64() || ctx.emission_context.TargetTriple.isARM() || ctx.emission_context.TargetTriple.isPPC64()) {
         if (nest) {
             emit_error(ctx, "cfunction: closures are not supported on this platform");
             JL_GC_POP();
             return jl_cgval_t();
         }
     }
-    const char *name = derive_sigt_name(fexpr_rt.typ);
+    const char *name = derive_sigt_name(fexpr_val.typ);
     Value *F = gen_cfun_wrapper(
             jl_Module, ctx.emission_context,
-            sig, fexpr_rt.constant, name,
+            sig, fexpr_val.constant, name,
             declrt, sigt,
             unionall_env, sparam_vals, &closure_types);
     bool outboxed;
     if (nest) {
         // F is actually an init_trampoline function that returns the real address
         // Now fill in the nest parameters
-        Value *fobj = boxed(ctx, fexpr_rt);
+        Value *fobj = boxed(ctx, fexpr_val);
         jl_svec_t *fill = jl_emptysvec;
         if (closure_types) {
             assert(ctx.spvals_ptr);
@@ -7722,7 +7722,7 @@ static jl_cgval_t emit_cfunction(jl_codectx_t &ctx, jl_value_t *output_type, con
             jl_aliasinfo_t ai = jl_aliasinfo_t::fromTBAA(ctx, tbaa);
             ai.decorateInst(ctx.builder.CreateStore(F, derived_strct));
             ai.decorateInst(ctx.builder.CreateStore(
-                ctx.builder.CreatePtrToInt(literal_pointer_val(ctx, fexpr_rt.constant), ctx.types().T_size),
+                ctx.builder.CreatePtrToInt(literal_pointer_val(ctx, fexpr_val.constant), ctx.types().T_size),
                 ctx.builder.CreateConstInBoundsGEP1_32(ctx.types().T_size, derived_strct, 1)));
             ai.decorateInst(ctx.builder.CreateStore(Constant::getNullValue(ctx.types().T_size),
                     ctx.builder.CreateConstInBoundsGEP1_32(ctx.types().T_size, derived_strct, 2)));

test/trimming/basic_jll.jl

@@ -1,10 +1,21 @@
 using Libdl
 using Zstd_jll # Note this uses the vendored older non-LazyLibrary version of Zstd_jll
 
+function print_string(fptr::Ptr{Cvoid})
+    println(Core.stdout, unsafe_string(ccall(fptr, Cstring, ())))
+end
+
 function @main(args::Vector{String})::Cint
+    # Test the basic "Hello, world!"
     println(Core.stdout, "Julia! Hello, world!")
-    fptr = dlsym(Zstd_jll.libzstd_handle, :ZSTD_versionString)
-    println(Core.stdout, unsafe_string(ccall(fptr, Cstring, ())))
+
+    # Make sure that JLL's are working as expected
     println(Core.stdout, unsafe_string(ccall((:ZSTD_versionString, libzstd), Cstring, ())))
+
+    # Add an indirection via `@cfunction`
+    cfunc = @cfunction(print_string, Cvoid, (Ptr{Cvoid},))
+    fptr = dlsym(Zstd_jll.libzstd_handle, :ZSTD_versionString)
+    ccall(cfunc, Cvoid, (Ptr{Cvoid},), fptr)
+
     return 0
 end