WIP: Add ABIOverwrite type for def field

Together with #54899, this PR is intending to replicate the functionality
of #54373, which allowed particular specializations to have a different
ABI signature than what would be suggested by the MethodInstance's
`specTypes` field. This PR handles that by adding a special `ABIOverwrite`
type, which, when placed in the `def` field of a `CodeInstance`
instructs the system to use the given signature instead.

Author: Keno

Compiler/src/stmtinfo.jl

@@ -71,7 +71,7 @@ function _add_edges_impl(edges::Vector{Any}, info::MethodMatchInfo, mi_edge::Boo
             mi = specialize_method(m) # don't allow `Method`-edge for this optimized format
             edge = mi
         else
-            mi = edge.def
+            mi = edge.def::MethodInstance
         end
         if mi.specTypes === m.spec_types
             add_one_edge!(edges, edge)
@@ -103,7 +103,7 @@ function add_one_edge!(edges::Vector{Any}, edge::MethodInstance)
     while i <= length(edges)
         edgeᵢ = edges[i]
         edgeᵢ isa Int && (i += 2 + edgeᵢ; continue)
-        edgeᵢ isa CodeInstance && (edgeᵢ = edgeᵢ.def)
+        edgeᵢ isa CodeInstance && (edgeᵢ = edgeᵢ.def::MethodInstance)
         edgeᵢ isa MethodInstance || (i += 1; continue)
         if edgeᵢ === edge && !(i > 1 && edges[i-1] isa Type)
             return # found existing covered edge
@@ -118,7 +118,7 @@ function add_one_edge!(edges::Vector{Any}, edge::CodeInstance)
     while i <= length(edges)
         edgeᵢ_orig = edgeᵢ = edges[i]
         edgeᵢ isa Int && (i += 2 + edgeᵢ; continue)
-        edgeᵢ isa CodeInstance && (edgeᵢ = edgeᵢ.def)
+        edgeᵢ isa CodeInstance && (edgeᵢ = edgeᵢ.def::MethodInstance)
         edgeᵢ isa MethodInstance || (i += 1; continue)
         if edgeᵢ === edge.def && !(i > 1 && edges[i-1] isa Type)
             if edgeᵢ_orig isa MethodInstance

base/boot.jl

@@ -465,6 +465,12 @@ struct InitError <: WrappedException
     error
 end
 
+struct ABIOverwrite
+    abi::Type
+    def::MethodInstance
+    ABIOverwrite(@nospecialize(abi::Type), def::MethodInstance) = new(abi, def)
+end
+
 struct PrecompilableError <: Exception end
 
 String(s::String) = s  # no constructor yet

base/show.jl

@@ -1353,7 +1353,13 @@ end
 show(io::IO, mi::Core.MethodInstance) = show_mi(io, mi)
 function show(io::IO, codeinst::Core.CodeInstance)
     print(io, "CodeInstance for ")
-    show_mi(io, codeinst.def)
+    def = codeinst.def
+    if isa(def, Core.ABIOverwrite)
+        show_mi(io, def.def)
+        print(io, " (ABI Overridden)")
+    else
+        show_mi(io, def::MethodInstance)
+    end
 end
 
 function show_mi(io::IO, mi::Core.MethodInstance, from_stackframe::Bool=false)

src/aotcompile.cpp

@@ -105,7 +105,7 @@ jl_get_llvm_mis_impl(void *native_code, size_t *num_elements, jl_method_instance
     assert(*num_elements == map.size());
     size_t i = 0;
     for (auto &ci : map) {
-        data[i++] = ci.first->def;
+        data[i++] = jl_get_ci_mi(ci.first);
     }
 }
 
@@ -455,14 +455,14 @@ static void compile_workqueue(jl_codegen_params_t &params, egal_set &method_root
                 if ((policy != CompilationPolicy::Default || params.params->trim) &&
                     jl_atomic_load_relaxed(&codeinst->inferred) == jl_nothing) {
                     // XXX: SOURCE_MODE_FORCE_SOURCE is wrong here (neither sufficient nor necessary)
-                    codeinst = jl_type_infer(codeinst->def, jl_atomic_load_relaxed(&codeinst->max_world), SOURCE_MODE_FORCE_SOURCE);
+                    codeinst = jl_type_infer(jl_get_ci_mi(codeinst), jl_atomic_load_relaxed(&codeinst->max_world), SOURCE_MODE_FORCE_SOURCE);
                 }
                 if (codeinst) {
                     orc::ThreadSafeModule result_m =
-                        jl_create_ts_module(name_from_method_instance(codeinst->def),
+                        jl_create_ts_module(name_from_method_instance(jl_get_ci_mi(codeinst)),
                             params.tsctx, params.DL, params.TargetTriple);
                     auto decls = jl_emit_codeinst(result_m, codeinst, NULL, params);
-                    record_method_roots(method_roots, codeinst->def);
+                    record_method_roots(method_roots, jl_get_ci_mi(codeinst));
                     if (result_m)
                         it = compiled_functions.insert(std::make_pair(codeinst, std::make_pair(std::move(result_m), std::move(decls)))).first;
                 }
@@ -501,7 +501,7 @@ static void compile_workqueue(jl_codegen_params_t &params, egal_set &method_root
             proto.decl->setLinkage(GlobalVariable::InternalLinkage);
             //protodecl->setAlwaysInline();
             jl_init_function(proto.decl, params.TargetTriple);
-            jl_method_instance_t *mi = codeinst->def;
+            jl_method_instance_t *mi = jl_get_ci_mi(codeinst);
             size_t nrealargs = jl_nparams(mi->specTypes); // number of actual arguments being passed
             bool is_opaque_closure = jl_is_method(mi->def.value) && mi->def.method->is_for_opaque_closure;
             // TODO: maybe this can be cached in codeinst->specfptr?
@@ -641,12 +641,12 @@ void *jl_create_native_impl(jl_array_t *methods, LLVMOrcThreadSafeModuleRef llvm
                         data->jl_fvar_map[codeinst] = std::make_tuple((uint32_t)-3, (uint32_t)-3);
                     }
                     else {
-                        orc::ThreadSafeModule result_m = jl_create_ts_module(name_from_method_instance(codeinst->def),
+                        orc::ThreadSafeModule result_m = jl_create_ts_module(name_from_method_instance(jl_get_ci_mi(codeinst)),
                                 params.tsctx, clone.getModuleUnlocked()->getDataLayout(),
                                 Triple(clone.getModuleUnlocked()->getTargetTriple()));
                         jl_llvm_functions_t decls = jl_emit_codeinst(result_m, codeinst, NULL, params);
                         JL_GC_PROMISE_ROOTED(codeinst->def); // analyzer seems confused
-                        record_method_roots(method_roots, codeinst->def);
+                        record_method_roots(method_roots, jl_get_ci_mi(codeinst));
                         if (result_m)
                             compiled_functions[codeinst] = {std::move(result_m), std::move(decls)};
                         else if (jl_options.trim != JL_TRIM_NO) {
@@ -2267,7 +2267,7 @@ void jl_get_llvmf_defn_impl(jl_llvmf_dump_t* dump, jl_method_instance_t *mi, jl_
         // To get correct names in the IR this needs to be at least 2
         output.temporary_roots = jl_alloc_array_1d(jl_array_any_type, 0);
         JL_GC_PUSH1(&output.temporary_roots);
-        auto decls = jl_emit_code(m, mi, src, output);
+        auto decls = jl_emit_code(m, mi, src, NULL, output);
         output.temporary_roots = nullptr;
         JL_GC_POP(); // GC the global_targets array contents now since reflection doesn't need it
 

src/builtins.c

@@ -1589,11 +1589,12 @@ JL_CALLABLE(jl_f_invoke)
         return jl_gf_invoke_by_method(m, args[0], &args[2], nargs - 1);
     } else if (jl_is_code_instance(argtypes)) {
         jl_code_instance_t *codeinst = (jl_code_instance_t*)args[1];
+        jl_method_instance_t *mi = jl_get_ci_mi(codeinst);
         jl_callptr_t invoke = jl_atomic_load_acquire(&codeinst->invoke);
         // N.B.: specTypes need not be a subtype of the method signature. We need to check both.
-        if (!jl_tuple1_isa(args[0], &args[2], nargs - 1, (jl_datatype_t*)codeinst->def->specTypes) ||
-            (jl_is_method(codeinst->def->def.value) && !jl_tuple1_isa(args[0], &args[2], nargs - 1, (jl_datatype_t*)codeinst->def->def.method->sig))) {
-            jl_type_error("invoke: argument type error", codeinst->def->specTypes, arg_tuple(args[0], &args[2], nargs - 1));
+        if (!jl_tuple1_isa(args[0], &args[2], nargs - 1, (jl_datatype_t*)mi->specTypes) ||
+            (jl_is_method(mi->def.value) && !jl_tuple1_isa(args[0], &args[2], nargs - 1, (jl_datatype_t*)mi->def.method->sig))) {
+            jl_type_error("invoke: argument type error", mi->specTypes, arg_tuple(args[0], &args[2], nargs - 1));
         }
         if (jl_atomic_load_relaxed(&codeinst->min_world) > jl_current_task->world_age ||
             jl_current_task->world_age > jl_atomic_load_relaxed(&codeinst->max_world)) {
@@ -1609,7 +1610,7 @@ JL_CALLABLE(jl_f_invoke)
             if (codeinst->owner != jl_nothing) {
                 jl_error("Failed to invoke or compile external codeinst");
             }
-            return jl_invoke(args[0], &args[2], nargs - 1, codeinst->def);
+            return jl_invoke(args[0], &args[2], nargs - 1, mi);
         }
     }
     if (!jl_is_tuple_type(jl_unwrap_unionall(argtypes)))