Set kernel metadata for all kernels, including deferred ones.

Author: maleadt

src/driver.jl

@@ -316,15 +316,16 @@ const __llvm_initialized = Ref(false)
     end
 
     @timeit_debug to "IR post-processing" begin
-        # mark the entry-point function (optimization may need it)
-        if deferred_codegen
-            # IDEA: save other parts of the CompileJob (so that we can reconstruct it
-            #       instead of setting it globally, which is incompatible with threading)?
+        # mark the kernel entry-point functions (optimization may need it)
+        if job.source.kernel
             # XXX: we want to save the actual function here, but due to our passes rewriting
             #      functions, and the inability to RAUW values with a different type, that
             #      metadata gets lost. So instead we save the function name. See also:
             #      https://discourse.llvm.org/t/replacing-module-metadata-uses-of-function/62431
-            push!(metadata(ir)["julia.entry"], MDNode([MDString(LLVM.name(entry); ctx)]; ctx))
+            push!(metadata(ir)["julia.kernel"], MDNode([MDString(LLVM.name(entry); ctx)]; ctx))
+
+            # IDEA: save all jobs, not only top-level kernels, and save other attributes
+            #       so that we can reconstruct the CompileJob instead of setting it globally
         end
 
         if optimize

src/irgen.jl

@@ -551,16 +551,19 @@ function add_kernel_state!(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_fn = string(operands(entry_md)[1])
-    entry = functions(mod)[entry_fn]
+    kernels = []
+    kernels_md = metadata(mod)["julia.kernel"]
+    for kernel_md in operands(kernels_md)
+        kernel_fn = string(operands(kernel_md)[1])
+        push!(kernels, functions(mod)[kernel_fn])
+    end
 
     # determine which functions need a kernel state argument
     #
     # previously, we add the argument to every function and relied on unused arg elim to
     # clean-up the IR. however, some libraries do Funny Stuff, e.g., libdevice bitcasting
     # function pointers. such IR is hard to rewrite, so instead be more conservative.
-    worklist = Set{LLVM.Function}([entry, state_intr])
+    worklist = Set{LLVM.Function}([state_intr, kernels...])
     worklist_length = 0
     while worklist_length != length(worklist)
         # iteratively discover functions that use the intrinsic or any function calling it