From bf017700d04e7fd5dcbe22d4761d8beedf5cf4ac Mon Sep 17 00:00:00 2001
From: Tim Besard <tim.besard@gmail.com>
Date: Mon, 26 Oct 2026 08:19:04 +0100
Subject: [PATCH] Cache the kernel lookup of `hipfunction` with GPUCompiler's
 `cached_lookup`

Every launch took a global lock, built the compiler configuration and job and
queried GPUCompiler's cache, and then looked the `HIPKernel` up in a dictionary
keyed on the kernel's callable. That dictionary kept every callable alive, e.g. a
closure launched over a fresh array added an entry on every launch.

Wrap the lookup in `GPUCompiler.cached_lookup` instead, keyed on the device and
the compiler options, and construct the `HIPKernel` on each launch. A lookup hit
goes from 2892 to 250 instructions, and `@roc` from about 9200 to 6600.
---
 Project.toml                     |  2 +-
 src/compiler/codegen.jl          | 29 ++++++++++++++---------------
 test/kernelabstractions_tests.jl |  3 +++
 3 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/Project.toml b/Project.toml
index 3ccc6f846..9021806a2 100644
--- a/Project.toml
+++ b/Project.toml
@@ -59,7 +59,7 @@ ChainRulesCore = "1"
 EnzymeCore = "0.8"
 ExprTools = "0.1"
 GPUArrays = "12.1"
-GPUCompiler = "2.13"
+GPUCompiler = "2.18"
 GPUToolbox = "3.3.1"
 KernelAbstractions = "0.10"
 KernelInterface = "0.4"
diff --git a/src/compiler/codegen.jl b/src/compiler/codegen.jl
index 3816cd47d..64e9463e7 100644
--- a/src/compiler/codegen.jl
+++ b/src/compiler/codegen.jl
@@ -50,9 +50,6 @@ mutable struct HIPResults
         Tuple{HIP.HIPDevice,HIP.HIPFunction}[])
 end
 
-# (objectid(source), hash(fun), f) => HIPKernel
-const _kernel_instances = Dict{Any, Any}()
-
 GPUCompiler.runtime_module(@nospecialize(::HIPCompilerJob)) = AMDGPU
 
 GPUCompiler.method_table(@nospecialize(::HIPCompilerJob)) = AMDGPU.method_table
@@ -296,19 +293,21 @@ The following kwargs are supported:
     values throw an `ArgumentError`.
 """
 function hipfunction(f::F, tt::TT = Tuple{}; kwargs...) where {F <: Core.Function, TT}
+    kw = values(kwargs)
+    dev = AMDGPU.device()
+    fun = GPUCompiler.cached_lookup(HIP.HIPFunction, F, tt, (dev, kw)) do
+        hipfunction_link(methodinstance(F, tt), dev, kw)
+    end
+    return Runtime.HIPKernel{F, tt}(f, fun)
+end
+
+# the part of `hipfunction` that doesn't depend on the kernel's types, kept out of the code
+# compiled for every kernel
+@noinline function hipfunction_link(source::Core.MethodInstance, dev::HIP.HIPDevice,
+                                    kw::NamedTuple)
     Base.@lock hipfunction_lock begin
-        dev = AMDGPU.device()
-        config = compiler_config(dev; kwargs...)
-        source = methodinstance(F, tt)
-        fun = hipfunction_lookup(source, config, dev)
-
-        key = (objectid(source), hash(fun), f)
-        kernel = get(_kernel_instances, key, nothing)
-        if kernel === nothing
-            kernel = Runtime.HIPKernel{F, tt}(f, fun)
-            _kernel_instances[key] = kernel
-        end
-        return kernel::Runtime.HIPKernel{F, tt}
+        config = compiler_config(dev; kw...)
+        return hipfunction_lookup(source, config, dev)
     end
 end
 
diff --git a/test/kernelabstractions_tests.jl b/test/kernelabstractions_tests.jl
index f77c2b6a9..889d30681 100644
--- a/test/kernelabstractions_tests.jl
+++ b/test/kernelabstractions_tests.jl
@@ -52,7 +52,10 @@ Testsuite.testsuite(
 
     # launching should not be much more expensive than with few arguments
     # (Julia 1.11 and older allocate a little per argument)
+    # warm up in the world of the measurement, as defining this closure advances the world
+    # on Julia 1.10, and the first launch in a new world looks up the kernel again
     Base.invokelatest() do
+        many_args_launch(kernel, out)
         few_args_launch(few_kernel, out)
         @test @allocated(many_args_launch(kernel, out)) <=
               @allocated(few_args_launch(few_kernel, out)) + 40*32