TEMP datadeps: Support CUDA execution

Author: jpsamaroo

ext/CUDAExt.jl

@@ -48,6 +48,22 @@ function Dagger.memory_space(x::CuArray)
     device_uuid = CUDA.uuid(dev)
     return CUDAVRAMMemorySpace(myid(), device_id, device_uuid)
 end
+function Dagger.aliasing(x::CuArray{T}) where T
+    space = Dagger.memory_space(x)
+    S = typeof(space)
+    cuptr = pointer(x)
+    rptr = Dagger.RemotePtr{Cvoid}(UInt64(cuptr), space)
+    return Dagger.ContiguousAliasing(Dagger.MemorySpan{S}(rptr, sizeof(T)*length(x)))
+end
+Dagger.pointer_in_space(ptr::UInt64, space::CUDAVRAMMemorySpace) =
+    CUDA.CuPtr{UInt8}(ptr)
+# TODO: Make async=true
+Dagger.unsafe_copyto_spaces!(to_space::CUDAVRAMMemorySpace, from_space::CUDAVRAMMemorySpace, to_ptr, from_ptr, len::UInt64) =
+    unsafe_copyto!(to_ptr, from_ptr, len; async=false)
+Dagger.unsafe_copyto_spaces!(to_space::CUDAVRAMMemorySpace, from_space::CPURAMMemorySpace, to_ptr, from_ptr, len::UInt64) =
+    unsafe_copyto!(to_ptr, from_ptr, len; async=false)
+Dagger.unsafe_copyto_spaces!(to_space::CPURAMMemorySpace, from_space::CUDAVRAMMemorySpace, to_ptr, from_ptr, len::UInt64) =
+    unsafe_copyto!(to_ptr, from_ptr, len; async=false)
 
 Dagger.memory_spaces(proc::CuArrayDeviceProc) = Set([CUDAVRAMMemorySpace(proc.owner, proc.device, proc.device_uuid)])
 Dagger.processors(space::CUDAVRAMMemorySpace) = Set([CuArrayDeviceProc(space.owner, space.device, space.device_uuid)])

src/datadeps/remainders.jl

@@ -143,6 +143,7 @@ function compute_remainder_for_arg!(state::DataDepsState,
         push!(target_ainfos, LocalMemorySpan.(spans))
     end
     nspans = length(first(target_ainfos))
+    @assert all(==(nspans), length.(target_ainfos)) "Aliasing info for $(typeof(arg_w.arg))[$(arg_w.dep_mod)] has different number of spans in different memory spaces"
 
     # FIXME: This is a hack to ensure that we don't miss any history generated by aliasing(...)
     for entry in state.arg_history[arg_w]
@@ -433,16 +434,18 @@ end
 
 # Main copy function for RemainderAliasing
 function move!(dep_mod::RemainderAliasing{S}, to_space::MemorySpace, from_space::MemorySpace, to::Chunk, from::Chunk) where S
+    # TODO: Support direct copy between GPU memory spaces
+
     # Copy the data from the source object
     copies = remotecall_fetch(root_worker_id(from_space), dep_mod) do dep_mod
         len = sum(span_tuple->span_len(span_tuple[1]), dep_mod.spans)
         copies = Vector{UInt8}(undef, len)
         offset = 1
         GC.@preserve copies begin
             for (from_span, _) in dep_mod.spans
-                from_ptr = Ptr{UInt8}(from_span.ptr)
+                from_ptr = pointer_in_space(from_span.ptr, from_space)
                 to_ptr = Ptr{UInt8}(pointer(copies, offset))
-                unsafe_copyto!(to_ptr, from_ptr, from_span.len)
+                unsafe_copyto_spaces!(CPURAMMemorySpace(), from_space, to_ptr, from_ptr, from_span.len)
                 offset += from_span.len
             end
         end
@@ -455,8 +458,8 @@ function move!(dep_mod::RemainderAliasing{S}, to_space::MemorySpace, from_space:
     GC.@preserve copies begin
         for (_, to_span) in dep_mod.spans
             from_ptr = Ptr{UInt8}(pointer(copies, offset))
-            to_ptr = Ptr{UInt8}(to_span.ptr)
-            unsafe_copyto!(to_ptr, from_ptr, to_span.len)
+            to_ptr = pointer_in_space(to_span.ptr, to_space)
+            unsafe_copyto_spaces!(to_space, CPURAMMemorySpace(), to_ptr, from_ptr, to_span.len)
             offset += to_span.len
         end
         @assert offset == length(copies)+1
@@ -467,3 +470,7 @@ function move!(dep_mod::RemainderAliasing{S}, to_space::MemorySpace, from_space:
 
     return
 end
+
+pointer_in_space(ptr::UInt64, space::CPURAMMemorySpace) = Ptr{UInt8}(ptr)
+unsafe_copyto_spaces!(to_space, from_space, to_ptr, from_ptr, len::UInt64) =
+    unsafe_copyto!(to_ptr, from_ptr, len)

src/memory-spaces.jl

@@ -1,6 +1,7 @@
 struct CPURAMMemorySpace <: MemorySpace
     owner::Int
 end
+CPURAMMemorySpace() = CPURAMMemorySpace(myid())
 root_worker_id(space::CPURAMMemorySpace) = space.owner
 
 memory_space(x) = CPURAMMemorySpace(myid())
@@ -87,7 +88,8 @@ function type_may_alias(::Type{T}) where T
     return false
 end
 
-may_alias(::MemorySpace, ::MemorySpace) = true
+may_alias(::MemorySpace, ::MemorySpace) = false
+may_alias(space1::M, space2::M) where M<:MemorySpace = space1 == space2
 may_alias(space1::CPURAMMemorySpace, space2::CPURAMMemorySpace) = space1.owner == space2.owner
 
 abstract type AbstractAliasing end
@@ -571,7 +573,7 @@ end
 function will_alias(x_span::MemorySpan, y_span::MemorySpan)
     may_alias(x_span.ptr.space, y_span.ptr.space) || return false
     # FIXME: Allow pointer conversion instead of just failing
-    @assert x_span.ptr.space == y_span.ptr.space
+    @assert x_span.ptr.space == y_span.ptr.space "Memory spans are in different spaces: $(x_span.ptr.space) vs. $(y_span.ptr.space)"
     x_end = x_span.ptr + x_span.len - 1
     y_end = y_span.ptr + y_span.len - 1
     return x_span.ptr <= y_end && y_span.ptr <= x_end