From 830c49b4204ccb54f751a8e00ebe3126f7b47c30 Mon Sep 17 00:00:00 2001
From: Juan Ignacio Polanco <juan-ignacio.polanco@cnrs.fr>
Date: Thu, 19 Dec 2024 11:04:59 +0100
Subject: [PATCH] Re-pin variable sized memory (#2599)

---
 lib/cudadrv/memory.jl | 18 ++++++++++++++----
 test/core/cudadrv.jl  | 14 ++++++++++++++
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/lib/cudadrv/memory.jl b/lib/cudadrv/memory.jl
index 13c4aa6f81..8fd4950d6b 100644
--- a/lib/cudadrv/memory.jl
+++ b/lib/cudadrv/memory.jl
@@ -659,9 +659,14 @@ end
 
 const __pin_lock = ReentrantLock()
 
+struct PinnedObject
+    ref::WeakRef
+    size::Int  # memory size in bytes
+end
+
 # - IdDict does not free the memory
 # - WeakRef dict does not unique the key by objectid
-const __pinned_objects = Dict{Tuple{CuContext,Ptr{Cvoid}}, WeakRef}()
+const __pinned_objects = Dict{Tuple{CuContext,Ptr{Cvoid}}, PinnedObject}()
 
 function pin(a::AbstractArray)
     ctx = context()
@@ -670,10 +675,15 @@ function pin(a::AbstractArray)
     Base.@lock __pin_lock begin
         # only pin an object once per context
         key = (ctx, convert(Ptr{Nothing}, ptr))
-        if haskey(__pinned_objects, key) && __pinned_objects[key].value !== nothing
-            return nothing
+        if haskey(__pinned_objects, key) && __pinned_objects[key].ref.value !== nothing
+            if sizeof(a) == __pinned_objects[key].size
+                return nothing
+            else
+                # if the object size has changed, unpin it first; it will be re-pinned with the new size
+                __unpin(ptr, ctx)
+            end
         end
-        __pinned_objects[key] = WeakRef(a)
+        __pinned_objects[key] = PinnedObject(WeakRef(a), sizeof(a))
     end
 
      __pin(ptr, sizeof(a))
diff --git a/test/core/cudadrv.jl b/test/core/cudadrv.jl
index 355a7beb98..2a1a1c9870 100644
--- a/test/core/cudadrv.jl
+++ b/test/core/cudadrv.jl
@@ -689,6 +689,20 @@ if attribute(device(), CUDA.DEVICE_ATTRIBUTE_HOST_REGISTER_SUPPORTED) != 0
     dA = CUDA.rand(UInt8, 512)
     copyto!(dA, hA)
     copyto!(hA, dA)
+
+    # memory copies with resized pinned memory (used to fail with CUDA_ERROR_INVALID_VALUE)
+    dA = rand(Float32, 100)
+    hA = Array(dA)
+    @test !CUDA.is_pinned(pointer(hA))
+    for n ∈ 100:2000
+        resize!(dA, n)
+        resize!(hA, n)
+        dA .= n
+        CUDA.pin(hA)
+        @test CUDA.is_pinned(pointer(hA))
+        copyto!(hA, dA)
+        copyto!(dA, hA)
+    end
 end
 
 end