From 93e1f440f3e2d622c09491ac2691df0d87f66d1d Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 9 Jan 2025 14:41:17 +0100 Subject: [PATCH] Alloc cache: Add type assertion. It assumes caching array values right now. --- src/host/alloc_cache.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/host/alloc_cache.jl b/src/host/alloc_cache.jl index 7ed435de..e2143efb 100644 --- a/src/host/alloc_cache.jl +++ b/src/host/alloc_cache.jl @@ -33,7 +33,7 @@ end function cached_alloc(f, key) cache = ALLOC_CACHE[] if cache === nothing - return f() + return f()::AbstractGPUArray end x = nothing @@ -41,7 +41,7 @@ function cached_alloc(f, 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