Initial support for in-place operations with unified memory (#21)

* Handle multiple GPUs
* Add support for in-place ops with unified memory

Author: ronisbr

src/metal.jl

@@ -3,10 +3,74 @@ import .Metal: MtlArray, MtlDevice
 
 struct MtlArrayDeviceProc <: Dagger.Processor
     owner::Int
-    device::MtlDevice
+    device_id::UInt64
+end
+
+# Assume that we can run anything.
+Dagger.iscompatible_func(proc::MtlArrayDeviceProc, opts, f) = true
+Dagger.iscompatible_arg(proc::MtlArrayDeviceProc, opts, x) = true
+
+# CPUs shouldn't process our array type.
+Dagger.iscompatible_arg(proc::Dagger.ThreadProc, opts, x::MtlArray) = false
+
+function Dagger.move(from_proc::OSProc, to_proc::MtlArrayDeviceProc, x::Chunk)
+    from_pid = from_proc.pid
+    to_pid = Dagger.get_parent(to_proc).pid
+    @assert myid() == to_pid
+
+    return Dagger.move(from_proc, to_proc, remotecall_fetch(x->poolget(x.handle), from_pid, x))
+end
+
+function Dagger.move(from_proc::MtlArrayDeviceProc, to_proc::OSProc, x::Chunk)
+    from_pid = Dagger.get_parent(from_proc).pid
+    to_pid = to_proc.pid
+    @assert myid() == to_pid
+
+    return remotecall_fetch(from_pid, x) do x
+        mtlarray = poolget(x.handle)
+        return Dagger.move(from_proc, to_proc, mtlarray)
+    end
+end
+
+function Dagger.move(
+    from_proc::OSProc,
+    to_proc::MtlArrayDeviceProc,
+    x::Array{T, N}
+) where {T, N}
+    # If we have unified memory, we can try casting the `Array` to `MtlArray`.
+    device = _get_metal_device(to_proc)
+
+    if (device !== nothing) && device.hasUnifiedMemory
+        marray = _cast_array_to_mtlarray(x, device)
+        marray !== nothing && return marray
+    end
+
+    return adapt(MtlArray, x)
+end
+
+function Dagger.move(from_proc::OSProc, to_proc::MtlArrayDeviceProc, x)
+    adapt(MtlArray, x)
+end
+
+function Dagger.move(
+    from_proc::MtlArrayDeviceProc,
+    to_proc::OSProc,
+    x::Array{T, N}
+) where {T, N}
+    # If we have unified memory, we can just cast the `MtlArray` to an `Array`.
+    device = _get_metal_device(from_proc)
+
+    if (device !== nothing) && device.hasUnifiedMemory
+        return unsafe_wrap(Array{T}, x, size(x))
+    else
+        return adapt(Array, x)
+    end
+end
+
+function Dagger.move(from_proc::MtlArrayDeviceProc, to_proc::OSProc, x)
+    adapt(Array, x)
 end
 
-@gpuproc(MtlArrayDeviceProc, MtlArray)
 Dagger.get_parent(proc::MtlArrayDeviceProc) = Dagger.OSProc(proc.owner)
 
 function Dagger.execute!(proc::MtlArrayDeviceProc, func, args...)
@@ -30,15 +94,58 @@ function Dagger.execute!(proc::MtlArrayDeviceProc, func, args...)
 end
 
 function Base.show(io::IO, proc::MtlArrayDeviceProc)
-    print(io, "MtlArrayDeviceProc on worker $(proc.owner), device ($(proc.device.name))")
+    print(io, "MtlArrayDeviceProc on worker $(proc.owner), device ($(something(_get_metal_device(proc)).name))")
 end
 
 processor(::Val{:Metal}) = MtlArrayDeviceProc
 cancompute(::Val{:Metal}) = length(Metal.devices()) >= 1
-kernel_backend(::MtlArrayDeviceProc) = Metal.current_device()
+kernel_backend(proc::MtlArrayDeviceProc) = _get_metal_device(proc)
+
+for dev in Metal.devices()
+    Dagger.add_processor_callback!("metal_device_$(dev.registryID)") do
+        MtlArrayDeviceProc(Distributed.myid(), dev.registryID)
+    end
+end
+
+################################################################################
+#                              Private functions
+################################################################################
+
+# Try casting the array `x` to an `MtlArray`. If the casting is not possible,
+# return `nothing`.
+function _cast_array_to_mtlarray(x::Array{T,N}, device::MtlDevice) where {T,N}
+    # Try creating the buffer without copying.
+    dims = size(x)
+    nbytes_array = prod(dims) * sizeof(T)
+    pagesize = ccall(:getpagesize, Cint, ())
+    num_pages = div(nbytes_array, pagesize, RoundUp)
+    nbytes = num_pages * pagesize
+
+    pbuf = Metal.MTL.mtDeviceNewBufferWithBytesNoCopy(
+        device,
+        pointer(x),
+        nbytes,
+        Metal.Shared | Metal.MTL.DefaultTracking | Metal.MTL.DefaultCPUCache
+    )
+
+    if pbuf != C_NULL
+        buf = MtlBuffer(pbuf)
+        marray = MtlArray{T,N}(buf, dims)
+        return marray
+    end
+
+    # If we reached here, the conversion was not possible.
+    return nothing
+end
+
+# Return the Metal device handler given the ID recorded in `proc`.
+function _get_metal_device(proc::MtlArrayDeviceProc)
+    devices = Metal.devices()
+    id = findfirst(dev -> dev.registryID == proc.device_id, devices)
 
-if length(Metal.devices()) >= 1
-    Dagger.add_processor_callback!("metal_device") do
-        MtlArrayDeviceProc(Distributed.myid(), Metal.current_device())
+    if devices === nothing
+        return nothing
+    else
+        return devices[id]
     end
 end

test/runtests.jl

@@ -22,6 +22,11 @@ end
         @show A
         A
     end
+
+    # Create a function to perform an in-place operation.
+    function addarray!(x)
+        x .= x .+ 1.0f0
+    end
 end
 
 function generate_thunks()
@@ -108,5 +113,51 @@ end
 
         # It seems KernelAbstractions does not support Metal.jl.
         @test_skip "KernelAbstractions"
+
+        @testset "In-place operations" begin
+            # Create a page-aligned array.
+            dims = (2, 2)
+            T = Float32
+            pagesize = ccall(:getpagesize, Cint, ())
+            addr = Ref(C_NULL)
+
+            ccall(
+                :posix_memalign,
+                Cint,
+                (Ptr{Ptr{Cvoid}}, Csize_t, Csize_t), addr,
+                pagesize,
+                prod(dims) * sizeof(T)
+            )
+
+            array = unsafe_wrap(
+                Array{T, length(dims)},
+                reinterpret(Ptr{T}, addr[]),
+                dims,
+                own = false
+            )
+
+            # Initialize the array.
+            array[1, 1] = 1
+            array[1, 2] = 2
+            array[2, 1] = 3
+            array[2, 2] = 4
+
+            # Perform the computation only on a local `MtlArrayDeviceProc`
+            t = Dagger.@spawn single=myid() proclist = [metalproc] addarray!(array)
+
+            # Fetch and check the results.
+            ret = fetch(t)
+
+            @test ret[1, 1] == 2.0f0
+            @test ret[1, 2] == 3.0f0
+            @test ret[2, 1] == 4.0f0
+            @test ret[2, 2] == 5.0f0
+
+            # Check if the operation happened in-place.
+            @test array[1, 1] == 2.0f0
+            @test array[1, 2] == 3.0f0
+            @test array[2, 1] == 4.0f0
+            @test array[2, 2] == 5.0f0
+        end
     end
 end