Store the actual function in metadata, not just its name.

Author: maleadt

Manifest.toml

@@ -39,15 +39,15 @@ version = "1.4.1"
 
 [[LLVM]]
 deps = ["CEnum", "LLVMExtra_jll", "Libdl", "Printf", "Unicode"]
-git-tree-sha1 = "d7f3db3b5f564016b5dfcc71a20506f8796f403a"
+git-tree-sha1 = "dd58421009014ff1ffacaa0db2a9a392114d75ee"
 uuid = "929cbde3-209d-540e-8aea-75f648917ca0"
-version = "4.10.0"
+version = "4.11.0"
 
 [[LLVMExtra_jll]]
 deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "Pkg", "TOML"]
-git-tree-sha1 = "00d23b26d194507028b9a1c2728a691ab9914262"
+git-tree-sha1 = "771bfe376249626d3ca12bcd58ba243d3f961576"
 uuid = "dad2f222-ce93-54a1-a47d-0025e8a3acab"
-version = "0.0.15+0"
+version = "0.0.16+0"
 
 [[LazyArtifacts]]
 deps = ["Artifacts", "Pkg"]

Project.toml

@@ -14,6 +14,6 @@ UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
 [compat]
 ExprTools = "0.1"
-LLVM = "4.10"
+LLVM = "4.11"
 TimerOutputs = "0.5"
 julia = "1.6"

src/driver.jl

@@ -318,13 +318,9 @@ const __llvm_initialized = Ref(false)
     @timeit_debug to "IR post-processing" begin
         # 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.kernel"], MDNode([MDString(LLVM.name(entry); ctx)]; ctx))
+            push!(metadata(ir)["julia.kernel"], MDNode([entry]; ctx))
 
-            # IDEA: save all jobs, not only top-level kernels, and save other attributes
+            # IDEA: save all jobs, not only kernels, and save other attributes
             #       so that we can reconstruct the CompileJob instead of setting it globally
         end
 

src/irgen.jl

@@ -510,6 +510,7 @@ function lower_byval(@nospecialize(job::CompilerJob), mod::LLVM.Module, f::LLVM.
     # NOTE: if we ever have legitimate uses of the old function, create a shim instead
     fn = LLVM.name(f)
     @assert isempty(uses(f))
+    replace_metadata_uses!(f, new_f)
     unsafe_delete!(mod, f)
     LLVM.name!(new_f, fn)
 
@@ -554,8 +555,7 @@ function add_kernel_state!(mod::LLVM.Module)
     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])
+        push!(kernels, Value(operands(kernel_md)[1]; ctx))
     end
 
     # determine which functions need a kernel state argument
@@ -675,6 +675,7 @@ function add_kernel_state!(mod::LLVM.Module)
                 error("old function still has uses")
             end
         end
+        replace_metadata_uses!(f, workmap[f])
         unsafe_delete!(mod, f)
     end
 

src/ptx.jl

@@ -198,10 +198,6 @@ function finish_ir!(@nospecialize(job::CompilerJob{PTXCompilerTarget}),
     if job.source.kernel
         # add metadata annotations for the assembler to the module
 
-        # NOTE: we do this here, rather than in finish_module!, because otherwise
-        #       the metadata is lost when functions are cloned. See also:
-        #       https://discourse.llvm.org/t/replacing-module-metadata-uses-of-function/62431
-
         # property annotations
         annotations = Metadata[entry]
 

src/spirv.jl

@@ -317,6 +317,7 @@ function wrap_byval(@nospecialize(job::CompilerJob), mod::LLVM.Module, f::LLVM.F
     # NOTE: if we ever have legitimate uses of the old function, create a shim instead
     fn = LLVM.name(f)
     @assert isempty(uses(f))
+    replace_metadata_uses!(f, new_f)
     unsafe_delete!(mod, f)
     LLVM.name!(new_f, fn)