Compare CUDA devices by UUIDs

Author: jpsamaroo

Project.toml

@@ -10,6 +10,7 @@ Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 MemPool = "f9f48841-c794-520a-933b-121f7ba6ed94"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
+UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
 [compat]
 Adapt = "1, 2, 3"

src/cu.jl

@@ -1,13 +1,15 @@
 using .CUDA
 import .CUDA: CuDevice, CuContext, devices, attribute
 
+using UUIDs
+
 export CuArrayDeviceProc
 
 "Represents a single CUDA GPU device."
 struct CuArrayDeviceProc <: Dagger.Processor
     owner::Int
-    #ctx::CuContext
     device::Int
+    device_uuid::UUID
 end
 @gpuproc(CuArrayDeviceProc, CuArray)
 Dagger.get_parent(proc::CuArrayDeviceProc) = Dagger.OSProc(proc.owner)
@@ -49,7 +51,15 @@ function Dagger.move(from::CuArrayDeviceProc, to::CuArrayDeviceProc, x::Dagger.C
         finalizer(arr) do arr
             CUDA.cuIpcCloseMemHandle(pointer(arr))
         end
-        # FIXME: Deal with to_proc being a different GPU
+        if from.device_uuid != to.device_uuid
+            CUDA.device!(to.device) do
+                to_arr = similar(arr)
+                copyto!(to_arr, arr)
+                to_arr
+            end
+        else
+            arr
+        end
     else
         # Different node, use DtoH, serialization, HtoD
         # TODO UCX
@@ -68,7 +78,7 @@ function Dagger.execute!(proc::CuArrayDeviceProc, func, args...)
     end)
 end
 Base.show(io::IO, proc::CuArrayDeviceProc) =
-    print(io, "CuArrayDeviceProc on worker $(proc.owner), device $(proc.device)")
+    print(io, "CuArrayDeviceProc on worker $(proc.owner), device $(proc.device), uuid $(proc.device_uuid)")
 
 processor(::Val{:CUDA}) = CuArrayDeviceProc
 cancompute(::Val{:CUDA}) = CUDA.has_cuda()
@@ -77,7 +87,7 @@ kernel_backend(::CuArrayDeviceProc) = CUDADevice()
 if CUDA.has_cuda()
     for dev in devices()
         Dagger.add_callback!(() -> begin
-            return CuArrayDeviceProc(Distributed.myid(), #=CuContext(dev),=# dev.handle)
+            return CuArrayDeviceProc(Distributed.myid(), dev.handle, CUDA.uuid(dev))
         end)
     end
 end