Re-pin variable sized memory (#2599)

Author: jipolanco

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))

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