Support func entry abi (#287)

Co-authored-by: Julian Samaroo <[email protected]>

Author: vchuravy

src/driver.jl

@@ -203,7 +203,11 @@ const __llvm_initialized = Ref(false)
 
     @timeit_debug to "IR generation" begin
         ir, compiled = irgen(job, method_instance; ctx)
-        entry_fn = compiled[method_instance].specfunc
+        if job.entry_abi === :specfunc
+            entry_fn = compiled[method_instance].specfunc
+        else
+            entry_fn = compiled[method_instance].func
+        end
         entry = functions(ir)[entry_fn]
     end
 

src/interface.jl

@@ -117,17 +117,37 @@ export CompilerJob
 
 # a specific invocation of the compiler, bundling everything needed to generate code
 
+"""
+    CompilerJob(target, source, params, entry_abi)
+
+Construct a `CompilerJob` for `source` that will be used to drive compilation for
+the given `target` and `params`. The `entry_abi` can be either `:specfunc` the default,
+or `:func`. `:specfunc` expects the arguments to be passed in registers, simple
+return values are returned in registers as well, and complex return values are returned
+on the stack using `sret`, the calling convention is `fastcc`. The `:func` abi is simpler
+with a calling convention of the first argument being the function itself (to support closures),
+the second argument being a pointer to a vector of boxed Julia values and the third argument
+being the number of values, the return value will also be boxed. The `:func` abi
+will internally call the `:specfunc` abi, but is generally easier to invoke directly.
+"""
 struct CompilerJob{T,P,F}
     target::T
     source::F
     params::P
+    entry_abi::Symbol
 
-    CompilerJob(target::AbstractCompilerTarget, source::FunctionSpec, params::AbstractCompilerParams) =
-        new{typeof(target), typeof(params), typeof(source)}(target, source, params)
+    function CompilerJob(target::AbstractCompilerTarget, source::FunctionSpec, params::AbstractCompilerParams, entry_abi::Symbol)
+        if entry_abi ∉ (:specfunc, :func)
+            error("Unknown entry_abi=$entry_abi")
+        end
+        new{typeof(target), typeof(params), typeof(source)}(target, source, params, entry_abi)
+    end
 end
+CompilerJob(target::AbstractCompilerTarget, source::FunctionSpec, params::AbstractCompilerParams; entry_abi=:specfunc) =
+    CompilerJob(target, source, params, entry_abi)
 
 Base.similar(@nospecialize(job::CompilerJob), @nospecialize(source::FunctionSpec)) =
-    CompilerJob(job.target, source, job.params)
+    CompilerJob(job.target, source, job.params, job.entry_abi)
 
 function Base.show(io::IO, @nospecialize(job::CompilerJob{T})) where {T}
     print(io, "CompilerJob of ", job.source, " for ", T)
@@ -137,6 +157,7 @@ function Base.hash(job::CompilerJob, h::UInt)
     h = hash(job.target, h)
     h = hash(job.source, h)
     h = hash(job.params, h)
+    h = hash(job.entry_abi, h)
     h
 end
 

src/irgen.jl

@@ -3,7 +3,11 @@
 function irgen(@nospecialize(job::CompilerJob), method_instance::Core.MethodInstance;
                ctx::Context)
     mod, compiled = @timeit_debug to "emission" compile_method_instance(job, method_instance; ctx)
-    entry_fn = compiled[method_instance].specfunc
+    if job.entry_abi === :specfunc
+        entry_fn = compiled[method_instance].specfunc
+    else
+        entry_fn = compiled[method_instance].func
+    end
 
     # clean up incompatibilities
     @timeit_debug to "clean-up" begin
@@ -16,8 +20,11 @@ function irgen(@nospecialize(job::CompilerJob), method_instance::Core.MethodInst
             end
 
             # remove the non-specialized jfptr functions
-            if startswith(LLVM.name(llvmf), "jfptr_")
-                unsafe_delete!(mod, llvmf)
+            # TODO: Do we need to remove these?
+            if job.entry_abi === :specfunc
+                if startswith(LLVM.name(llvmf), "jfptr_")
+                    unsafe_delete!(mod, llvmf)
+                end
             end
         end
 
@@ -55,9 +62,16 @@ function irgen(@nospecialize(job::CompilerJob), method_instance::Core.MethodInst
         LLVM.name!(entry, mangle_call(entry, job.source.tt))
     end
     entry = process_entry!(job, mod, entry)
+    if job.entry_abi === :specfunc
+        func = compiled[method_instance].func
+        specfunc = LLVM.name(entry)
+    else
+        func = LLVM.name(entry)
+        specfunc = compiled[method_instance].specfunc
+    end
+
     compiled[method_instance] =
-        (; compiled[method_instance].ci, compiled[method_instance].func,
-           specfunc=LLVM.name(entry))
+        (; compiled[method_instance].ci, func, specfunc)
 
     # minimal required optimization
     @timeit_debug to "rewrite" ModulePassManager() do pm

src/reflection.jl

@@ -102,10 +102,10 @@ The following keyword arguments are supported:
 See also: [`@device_code_llvm`](@ref), `InteractiveUtils.code_llvm`
 """
 function code_llvm(io::IO, @nospecialize(job::CompilerJob); optimize::Bool=true, raw::Bool=false,
-                   debuginfo::Symbol=:default, dump_module::Bool=false)
+                   debuginfo::Symbol=:default, dump_module::Bool=false, kwargs...)
     # NOTE: jl_dump_function_ir supports stripping metadata, so don't do it in the driver
     str = JuliaContext() do ctx
-        ir, meta = codegen(:llvm, job; optimize=optimize, strip=false, validate=false, ctx)
+        ir, meta = codegen(:llvm, job; optimize=optimize, strip=false, validate=false, ctx, kwargs...)
         ccall(:jl_dump_function_ir, Ref{String},
               (LLVM.API.LLVMValueRef, Bool, Bool, Ptr{UInt8}),
               meta.entry, !raw, dump_module, debuginfo)

test/definitions/native.jl

@@ -18,11 +18,11 @@ end
 
 GPUCompiler.method_table(@nospecialize(job::NativeCompilerJob)) = method_table
 
-function native_job(@nospecialize(func), @nospecialize(types); kernel::Bool=false, kwargs...)
+function native_job(@nospecialize(func), @nospecialize(types); kernel::Bool=false, entry_abi=:specfunc, kwargs...)
     source = FunctionSpec(func, Base.to_tuple_type(types), kernel)
     target = NativeCompilerTarget(always_inline=true)
     params = TestCompilerParams()
-    CompilerJob(target, source, params), kwargs
+    CompilerJob(target, source, params, entry_abi), kwargs
 end
 
 function native_code_typed(@nospecialize(func), @nospecialize(types); kwargs...)

test/native.jl

@@ -135,6 +135,15 @@ end
     native_code_llvm(devnull, D32593, Tuple{Ptr{D32593_struct}})
 end
 
+@testset "slow abi" begin
+    x = 2
+    f = () -> x+1
+    ir = sprint(io->native_code_llvm(io, f, Tuple{}, entry_abi=:func, dump_module=true))
+    @test occursin(r"define nonnull {}\* @jfptr", ir)
+    @test occursin(r"define internal fastcc .+ @julia", ir)
+    @test occursin(r"call fastcc .+ @julia", ir)
+end
+
 end
 
 ############################################################################################