Add storage query to get an array's DataRef. (#567)

This makes it possible to implement `unsafe_free!` and other calls
directly in GPUArrays.

Author: maleadt

Project.toml

@@ -1,6 +1,6 @@
 name = "GPUArrays"
 uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
-version = "11.0.0"
+version = "11.1.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

lib/JLArrays/Project.toml

@@ -11,6 +11,6 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 
 [compat]
 Adapt = "2.0, 3.0, 4.0"
-GPUArrays = "11"
+GPUArrays = "11.1"
 Random = "1"
 julia = "1.8"

lib/JLArrays/src/JLArrays.jl

@@ -105,7 +105,7 @@ mutable struct JLArray{T, N} <: AbstractGPUArray{T, N}
     end
 end
 
-unsafe_free!(a::JLArray) = GPUArrays.unsafe_free!(a.data)
+GPUArrays.storage(a::JLArray) = a.data
 
 # conversion of untyped data to a typed Array
 function typed_data(x::JLArray{T}) where {T}

src/host/abstractarray.jl

@@ -3,7 +3,7 @@
 
 # storage handling
 
-export DataRef
+export DataRef, unsafe_free!
 
 # DataRef provides a helper class to manage the storage of an array.
 #
@@ -92,6 +92,19 @@ function unsafe_free!(ref::DataRef, args...)
     return
 end
 
+# array methods
+
+storage(x::AbstractGPUArray) = error("Not implemented") # COV_EXCL_LINE
+
+"""
+    unsafe_free!(a::GPUArray)
+
+Release the memory of an array for reuse by future allocations. This operation is
+performed automatically by the GC when an array goes out of scope, but can be called
+earlier to reduce pressure on the memory allocator.
+"""
+unsafe_free!(x::AbstractGPUArray) = unsafe_free!(storage(x))
+
 
 # input/output
 

test/testsuite/base.jl

@@ -21,6 +21,14 @@ end
 end
 
 @testsuite "base" (AT, eltypes)->begin
+    if AT <: AbstractGPUArray
+        @testset "storage" begin
+          x = AT(rand(Float32, 10))
+          @test GPUArrays.storage(x) isa GPUArrays.DataRef
+          GPUArrays.unsafe_free!(x)
+        end
+    end
+
     @testset "copy!" begin
         for (dst, src,) in (
                             (rand(Float32, (10,)),   rand(Float32, (10,))),   # vectors