Add tests

lib/JLArrays/src/JLArrays.jl

@@ -88,12 +88,20 @@ mutable struct JLArray{T, N} <: AbstractGPUArray{T, N}
     function JLArray{T,N}(::UndefInitializer, dims::Dims{N}) where {T,N}
         check_eltype(T)
         maxsize = prod(dims) * sizeof(T)
-        data = Vector{UInt8}(undef, maxsize)
-        ref = DataRef(data) do data
-            resize!(data, 0)
+
+        function _alloc_f()
+            data = Vector{UInt8}(undef, maxsize)
+            ref = DataRef(data) do data
+                resize!(data, 0)
+            end
+            obj = new{T,N}(ref, 0, dims)
+            finalizer(unsafe_free!, obj)
         end
-        obj = new{T,N}(ref, 0, dims)
-        finalizer(unsafe_free!, obj)
+
+        name = GPUArrays.CacheAllocatorName[]
+        return name == :none ?
+            _alloc_f() :
+            GPUArrays.alloc!(_alloc_f, JLBackend(), name, T, dims)::JLArray{T, N}
     end
 
     # low-level constructor for wrapping existing data
@@ -387,4 +395,12 @@ Adapt.adapt_storage(::JLBackend, a::Array) = Adapt.adapt(JLArrays.JLArray, a)
 Adapt.adapt_storage(::JLBackend, a::JLArrays.JLArray) = a
 Adapt.adapt_storage(::KernelAbstractions.CPU, a::JLArrays.JLArray) = convert(Array, a)
 
+# Caching Allocator.
+
+const JLACacheAllocator = GPUArrays.PerDeviceCacheAllocator(JLArray; free_immediately=false)
+
+GPUArrays.cache_allocator(::JLBackend) = JLACacheAllocator
+
+GPUArrays.device(::JLBackend) = 1
+
 end

src/host/allocations_cache.jl

@@ -69,19 +69,19 @@ function free_busy!(cache::CacheAllocator; free_immediately::Bool)
         busy_pool = get_pool!(cache, :busy, uid)
         isempty(busy_pool) && continue
 
-        free_pool = get_pool!(cache, :free, uid)
         Base.@lock cache.lock begin
             if free_immediately
                 map(unsafe_free!, busy_pool)
             else
+                free_pool = get_pool!(cache, :free, uid)
                 append!(free_pool, busy_pool)
             end
             empty!(busy_pool)
         end
     end
 end
 
-struct PerDeviceCacheAllocator{T <: AbstractGPUArray}
+mutable struct PerDeviceCacheAllocator{T <: AbstractGPUArray}
     lock::ReentrantLock
     caches::Dict{UInt64, Dict{Symbol, CacheAllocator{T}}}
     free_immediately::Bool
@@ -104,7 +104,7 @@ function named_cache_allocator!(pdcache::PerDeviceCacheAllocator{T}, device, nam
     named_cache = get(dev_cache, name, nothing)
     if named_cache ≡ nothing
         named_cache = CacheAllocator(T)
-        Base.@lock dev_cache.lock dev_cache[name] = named_cache
+        Base.@lock pdcache.lock dev_cache[name] = named_cache
     end
     return named_cache
 end

test/runtests.jl

@@ -48,6 +48,9 @@ include("setup.jl")     # make sure everything is precompiled
 const tests = []
 const test_runners = Dict()
 for AT in (JLArray, Array), name in keys(TestSuite.tests)
+    # Disable for now.
+    name == "Caching Allocator" && continue
+
     push!(tests, "$(AT)/$name")
     test_runners["$(AT)/$name"] = ()->TestSuite.tests[name](AT)
 end

test/testsuite.jl

@@ -93,6 +93,7 @@ include("testsuite/math.jl")
 include("testsuite/random.jl")
 include("testsuite/uniformscaling.jl")
 include("testsuite/statistics.jl")
+include("testsuite/caching_allocator.jl")
 
 """
 Runs the entire GPUArrays test suite on array type `AT`

test/testsuite/caching_allocator.jl

@@ -0,0 +1,82 @@
+@testsuite "Caching Allocator" (AT, eltypes) -> begin
+    # Hacky way to get KA backend from AT.
+    kab = KernelAbstractions.get_backend(AT(Array{Int}(undef, 0)))
+    device = GPUArrays.device(kab)
+
+    @testset "free_immediately=false" begin
+        pdcache = GPUArrays.cache_allocator(kab)
+        pdcache.free_immediately = false
+        named_cache = GPUArrays.named_cache_allocator!(pdcache, device, :cache)
+
+        T = Float32
+        dims = (1, 2, 3)
+        key = hash((T, dims))
+
+        GPUArrays.@cache_scope kab :cache begin
+            x1 = AT(zeros(T, dims))
+        end
+        @test sizeof(pdcache, device, :cache) == sizeof(Float32) * prod(dims)
+        @test length(named_cache.free[key]) == 1
+        @test length(named_cache.busy[key]) == 0
+        @test x1 === named_cache.free[key][1]
+
+        # Second allocation does not allocate - cache stays the same in size.
+
+        GPUArrays.@cache_scope kab :cache begin
+            x2 = AT(zeros(T, dims))
+
+            # Does not go to cache.
+            GPUArrays.@no_cache_scope begin
+                x_free = AT(zeros(T, dims))
+            end
+        end
+        @test sizeof(pdcache, device, :cache) == sizeof(Float32) * prod(dims)
+        @test length(named_cache.free[key]) == 1
+        @test length(named_cache.busy[key]) == 0
+        @test x2 === x1
+        @test x2 === named_cache.free[key][1]
+        @test x_free !== x2
+
+        # Third allocation of different type - cache grows.
+
+        T2 = Int32
+        key2 = hash((T2, dims))
+        GPUArrays.@cache_scope kab :cache begin
+            x3 = AT(zeros(T2, dims))
+        end
+        @test sizeof(pdcache, device, :cache) == (sizeof(Float32) + sizeof(Int32)) * prod(dims)
+        @test length(named_cache.free[key]) == 1
+        @test length(named_cache.free[key2]) == 1
+        @test x3 === named_cache.free[key2][1]
+
+        # Freeing all memory held by cache.
+
+        GPUArrays.invalidate_cache_allocator!(kab, :cache)
+        @test sizeof(pdcache, device, :cache) == 0
+    end
+
+    @testset "free_immediately=true" begin
+        pdcache = GPUArrays.cache_allocator(kab)
+        pdcache.free_immediately = true
+        named_cache = GPUArrays.named_cache_allocator!(pdcache, device, :cache2)
+
+        T = Float32
+        dims = (1, 2, 3)
+        key = hash((T, dims))
+
+        @test sizeof(pdcache, device, :cache2) == 0
+
+        GPUArrays.@cache_scope kab :cache2 begin
+            x1 = AT(zeros(T, dims))
+
+            @test !haskey(named_cache.free, key)
+            @test length(named_cache.busy[key]) == 1
+            @test sizeof(pdcache, device, :cache2) == sizeof(Float32) * prod(dims)
+        end
+
+        # `free` was never even used with `free_immediately=true`.
+        @test !haskey(named_cache.free, key)
+        @test length(named_cache.busy[key]) == 0
+        @test sizeof(pdcache, device, :cache2) == 0
+    end
+end