Add support for calling runtime functions post-optimization.

Author: maleadt

src/rtlib.jl

@@ -28,7 +28,8 @@ end
 
 ## higher-level functionality to work with runtime functions
 
-function LLVM.call!(builder, rt::Runtime.RuntimeMethodInstance, args=LLVM.Value[])
+function LLVM.call!(builder, rt::Runtime.RuntimeMethodInstance, args=LLVM.Value[];
+                    state::Type=Nothing)
     bb = position(builder)
     f = LLVM.parent(bb)
     mod = LLVM.parent(f)
@@ -39,10 +40,18 @@ function LLVM.call!(builder, rt::Runtime.RuntimeMethodInstance, args=LLVM.Value[
         f = functions(mod)[rt.llvm_name]
         ft = eltype(llvmtype(f))
     else
-        ft = convert(LLVM.FunctionType, rt; ctx)
+        ft = convert(LLVM.FunctionType, rt; ctx, state)
         f = LLVM.Function(mod, rt.llvm_name, ft)
     end
 
+    # we may be calling this function after kernel state lowering,
+    # in which case we need to manually get and pass the state.
+    args = Value[args...]
+    if state !== Nothing
+        state_val = kernel_state_argument(f, state)
+        pushfirst!(args, state_val)
+    end
+
     # runtime functions are written in Julia, while we're calling from LLVM,
     # this often results in argument type mismatches. try to fix some here.
     for (i,arg) in enumerate(args)

src/runtime.jl

@@ -33,13 +33,21 @@ struct RuntimeMethodInstance
     llvm_name::String
 end
 
-function Base.convert(::Type{LLVM.FunctionType}, rt::RuntimeMethodInstance; ctx::LLVM.Context)
+function Base.convert(::Type{LLVM.FunctionType}, rt::RuntimeMethodInstance;
+                      ctx::LLVM.Context, state::Type=Nothing)
     types = if rt.llvm_types === nothing
         LLVMType[convert(LLVMType, typ; ctx, allow_boxed=true) for typ in rt.types]
     else
         rt.llvm_types(ctx)
     end
 
+    # if we're running post-optimization, prepend the kernel state to the argument list
+    if state !== Nothing
+        T_state = convert(LLVMType, state; ctx)
+        T_ptr_state = LLVM.PointerType(T_state)
+        pushfirst!(types, T_ptr_state)
+    end
+
     return_type = if rt.llvm_return_type === nothing
         convert(LLVMType, rt.return_type; ctx, allow_boxed=true)
     else