Support GPUArrays allocations cache (#2593)

Co-authored-by: Tim Besard <[email protected]>
JuliaGPU · Jan 9, 2025 · 1389800 · 1389800
1 parent 73b28c8
commit 1389800
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/Project.toml b/Project.toml
@@ -59,7 +59,7 @@ Crayons = "4"
 DataFrames = "1"
 EnzymeCore = "0.8.2"
 ExprTools = "0.1"
-GPUArrays = "11.1"
+GPUArrays = "11.2"
 GPUCompiler = "0.24, 0.25, 0.26, 0.27, 1"
 KernelAbstractions = "0.9.2"
 LLVM = "9.1"

diff --git a/src/array.jl b/src/array.jl
@@ -71,9 +71,13 @@ mutable struct CuArray{T,N,M} <: AbstractGPUArray{T,N}
     else
       maxsize
     end
-    data = DataRef(pool_free, pool_alloc(M, bufsize))
-    obj = new{T,N,M}(data, maxsize, 0, dims)
-    finalizer(unsafe_free!, obj)
+
+    GPUArrays.cached_alloc((CuArray, CUDA.device(), T, bufsize, M)) do
+        data = DataRef(pool_free, pool_alloc(M, bufsize))
+        obj = new{T,N,M}(data, maxsize, 0, dims)
+        finalizer(unsafe_free!, obj)
+        return obj
+    end::CuArray{T, N, M}
   end
 
   function CuArray{T,N}(data::DataRef{Managed{M}}, dims::Dims{N};