add back argtypes

Author: AntonOresten

ext/CUDAExt.jl

@@ -9,7 +9,7 @@ using CUDA_Compiler_jll
 public launch
 
 # Compilation cache - stores CuFunction directly to avoid re-loading CuModule
-const _compilation_cache = Dict{Any, Any}()  # (method, sm_arch, opt_level) => CuFunction
+const _compilation_cache = Dict{Any, Any}()  # (method, argtypes, sm_arch, opt_level) => CuFunction
 
 """
     launch(f, grid, args...; name=nothing, sm_arch=default_sm_arch(), opt_level=3)
@@ -65,7 +65,7 @@ function cuTile.launch(@nospecialize(f), grid, args...;
     method = which(f, argtypes)
 
     # Check compilation cache - returns CuFunction directly
-    cache_key = (method, sm_arch, opt_level)
+    cache_key = (method, argtypes, sm_arch, opt_level)
     cufunc = get(_compilation_cache, cache_key, nothing)
     if cufunc === nothing || cuTile.compile_hook[] !== nothing
         cubin = compile(f, argtypes; name, sm_arch, opt_level)