Skip to content

Commit d0115fd

Browse files
pxl-thpxl-th
committed
Remove 'free_immediately' param
1 parent d6a74b0 commit d0115fd

File tree

4 files changed

+61
-101
lines changed

4 files changed

+61
-101
lines changed

‎lib/JLArrays/src/JLArrays.jl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -397,7 +397,7 @@ Adapt.adapt_storage(::KernelAbstractions.CPU, a::JLArrays.JLArray) = convert(Arr
397397

398398
# Caching Allocator.
399399

400-
const JLACacheAllocator = GPUArrays.AllocCache.PerDeviceCacheAllocator(JLArray; free_immediately=false)
400+
const JLACacheAllocator = GPUArrays.AllocCache.PerDeviceCacheAllocator(JLArray)
401401

402402
GPUArrays.AllocCache.cache_allocator(::Type{<: JLArray}) = JLACacheAllocator
403403

‎src/host/allocations_cache.jl‎

Lines changed: 17 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -33,53 +33,39 @@ function get_pool!(cache::CacheAllocator{T}, pool::Symbol, uid::UInt64) where T
3333
end
3434

3535
"""
36-
alloc!(alloc_f, cache::CacheAllocator, key; skip_free::Bool)
36+
alloc!(alloc_f, cache::CacheAllocator, key)
3737
3838
Attempt to retrieve cached allocation from `cache` using `key` for searching.
3939
If no such allocation is found, execute `alloc_f` that does actual allocation,
4040
store it in cache for future use and return it.
41-
42-
`skip_free::Bool` is used together with `PerDeviceCacheAllocator.free_immediately`.
43-
When `true` arrays are bulk-freed instead of stored in cache.
44-
In this case `alloc!` will avoid looking into "free" part of `cache`
45-
and execute `alloc_f` immediately, storing allocation for future bulk-freeing.
4641
"""
47-
function alloc!(alloc_f, cache::CacheAllocator, key; skip_free::Bool)
42+
function alloc!(alloc_f, cache::CacheAllocator, key)
4843
x = nothing
4944
uid = hash(key)
5045
busy_pool = get_pool!(cache, :busy, uid)
51-
52-
if skip_free
53-
x = alloc_f()
54-
else
55-
free_pool = get_pool!(cache, :free, uid)
56-
isempty(free_pool) && (x = alloc_f())
57-
58-
while !isempty(free_pool) && x ≡ nothing
59-
tmp = Base.@lock cache.lock pop!(free_pool)
60-
# Array was manually freed via `unsafe_free!`.
61-
GPUArrays.storage(tmp).freed && continue
62-
x = tmp
63-
end
46+
free_pool = get_pool!(cache, :free, uid)
47+
isempty(free_pool) && (x = alloc_f())
48+
49+
while !isempty(free_pool) && x ≡ nothing
50+
tmp = Base.@lock cache.lock pop!(free_pool)
51+
# Array was manually freed via `unsafe_free!`.
52+
GPUArrays.storage(tmp).freed && continue
53+
x = tmp
6454
end
6555

6656
x ≡ nothing && (x = alloc_f())
6757
Base.@lock cache.lock push!(busy_pool, x)
6858
return x
6959
end
7060

71-
function free_busy!(cache::CacheAllocator; free_immediately::Bool)
61+
function free_busy!(cache::CacheAllocator)
7262
for uid in cache.busy.keys
7363
busy_pool = get_pool!(cache, :busy, uid)
7464
isempty(busy_pool) && continue
7565

7666
Base.@lock cache.lock begin
77-
if free_immediately
78-
map(unsafe_free!, busy_pool)
79-
else
80-
free_pool = get_pool!(cache, :free, uid)
81-
append!(free_pool, busy_pool)
82-
end
67+
free_pool = get_pool!(cache, :free, uid)
68+
append!(free_pool, busy_pool)
8369
empty!(busy_pool)
8470
end
8571
end
@@ -88,11 +74,10 @@ end
8874
mutable struct PerDeviceCacheAllocator{T <: AbstractGPUArray}
8975
lock::ReentrantLock
9076
caches::Dict{UInt64, Dict{Symbol, CacheAllocator{T}}}
91-
free_immediately::Bool
9277
end
9378

94-
PerDeviceCacheAllocator(::Type{T}; free_immediately::Bool) where T <: AbstractGPUArray =
95-
PerDeviceCacheAllocator(ReentrantLock(), Dict{UInt64, Dict{Symbol, CacheAllocator{T}}}(), free_immediately)
79+
PerDeviceCacheAllocator(::Type{T}) where T <: AbstractGPUArray =
80+
PerDeviceCacheAllocator(ReentrantLock(), Dict{UInt64, Dict{Symbol, CacheAllocator{T}}}())
9681

9782
function named_cache_allocator!(pdcache::PerDeviceCacheAllocator{T}, device, name::Symbol) where T
9883
h = hash(device)
@@ -116,7 +101,7 @@ end
116101
function alloc!(alloc_f, AT::Type{<: AbstractGPUArray}, name::Symbol, key)
117102
pdcache = cache_allocator(AT)
118103
cache = named_cache_allocator!(pdcache, device(AT), name)
119-
alloc!(alloc_f, cache, key; skip_free=pdcache.free_immediately)
104+
alloc!(alloc_f, cache, key)
120105
end
121106

122107
function Base.sizeof(pdcache::PerDeviceCacheAllocator, device, name::Symbol)
@@ -172,7 +157,7 @@ end
172157

173158
function free_busy!(AT::Type{<: AbstractGPUArray}, name::Symbol)
174159
pdcache = cache_allocator(AT)
175-
free_busy!(named_cache_allocator!(pdcache, device(AT), name); pdcache.free_immediately)
160+
free_busy!(named_cache_allocator!(pdcache, device(AT), name))
176161
end
177162

178163
"""

‎test/testsuite.jl‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -93,8 +93,7 @@ include("testsuite/math.jl")
9393
include("testsuite/random.jl")
9494
include("testsuite/uniformscaling.jl")
9595
include("testsuite/statistics.jl")
96-
# TODO re-enable once backends support it.
97-
# include("testsuite/caching_allocator.jl")
96+
include("testsuite/caching_allocator.jl")
9897

9998
"""
10099
Runs the entire GPUArrays test suite on array type `AT`

‎test/testsuite/caching_allocator.jl‎

Lines changed: 42 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -1,80 +1,56 @@
11
@testsuite "Caching Allocator" (AT, eltypes) -> begin
22
device = GPUArrays.AllocCache.device(AT)
3+
pdcache = GPUArrays.AllocCache.cache_allocator(AT)
4+
named_cache = GPUArrays.AllocCache.named_cache_allocator!(pdcache, device, :cache)
35

4-
@testset "free_immediately=false" begin
5-
pdcache = GPUArrays.AllocCache.cache_allocator(AT)
6-
pdcache.free_immediately = false
7-
named_cache = GPUArrays.AllocCache.named_cache_allocator!(pdcache, device, :cache)
6+
T = Float32
7+
dims = (1, 2, 3)
88

9-
T = Float32
10-
dims = (1, 2, 3)
11-
key = hash((T, dims))
9+
GPUArrays.AllocCache.@enable AT :cache begin
10+
x1 = AT(zeros(T, dims))
11+
end
12+
@test sizeof(pdcache, device, :cache) == sizeof(Float32) * prod(dims)
13+
@test length(named_cache.free) == 1
1214

13-
GPUArrays.AllocCache.@enable AT :cache begin
14-
x1 = AT(zeros(T, dims))
15-
end
16-
@test sizeof(pdcache, device, :cache) == sizeof(Float32) * prod(dims)
17-
@test length(named_cache.free[key]) == 1
18-
@test length(named_cache.busy[key]) == 0
19-
@test x1 === named_cache.free[key][1]
15+
key = first(keys(named_cache.free))
16+
@test length(named_cache.free[key]) == 1
17+
@test length(named_cache.busy[key]) == 0
18+
@test x1 === named_cache.free[key][1]
2019

21-
# Second allocation does not allocate - cache stays the same in size.
20+
# Second allocation does not allocate - cache stays the same in size.
2221

23-
GPUArrays.AllocCache.@enable AT :cache begin
24-
x2 = AT(zeros(T, dims))
22+
GPUArrays.AllocCache.@enable AT :cache begin
23+
x2 = AT(zeros(T, dims))
2524

26-
# Does not go to cache.
27-
GPUArrays.AllocCache.@disable begin
28-
x_free = AT(zeros(T, dims))
29-
end
25+
# Does not go to cache.
26+
GPUArrays.AllocCache.@disable begin
27+
x_free = AT(zeros(T, dims))
3028
end
31-
@test sizeof(pdcache, device, :cache) == sizeof(Float32) * prod(dims)
32-
@test length(named_cache.free[key]) == 1
33-
@test length(named_cache.busy[key]) == 0
34-
@test x2 === x1
35-
@test x2 === named_cache.free[key][1]
36-
@test x_free !== x2
37-
38-
# Third allocation of different type - cache grows.
39-
40-
T2 = Int32
41-
key2 = hash((T2, dims))
42-
GPUArrays.AllocCache.@enable AT :cache begin
43-
x3 = AT(zeros(T2, dims))
44-
end
45-
@test sizeof(pdcache, device, :cache) == (sizeof(Float32) + sizeof(Int32)) * prod(dims)
46-
@test length(named_cache.free[key]) == 1
47-
@test length(named_cache.free[key2]) == 1
48-
@test x3 === named_cache.free[key2][1]
49-
50-
# Freeing all memory held by cache.
51-
52-
GPUArrays.AllocCache.invalidate!(AT, :cache)
53-
@test sizeof(pdcache, device, :cache) == 0
5429
end
30+
@test sizeof(pdcache, device, :cache) == sizeof(Float32) * prod(dims)
31+
@test length(named_cache.free[key]) == 1
32+
@test length(named_cache.busy[key]) == 0
33+
@test x2 === x1
34+
@test x2 === named_cache.free[key][1]
35+
@test x_free !== x2
36+
37+
# Third allocation of different type - cache grows.
38+
39+
T2 = Int32
40+
key2 = hash((T2, dims))
41+
GPUArrays.AllocCache.@enable AT :cache begin
42+
x3 = AT(zeros(T2, dims))
43+
end
44+
@test sizeof(pdcache, device, :cache) == (sizeof(Float32) + sizeof(Int32)) * prod(dims)
5545

56-
@testset "free_immediately=true" begin
57-
pdcache = GPUArrays.AllocCache.cache_allocator(AT)
58-
pdcache.free_immediately = true
59-
named_cache = GPUArrays.AllocCache.named_cache_allocator!(pdcache, device, :cache2)
60-
61-
T = Float32
62-
dims = (1, 2, 3)
63-
key = hash((T, dims))
64-
65-
@test sizeof(pdcache, device, :cache2) == 0
66-
67-
GPUArrays.AllocCache.@enable AT :cache2 begin
68-
x1 = AT(zeros(T, dims))
46+
_keys = collect(keys(named_cache.free))
47+
key2 = _keys[findfirst(i -> i != key, _keys)]
48+
@test length(named_cache.free[key]) == 1
49+
@test length(named_cache.free[key2]) == 1
50+
@test x3 === named_cache.free[key2][1]
6951

70-
@test !haskey(named_cache.free, key)
71-
@test length(named_cache.busy[key]) == 1
72-
@test sizeof(pdcache, device, :cache2) == sizeof(Float32) * prod(dims)
73-
end
52+
# Freeing all memory held by cache.
7453

75-
# `free` was never even used with `free_immediately=true`.
76-
@test !haskey(named_cache.free, key)
77-
@test length(named_cache.busy[key]) == 0
78-
@test sizeof(pdcache, device, :cache2) == 0
79-
end
54+
GPUArrays.AllocCache.invalidate!(AT, :cache)
55+
@test sizeof(pdcache, device, :cache) == 0
8056
end

0 commit comments

Comments
 (0)