Add helpers for iterating kernels in a module.

Author: maleadt

src/driver.jl

@@ -322,10 +322,7 @@ const __llvm_initialized = Ref(false)
     @tracepoint "IR post-processing" begin
         # mark the kernel entry-point functions (optimization may need it)
         if job.config.kernel
-            push!(metadata(ir)["julia.kernel"], MDNode([entry]))
-
-            # IDEA: save all jobs, not only kernels, and save other attributes
-            #       so that we can reconstruct the CompileJob instead of setting it globally
+            mark_kernel!(entry)
         end
 
         if job.config.toplevel

src/irgen.jl

@@ -526,18 +526,12 @@ function add_kernel_state!(mod::LLVM.Module)
     state_intr = kernel_state_intr(mod, T_state)
     state_intr_ft = LLVM.FunctionType(T_state)
 
-    kernels = []
-    kernels_md = metadata(mod)["julia.kernel"]
-    for kernel_md in operands(kernels_md)
-        push!(kernels, Value(operands(kernel_md)[1]))
-    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}([state_intr, kernels...])
+    worklist = Set{LLVM.Function}([state_intr, kernels(mod)...])
     worklist_length = 0
     while worklist_length != length(worklist)
         # iteratively discover functions that use the intrinsic or any function calling it

src/metal.jl

@@ -48,15 +48,10 @@ isintrinsic(@nospecialize(job::CompilerJob{MetalCompilerTarget}), fn::String) =
     return startswith(fn, "air.")
 
 function finish_linked_module!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::LLVM.Module)
-    if haskey(metadata(mod), "julia.kernel")
-        kernels_md = metadata(mod)["julia.kernel"]
-        for kernel_md in operands(kernels_md)
-            f = LLVM.Value(operands(kernel_md)[1])::LLVM.Function
-
-            # update calling conventions
-            f = pass_by_reference!(job, mod, f)
-            f = add_input_arguments!(job, mod, f, kernel_intrinsics)
-        end
+    for f in kernels(mod)
+        # update calling conventions
+        f = pass_by_reference!(job, mod, f)
+        f = add_input_arguments!(job, mod, f, kernel_intrinsics)
     end
 
     # emit the AIR and Metal version numbers as constants in the module. this makes it

src/utils.jl

@@ -155,3 +155,30 @@ function prune_constexpr_uses!(root::LLVM.Value)
         end
     end
 end
+
+
+## kernel metadata handling
+
+# kernels are encoded in the IR using the julia.kernel metadata.
+
+# IDEA: don't only mark kernels, but all jobs, and save all attributes of the CompileJob
+#       so that we can reconstruct the CompileJob instead of setting it globally
+
+# mark a function as kernel
+function mark_kernel!(f::LLVM.Function)
+    mod = LLVM.parent(f)
+    push!(metadata(mod)["julia.kernel"], MDNode([f]))
+    return f
+end
+
+# iterate over all kernels in the module
+function kernels(mod::LLVM.Module)
+    vals = LLVM.Function[]
+    if haskey(metadata(mod), "julia.kernel")
+        kernels_md = metadata(mod)["julia.kernel"]
+        for kernel_md in operands(kernels_md)
+            push!(vals, LLVM.Value(operands(kernel_md)[1]))
+        end
+    end
+    return vals
+end