From b97643cbe1f8ae5667d59dfb851661fff8d2f158 Mon Sep 17 00:00:00 2001
From: Tim Besard <tim.besard@gmail.com>
Date: Wed, 30 Oct 2024 12:33:54 +0100
Subject: [PATCH] Add `storage` query to get an array's DataRef. (#567)

This makes it possible to implement `unsafe_free!` and other calls
directly in GPUArrays.
---
 Project.toml                 |  2 +-
 lib/JLArrays/Project.toml    |  2 +-
 lib/JLArrays/src/JLArrays.jl |  2 +-
 src/host/abstractarray.jl    | 15 ++++++++++++++-
 test/testsuite/base.jl       |  8 ++++++++
 5 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/Project.toml b/Project.toml
index a3f79030..551c9edc 100644
--- a/Project.toml
+++ b/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"
diff --git a/lib/JLArrays/Project.toml b/lib/JLArrays/Project.toml
index ce9e8a87..ffb70ad2 100644
--- a/lib/JLArrays/Project.toml
+++ b/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"
diff --git a/lib/JLArrays/src/JLArrays.jl b/lib/JLArrays/src/JLArrays.jl
index 69268b82..dbec3d25 100644
--- a/lib/JLArrays/src/JLArrays.jl
+++ b/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}
diff --git a/src/host/abstractarray.jl b/src/host/abstractarray.jl
index a4a86d67..150cd88b 100644
--- a/src/host/abstractarray.jl
+++ b/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
 
diff --git a/test/testsuite/base.jl b/test/testsuite/base.jl
index 0e134398..7c3cfe07 100644
--- a/test/testsuite/base.jl
+++ b/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