datadeps: Fix incorrect aliasing behavior and add remainder copies

.buildkite/pipeline.yml

@@ -91,7 +91,7 @@ steps:
           codecov: true
 
   - label: Julia 1.11 (CUDA)
-    timeout_in_minutes: 20
+    timeout_in_minutes: 30
     <<: *gputest
     plugins:
       - JuliaCI/julia#v1:

ext/CUDAExt.jl

@@ -48,6 +48,13 @@ 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.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)])
@@ -75,6 +82,8 @@ function with_context!(space::CUDAVRAMMemorySpace)
     @assert Dagger.root_worker_id(space) == myid()
     with_context!(space.device)
 end
+Dagger.with_context!(proc::CuArrayDeviceProc) = with_context!(proc)
+Dagger.with_context!(space::CUDAVRAMMemorySpace) = with_context!(space)
 function with_context(f, x)
     old_ctx = context()
     old_stream = stream()

ext/IntelExt.jl

@@ -46,6 +46,13 @@ function Dagger.memory_space(x::oneArray)
     return IntelVRAMMemorySpace(myid(), device_id)
 end
 _device_id(dev::ZeDevice) = findfirst(other_dev->other_dev === dev, collect(oneAPI.devices()))
+function Dagger.aliasing(x::oneArray{T}) where T
+    space = Dagger.memory_space(x)
+    S = typeof(space)
+    gpu_ptr = pointer(x)
+    rptr = Dagger.RemotePtr{Cvoid}(UInt64(gpu_ptr), space)
+    return Dagger.ContiguousAliasing(Dagger.MemorySpan{S}(rptr, sizeof(T)*length(x)))
+end
 
 Dagger.memory_spaces(proc::oneArrayDeviceProc) = Set([IntelVRAMMemorySpace(proc.owner, proc.device_id)])
 Dagger.processors(space::IntelVRAMMemorySpace) = Set([oneArrayDeviceProc(space.owner, space.device_id)])
@@ -68,6 +75,8 @@ function with_context!(space::IntelVRAMMemorySpace)
     @assert Dagger.root_worker_id(space) == myid()
     with_context!(space.device_id)
 end
+Dagger.with_context!(proc::oneArrayDeviceProc) = with_context!(proc)
+Dagger.with_context!(space::IntelVRAMMemorySpace) = with_context!(space)
 function with_context(f, x)
     old_drv = driver()
     old_dev = device()

ext/MetalExt.jl

@@ -43,6 +43,13 @@ function Dagger.memory_space(x::MtlArray)
     return MetalVRAMMemorySpace(myid(), device_id)
 end
 _device_id(dev::MtlDevice) = findfirst(other_dev->other_dev === dev, Metal.devices())
+function Dagger.aliasing(x::MtlArray{T}) where T
+    space = Dagger.memory_space(x)
+    S = typeof(space)
+    gpu_ptr = pointer(x)
+    rptr = Dagger.RemotePtr{Cvoid}(UInt64(gpu_ptr), space)
+    return Dagger.ContiguousAliasing(Dagger.MemorySpan{S}(rptr, sizeof(T)*length(x)))
+end
 
 Dagger.memory_spaces(proc::MtlArrayDeviceProc) = Set([MetalVRAMMemorySpace(proc.owner, proc.device_id)])
 Dagger.processors(space::MetalVRAMMemorySpace) = Set([MtlArrayDeviceProc(space.owner, space.device_id)])
@@ -66,6 +73,8 @@ end
 function with_context!(space::MetalVRAMMemorySpace)
     @assert Dagger.root_worker_id(space) == myid()
 end
+Dagger.with_context!(proc::MtlArrayDeviceProc) = with_context!(proc)
+Dagger.with_context!(space::MetalVRAMMemorySpace) = with_context!(space)
 function with_context(f, x)
     with_context!(x)
     return f()

ext/OpenCLExt.jl

@@ -44,6 +44,13 @@ function Dagger.memory_space(x::CLArray)
     idx = findfirst(==(queue), QUEUES)
     return CLMemorySpace(myid(), idx)
 end
+function Dagger.aliasing(x::CLArray{T}) where T
+    space = Dagger.memory_space(x)
+    S = typeof(space)
+    gpu_ptr = pointer(x)
+    rptr = Dagger.RemotePtr{Cvoid}(UInt64(gpu_ptr), space)
+    return Dagger.ContiguousAliasing(Dagger.MemorySpan{S}(rptr, sizeof(T)*length(x)))
+end
 
 Dagger.memory_spaces(proc::CLArrayDeviceProc) = Set([CLMemorySpace(proc.owner, proc.device)])
 Dagger.processors(space::CLMemorySpace) = Set([CLArrayDeviceProc(space.owner, space.device)])
@@ -71,6 +78,8 @@ function with_context!(space::CLMemorySpace)
     @assert Dagger.root_worker_id(space) == myid()
     with_context!(space.device)
 end
+Dagger.with_context!(proc::CLArrayDeviceProc) = with_context!(proc)
+Dagger.with_context!(space::CLMemorySpace) = with_context!(space)
 function with_context(f, x)
     old_ctx = cl.context()
     old_queue = cl.queue()

ext/ROCExt.jl

@@ -39,6 +39,13 @@ end
 Dagger.root_worker_id(space::ROCVRAMMemorySpace) = space.owner
 Dagger.memory_space(x::ROCArray) =
     ROCVRAMMemorySpace(myid(), AMDGPU.device(x).device_id)
+function Dagger.aliasing(x::ROCArray{T}) where T
+    space = Dagger.memory_space(x)
+    S = typeof(space)
+    gpu_ptr = pointer(x)
+    rptr = Dagger.RemotePtr{Cvoid}(UInt64(gpu_ptr), space)
+    return Dagger.ContiguousAliasing(Dagger.MemorySpan{S}(rptr, sizeof(T)*length(x)))
+end
 
 Dagger.memory_spaces(proc::ROCArrayDeviceProc) = Set([ROCVRAMMemorySpace(proc.owner, proc.device_id)])
 Dagger.processors(space::ROCVRAMMemorySpace) = Set([ROCArrayDeviceProc(space.owner, space.device_id)])
@@ -67,6 +74,8 @@ function with_context!(space::ROCVRAMMemorySpace)
     @assert Dagger.root_worker_id(space) == myid()
     with_context!(space.device_id)
 end
+Dagger.with_context!(proc::ROCArrayDeviceProc) = with_context!(proc)
+Dagger.with_context!(space::ROCVRAMMemorySpace) = with_context!(space)
 function with_context(f, x)
     old_ctx = context()
     old_device = AMDGPU.device()

src/Dagger.jl

@@ -73,6 +73,9 @@ include("utils/fetch.jl")
 include("utils/chunks.jl")
 include("utils/logging.jl")
 include("submission.jl")
+abstract type MemorySpace end
+include("utils/memory-span.jl")
+include("utils/interval_tree.jl")
 include("memory-spaces.jl")
 
 # Task scheduling
@@ -83,7 +86,12 @@ include("utils/caching.jl")
 include("sch/Sch.jl"); using .Sch
 
 # Data dependency task queue
-include("datadeps.jl")
+include("datadeps/aliasing.jl")
+include("datadeps/chunkview.jl")
+include("datadeps/remainders.jl")
+include("datadeps/queue.jl")
+
+# Stencils
 include("utils/haloarray.jl")
 include("stencil.jl")
 

src/argument.jl

@@ -20,6 +20,7 @@ function pos_kw(pos::ArgPosition)
     @assert pos.kw != :NULL
     return pos.kw
 end
+
 mutable struct Argument
     pos::ArgPosition
     value
@@ -41,6 +42,35 @@ function Base.iterate(arg::Argument, state::Bool)
         return nothing
     end
 end
-
 Base.copy(arg::Argument) = Argument(ArgPosition(arg.pos), arg.value)
 chunktype(arg::Argument) = chunktype(value(arg))
+
+mutable struct TypedArgument{T}
+    pos::ArgPosition
+    value::T
+end
+TypedArgument(pos::Integer, value::T) where T = TypedArgument{T}(ArgPosition(true, pos, :NULL), value)
+TypedArgument(kw::Symbol, value::T) where T = TypedArgument{T}(ArgPosition(false, 0, kw), value)
+Base.setproperty!(arg::TypedArgument, name::Symbol, value::T) where T =
+    throw(ArgumentError("Cannot set properties of TypedArgument"))
+ispositional(arg::TypedArgument) = ispositional(arg.pos)
+iskw(arg::TypedArgument) = iskw(arg.pos)
+pos_idx(arg::TypedArgument) = pos_idx(arg.pos)
+pos_kw(arg::TypedArgument) = pos_kw(arg.pos)
+raw_position(arg::TypedArgument) = raw_position(arg.pos)
+value(arg::TypedArgument) = arg.value
+valuetype(arg::TypedArgument{T}) where T = T
+Base.iterate(arg::TypedArgument) = (arg.pos, true)
+function Base.iterate(arg::TypedArgument, state::Bool)
+    if state
+        return (arg.value, false)
+    else
+        return nothing
+    end
+end
+Base.copy(arg::TypedArgument{T}) where T = TypedArgument{T}(ArgPosition(arg.pos), arg.value)
+chunktype(arg::TypedArgument) = chunktype(value(arg))
+
+Argument(arg::TypedArgument) = Argument(arg.pos, arg.value)
+
+const AnyArgument = Union{Argument, TypedArgument}