CUDA: Fix multi-GPU data movement

Author: jpsamaroo

ext/CUDAExt.jl

@@ -80,6 +80,24 @@ function Dagger.move(from::CuArrayDeviceProc, to::CuArrayDeviceProc, x::Dagger.C
     end
 end
 
+function Dagger.move(from_proc::CPUProc, to_proc::CuArrayDeviceProc, x::CuArray)
+    # TODO: No extra allocations here
+    if CUDA.device(x) == collect(CUDA.devices())[to_proc.device+1]
+        return x
+    end
+    DaggerGPU.with_device(to_proc) do
+        _x = similar(x)
+        copyto!(_x, x)
+        return _x
+    end
+end
+
+function Dagger.move(from_proc::CuArrayDeviceProc, to_proc::CPUProc, x::CuArray{T,N}) where {T,N}
+    _x = Array{T,N}(undef, size(x))
+    copyto!(_x, x)
+    return _x
+end
+
 function Dagger.execute!(proc::CuArrayDeviceProc, f, args...; kwargs...)
     @nospecialize f args kwargs
     tls = Dagger.get_tls()

test/runtests.jl

@@ -79,28 +79,37 @@ end
             CuArrayDeviceProc
         end
         @test DaggerGPU.processor(:CUDA) === cuproc
-        b = generate_thunks()
-        c = Dagger.with_options(;scope=Dagger.scope(cuda_gpu=1)) do
-            @test fetch(Dagger.@spawn isongpu(b))
-            Dagger.@spawn sum(b)
+        ndevices = length(collect(CUDA.devices()))
+
+        @testset "Arrays (GPU $gpu)" for gpu in 1:min(ndevices, 2)
+            b = generate_thunks()
+            c = Dagger.with_options(;scope=Dagger.scope(cuda_gpu=gpu)) do
+                @test fetch(Dagger.@spawn isongpu(b))
+                Dagger.@spawn sum(b)
+            end
+            @test !fetch(Dagger.@spawn isongpu(b))
+            @test fetch(Dagger.@spawn identity(c)) == 20
         end
-        @test !fetch(Dagger.@spawn isongpu(b))
-        @test fetch(Dagger.@spawn identity(c)) == 20
 
-        @testset "KernelAbstractions" begin
+        @testset "KernelAbstractions (GPU $gpu)" for gpu in 1:min(ndevices, 2)
             A = rand(Float32, 8)
-            DA, T = Dagger.with_options(;scope=Dagger.scope(cuda_gpu=1)) do
+            DA, T = Dagger.with_options(;scope=Dagger.scope(cuda_gpu=gpu)) do
                 fetch(Dagger.@spawn fill_thunk(A, 2.3f0))
             end
             @test all(DA .== 2.3f0)
             @test T <: CuArray
 
-            A = CUDA.rand(128)
-            B = CUDA.zeros(128)
-            Dagger.with_options(;scope=Dagger.scope(worker=1,cuda_gpu=1)) do
+            local A, B
+            CUDA.device!(gpu-1) do
+                A = CUDA.rand(128)
+                B = CUDA.zeros(128)
+            end
+            Dagger.with_options(;scope=Dagger.scope(worker=1,cuda_gpu=gpu)) do
                 fetch(Dagger.@spawn Kernel(copy_kernel)(B, A; ndrange=length(A)))
             end
-            @test all(B .== A)
+            CUDA.device!(gpu-1) do
+                @test all(B .== A)
+            end
         end
     end
 end