Remove kernel state special casing now that the transformation happens late.

Author: maleadt

src/irgen.jl

@@ -404,21 +404,14 @@ function lower_byval(@nospecialize(job::CompilerJob), mod::LLVM.Module, f::LLVM.
         # XXX: byval is not round-trippable on LLVM < 12 (see maleadt/LLVM.jl#186)
         #      so we need to re-classify the Julia arguments.
         #      remove this once we only support 1.7.
-        has_kernel_state = kernel_state_type(job) !== Nothing
-        orig_ft = if has_kernel_state
-            # the kernel state has been added here already, so strip the first parameter
-            LLVM.FunctionType(LLVM.return_type(ft), parameters(ft)[2:end]; vararg=isvararg(ft))
-        else
-            ft
-        end
-        args = classify_arguments(job, orig_ft)
+        args = classify_arguments(job, ft)
         filter!(args) do arg
             arg.cc != GHOST
         end
         for arg in args
             if arg.cc == BITS_REF
                 # NOTE: +1 since this pass runs after introducing the kernel state
-                byval[arg.codegen.i+has_kernel_state] = true
+                byval[arg.codegen.i] = true
             end
         end
     end

src/optim.jl

@@ -255,8 +255,7 @@ function optimize!(@nospecialize(job::CompilerJob), mod::LLVM.Module)
     ModulePassManager() do pm
         addTargetPasses!(pm, tm, triple)
 
-        # - remove unused kernel state arguments
-        # - simplify function calls that don't use the returned value
+        # simplify function calls that don't use the returned value
         dead_arg_elimination!(pm)
 
         run!(pm, mod)
@@ -356,16 +355,8 @@ function lower_gc_frame!(fun::LLVM.Function)
 
             # replace with PTX alloc_obj
             Builder(ctx) do builder
-                # NOTE: this happens late during the pipeline, where we may have to
-                #       pass a kernel state arguments to the runtime function.
-                state = if job.source.kernel
-                    kernel_state_type(job)
-                else
-                    Nothing
-                end
-
                 position!(builder, call)
-                ptr = call!(builder, Runtime.get(:gc_pool_alloc), [sz]; state)
+                ptr = call!(builder, Runtime.get(:gc_pool_alloc), [sz])
                 replace_uses!(call, ptr)
             end
 

src/rtlib.jl

@@ -28,8 +28,7 @@ end
 
 ## higher-level functionality to work with runtime functions
 
-function LLVM.call!(builder, rt::Runtime.RuntimeMethodInstance, args=LLVM.Value[];
-                    state::Type=Nothing)
+function LLVM.call!(builder, rt::Runtime.RuntimeMethodInstance, args=LLVM.Value[])
     bb = position(builder)
     f = LLVM.parent(bb)
     mod = LLVM.parent(f)
@@ -40,21 +39,10 @@ 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, state)
+        ft = convert(LLVM.FunctionType, rt; ctx)
         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
-        T_state = convert(LLVMType, state; ctx)
-
-        state_intr = kernel_state_intr(mod, T_state)
-        state_val = call!(builder, state_intr, Value[], "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

@@ -34,19 +34,13 @@ struct RuntimeMethodInstance
 end
 
 function Base.convert(::Type{LLVM.FunctionType}, rt::RuntimeMethodInstance;
-                      ctx::LLVM.Context, state::Type=Nothing)
+                      ctx::LLVM.Context)
     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)
-        pushfirst!(types, T_state)
-    end
-
     return_type = if rt.llvm_return_type === nothing
         convert(LLVMType, rt.return_type; ctx, allow_boxed=true)
     else

src/spirv.jl

@@ -232,26 +232,15 @@ function wrap_byval(@nospecialize(job::CompilerJob), mod::LLVM.Module, f::LLVM.F
         end
     else
         # XXX: byval is not round-trippable on LLVM < 12 (see maleadt/LLVM.jl#186)
-        has_kernel_state = kernel_state_type(job) !== Nothing
-        orig_ft = if has_kernel_state
-            # the kernel state has been added here already, so strip the first parameter
-            LLVM.FunctionType(LLVM.return_type(ft), parameters(ft)[2:end]; vararg=isvararg(ft))
-        else
-            ft
-        end
-        args = classify_arguments(job, orig_ft)
+        args = classify_arguments(job, ft)
         filter!(args) do arg
             arg.cc != GHOST
         end
         for arg in args
             if arg.cc == BITS_REF
-                # NOTE: +1 since this pass runs after introducing the kernel state
-                byval[arg.codegen.i+has_kernel_state] = true
+                byval[arg.codegen.i] = true
             end
         end
-        if has_kernel_state
-            byval[1] = true
-        end
     end
 
     # generate the wrapper function type & definition