Replace PerDevice cache with GPUToolbox memoization (#268)

Author: maleadt

Project.toml

@@ -15,6 +15,7 @@ CUDA_Tile_jll = "2068806d-a867-5dbd-af0e-42c2eb5d895d"
 CompilerCaching = "9db33cc3-5358-4881-8759-fa4194144afd"
 EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+GPUToolbox = "096a3bc2-3ced-46d0-87f4-dd12716f4bfc"
 IRStructurizer = "93e32bba-5bb8-402b-805d-ffb066edee93"
 LMDB = "11f193de-5e89-5f17-923a-7d207d56daf9"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
@@ -42,6 +43,7 @@ CompilerCaching = "0.2.4"
 DLFP8Types = "0.1.0"
 EnumX = "1.0"
 GPUArrays = "11"
+GPUToolbox = "3"
 IRStructurizer = "0.6.2"
 LMDB = "3"
 Microfloats = "0.2"

src/launch.jl

@@ -6,8 +6,9 @@
 # `(MethodInstance, sm_arch, opt_level, num_ctas, occupancy, num_worker_warps, bytecode_version)`.
 
 using CUDACore: CUDACore, CuArray, CuModule, CuFunction, cudacall, device, capability,
-                AbstractBackend, AbstractKernel, kernel_convert, kernel_compile, PerDevice
+                deviceid, AbstractBackend, AbstractKernel, kernel_convert, kernel_compile
 using CUDA_Compiler_jll
+using GPUToolbox: @memoize
 using Preferences: @load_preference
 
 using Adapt: Adapt, adapt
@@ -176,8 +177,6 @@ function run_and_collect(cmd)
     return proc, log
 end
 
-const tile_ir_support = PerDevice{Union{Nothing, VersionNumber}}()
-
 const toolkit_version_cache =
     Base.Lockable(Base.RefValue{Union{Nothing, Some{Union{Nothing, String}}}}(nothing))
 
@@ -339,7 +338,7 @@ function check_tile_ir_support()
     end
 
     dev = device()
-    ver = get!(tile_ir_support, dev) do
+    ver = @memoize index=deviceid(dev)+1 begin
         ver = bytecode_version()
 
         cap = capability(dev)
@@ -356,7 +355,7 @@ function check_tile_ir_support()
         end
 
         return ver
-    end
+    end::Union{Nothing, VersionNumber}
 
     if ver === nothing
         error("CUDA Tile is not supported on the current device")