Teach DataRef about cached references to support unsafe_free.

Author: maleadt

src/host/abstractarray.jl

@@ -57,11 +57,12 @@ end
 mutable struct DataRef{D}
     rc::RefCounted{D}
     freed::Bool
+    cached::Bool
 end
 
 function DataRef(finalizer, ref::D) where {D}
     rc = RefCounted{D}(ref, finalizer, Threads.Atomic{Int}(1))
-    DataRef{D}(rc, false)
+    DataRef{D}(rc, false, false)
 end
 DataRef(ref; kwargs...) = DataRef(nothing, ref; kwargs...)
 
@@ -79,10 +80,16 @@ function Base.copy(ref::DataRef{D}) where {D}
         throw(ArgumentError("Attempt to copy a freed reference."))
     end
     retain(ref.rc)
-    return DataRef{D}(ref.rc, false)
+    # copies of cached references are not managed by the cache, so
+    # we need to mark them as such to make sure their refcount can drop.
+    return DataRef{D}(ref.rc, false, false)
 end
 
 function unsafe_free!(ref::DataRef)
+    if ref.cached
+        # lifetimes of cached references are tied to the cache.
+        return
+    end
     if ref.freed
         # multiple frees *of the same object* are allowed.
         # we should only ever call `release` once per object, though,

src/host/alloc_cache.jl

@@ -44,15 +44,12 @@ function cached_alloc(f, key)
 
         if !isempty(free_pool)
             ref = Base.@lock cache.lock pop!(free_pool)
-            @assert !ref.freed
         end
     end
 
     if ref === nothing
         ref = f()::DataRef
-
-        # increase the refcount of the ref to prevent finalizers from freeing it
-        retain(ref.rc)
+        ref.cached = true
     end
 
     Base.@lock cache.lock begin
@@ -80,16 +77,11 @@ end
 function unsafe_free!(cache::AllocCache)
     Base.@lock cache.lock begin
         for pool in values(cache.busy)
-            isempty(pool) || error(
-                "Invalidating allocations cache that's currently in use. " *
-                    "Invalidating inside `@cached` is not allowed."
-            )
+            isempty(pool) || error("Cannot invalidate a cache that's in active use")
         end
         for pool in values(cache.free), ref in pool
-            # release our hold on the underlying data
-            release(ref.rc)
-
-            # early-release the reference
+            # release the reference
+            ref.cached = false
             unsafe_free!(ref)
         end
         empty!(cache.free)

test/testsuite/alloc_cache.jl

@@ -5,75 +5,87 @@
         # first allocation populates the cache
         T, dims = Float32, (1, 2, 3)
         GPUArrays.@cached cache begin
-            x1 = AT(zeros(T, dims))
+            cached1 = AT(zeros(T, dims))
         end
-        @test sizeof(cache) == sizeof(T) * prod(dims)
+        @test sizeof(cache) == sizeof(cached1)
         key = first(keys(cache.free))
         @test length(cache.free[key]) == 1
         @test length(cache.busy[key]) == 0
-        @test cache.free[key][1] === GPUArrays.storage(x1)
+        @test cache.free[key][1] === GPUArrays.storage(cached1)
 
         # second allocation hits the cache
         GPUArrays.@cached cache begin
-            x2 = AT(zeros(T, dims))
+            cached2 = AT(zeros(T, dims))
 
             # explicitly uncached ones don't
-            GPUArrays.@uncached x_free = AT(zeros(T, dims))
+            GPUArrays.@uncached uncached = AT(zeros(T, dims))
         end
-        @test sizeof(cache) == sizeof(T) * prod(dims)
+        @test sizeof(cache) == sizeof(cached2)
         key = first(keys(cache.free))
         @test length(cache.free[key]) == 1
         @test length(cache.busy[key]) == 0
-        @test cache.free[key][1] === GPUArrays.storage(x2)
-        @test x_free !== x2
+        @test cache.free[key][1] === GPUArrays.storage(cached2)
+        @test uncached !== cached2
 
         # compatible shapes should also hit the cache
         dims = (3, 2, 1)
         GPUArrays.@cached cache begin
-            x3 = AT(zeros(T, dims))
+            cached3 = AT(zeros(T, dims))
         end
-        @test sizeof(cache) == sizeof(T) * prod(dims)
+        @test sizeof(cache) == sizeof(cached3)
         key = first(keys(cache.free))
         @test length(cache.free[key]) == 1
         @test length(cache.busy[key]) == 0
-        @test cache.free[key][1] === GPUArrays.storage(x3)
+        @test cache.free[key][1] === GPUArrays.storage(cached3)
 
         # as should compatible eltypes
         T = Int32
         GPUArrays.@cached cache begin
-            x4 = AT(zeros(T, dims))
+            cached4 = AT(zeros(T, dims))
         end
-        @test sizeof(cache) == sizeof(T) * prod(dims)
+        @test sizeof(cache) == sizeof(cached4)
         key = first(keys(cache.free))
         @test length(cache.free[key]) == 1
         @test length(cache.busy[key]) == 0
-        @test cache.free[key][1] === GPUArrays.storage(x4)
+        @test cache.free[key][1] === GPUArrays.storage(cached4)
 
         # different shapes should trigger a new allocation
         dims = (2, 2)
         GPUArrays.@cached cache begin
-            x5 = AT(zeros(T, dims))
+            cached5 = AT(zeros(T, dims))
 
-            # we're allowed to early free arrays, which shouldn't release the underlying data
-            GPUArrays.unsafe_free!(x5)
+            # we're allowed to early free arrays, which should be a no-op for cached data
+            GPUArrays.unsafe_free!(cached5)
         end
+        @test sizeof(cache) == sizeof(cached4) + sizeof(cached5)
         _keys = collect(keys(cache.free))
         key2 = _keys[findfirst(i -> i != key, _keys)]
         @test length(cache.free[key]) == 1
         @test length(cache.free[key2]) == 1
-        @test cache.free[key2][1] === GPUArrays.storage(x5)
+        @test cache.free[key2][1] === GPUArrays.storage(cached5)
+
+        # we should be able to re-use the early-freed
+        GPUArrays.@cached cache begin
+            cached5 = AT(zeros(T, dims))
+        end
 
         # freeing all memory held by cache should free all allocations
-        @test !GPUArrays.storage(x1).freed
-        @test GPUArrays.storage(x5).freed
-        @test GPUArrays.storage(x5).rc.count[] == 1 # the ref appears freed, but the data isn't
-        @test !GPUArrays.storage(x_free).freed
+        @test !GPUArrays.storage(cached1).freed
+        @test GPUArrays.storage(cached1).cached
+        @test !GPUArrays.storage(cached5).freed
+        @test GPUArrays.storage(cached5).cached
+        @test !GPUArrays.storage(uncached).freed
+        @test !GPUArrays.storage(uncached).cached
         GPUArrays.unsafe_free!(cache)
         @test sizeof(cache) == 0
-        @test GPUArrays.storage(x1).freed
-        @test GPUArrays.storage(x1).rc.count[] == 0
-        @test GPUArrays.storage(x5).freed
-        @test GPUArrays.storage(x5).rc.count[] == 0
-        @test !GPUArrays.storage(x_free).freed
+        @test GPUArrays.storage(cached1).freed
+        @test !GPUArrays.storage(cached1).cached
+        @test GPUArrays.storage(cached5).freed
+        @test !GPUArrays.storage(cached5).cached
+        @test !GPUArrays.storage(uncached).freed
+        ## test that the underlying data was freed as well
+        @test GPUArrays.storage(cached1).rc.count[] == 0
+        @test GPUArrays.storage(cached5).rc.count[] == 0
+        @test GPUArrays.storage(uncached).rc.count[] == 1
     end
 end