Alloc cache: Add type assertion. (#581)

It assumes caching array values right now.
JuliaGPU · Jan 9, 2025 · b6bf8da · b6bf8da
1 parent 64482de
commit b6bf8da
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/host/alloc_cache.jl b/src/host/alloc_cache.jl
@@ -33,15 +33,15 @@ end
 function cached_alloc(f, key)
     cache = ALLOC_CACHE[]
     if cache === nothing
-        return f()
+        return f()::AbstractGPUArray
     end
 
     x = nothing
     uid = hash(key)
 
     busy_pool = get_pool!(cache, :busy, uid)
     free_pool = get_pool!(cache, :free, uid)
-    isempty(free_pool) && (x = f())
+    isempty(free_pool) && (x = f()::AbstractGPUArray)
 
     while !isempty(free_pool) && x ≡ nothing
         tmp = Base.@lock cache.lock pop!(free_pool)
@@ -50,7 +50,7 @@ function cached_alloc(f, key)
         x = tmp
     end
 
-    x ≡ nothing && (x = f())
+    x ≡ nothing && (x = f()::AbstractGPUArray)
     Base.@lock cache.lock push!(busy_pool, x)
     return x
 end