Skip to content

Commit b6bf8da

Browse files
authored
Alloc cache: Add type assertion. (#581)
It assumes caching array values right now.
1 parent 64482de commit b6bf8da

File tree

1 file changed

+3
-3
lines changed

1 file changed

+3
-3
lines changed

src/host/alloc_cache.jl

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -33,15 +33,15 @@ end
3333
function cached_alloc(f, key)
3434
cache = ALLOC_CACHE[]
3535
if cache === nothing
36-
return f()
36+
return f()::AbstractGPUArray
3737
end
3838

3939
x = nothing
4040
uid = hash(key)
4141

4242
busy_pool = get_pool!(cache, :busy, uid)
4343
free_pool = get_pool!(cache, :free, uid)
44-
isempty(free_pool) && (x = f())
44+
isempty(free_pool) && (x = f()::AbstractGPUArray)
4545

4646
while !isempty(free_pool) && x nothing
4747
tmp = Base.@lock cache.lock pop!(free_pool)
@@ -50,7 +50,7 @@ function cached_alloc(f, key)
5050
x = tmp
5151
end
5252

53-
x nothing && (x = f())
53+
x nothing && (x = f()::AbstractGPUArray)
5454
Base.@lock cache.lock push!(busy_pool, x)
5555
return x
5656
end

0 commit comments

Comments
 (0)