Update to new alloc cache interface. (#2614)

Author: maleadt

Project.toml

@@ -59,7 +59,7 @@ Crayons = "4"
 DataFrames = "1"
 EnzymeCore = "0.8.2"
 ExprTools = "0.1"
-GPUArrays = "11.2"
+GPUArrays = "11.2.1"
 GPUCompiler = "0.24, 0.25, 0.26, 0.27, 1"
 KernelAbstractions = "0.9.2"
 LLVM = "9.1"

src/array.jl

@@ -72,19 +72,20 @@ mutable struct CuArray{T,N,M} <: AbstractGPUArray{T,N}
       maxsize
     end
 
-    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}
+    data = GPUArrays.cached_alloc((CuArray, device(), M, bufsize)) do
+        DataRef(pool_free, pool_alloc(M, bufsize))
+    end
+    obj = new{T,N,M}(data, maxsize, 0, dims)
+    finalizer(unsafe_free!, obj)
+    return obj
   end
 
   function CuArray{T,N}(data::DataRef{Managed{M}}, dims::Dims{N};
                         maxsize::Int=prod(dims) * sizeof(T), offset::Int=0) where {T,N,M}
     check_eltype(T)
     obj = new{T,N,M}(data, maxsize, offset, dims)
     finalizer(unsafe_free!, obj)
+    return obj
   end
 end
 

src/memory.jl

@@ -516,6 +516,8 @@ mutable struct Managed{M}
   end
 end
 
+Base.sizeof(managed::Managed) = sizeof(managed.mem)
+
 # wait for the current owner of memory to finish processing
 function synchronize(managed::Managed)
   synchronize(managed.stream)