Convert initial kernel state transformation to a late pass.

Author: maleadt

src/interface.jl

@@ -233,15 +233,7 @@ optimize_module!(@nospecialize(job::CompilerJob), mod::LLVM.Module) = return
 
 # finalization of the module, before deferred codegen and optimization
 function finish_module!(@nospecialize(job::CompilerJob), mod::LLVM.Module, entry::LLVM.Function)
-    ctx = context(mod)
-    entry_fn = LLVM.name(entry)
-
-    # add the kernel state, and lower calls to the `julia.gpu.state_getter` intrinsic.
-    if job.source.kernel
-        add_kernel_state!(job, mod, entry)
-    end
-
-    return functions(mod)[entry_fn]
+    return entry
 end
 
 # final processing of the IR, right before validation and machine-code generation

src/irgen.jl

@@ -535,10 +535,9 @@ end
 # so that the julia.gpu.state_getter` can be simplified to return an opaque pointer.
 
 # add a state argument to every function in the module, starting from the kernel entry point
-function add_kernel_state!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
-                           entry::LLVM.Function)
+function add_kernel_state!(mod::LLVM.Module)
+    job = current_job::CompilerJob
     ctx = context(mod)
-    entry_fn = LLVM.name(entry)
 
     # check if we even need a kernel state argument
     state = kernel_state_type(job)
@@ -552,6 +551,11 @@ function add_kernel_state!(@nospecialize(job::CompilerJob), mod::LLVM.Module,
     # this is both for extern uses, and to make this transformation a two-step process.
     state_intr = kernel_state_intr(mod, T_state)
 
+    entry_md = operands(metadata(mod)["julia.entry"])[1]
+    entry = Value(operands(entry_md)[1]; ctx)
+    # XXX: this metadata will be invalid after the replacement here (it'll be null).
+    #      how do we replace Metadata uses? Normally RAUW, but it asserts type equality
+
     # determine which functions need a kernel state argument
     #
     # previously, we add the argument to every function and relied on unused arg elim to

src/optim.jl

@@ -192,6 +192,8 @@ function optimize!(@nospecialize(job::CompilerJob), mod::LLVM.Module)
         if job.source.kernel
             # GC lowering is the last pass that may introduce calls to the runtime library,
             # and thus additional uses of the kernel state intrinsic.
+            # TODO: now that all kernel state-related passes are being run here, merge some?
+            add!(pm, ModulePass("AddKernelState", add_kernel_state!))
             add!(pm, FunctionPass("LowerKernelState", lower_kernel_state!))
             add!(pm, ModulePass("CleanupKernelState", cleanup_kernel_state!))
         end