CUDA tasking using default alignment constraint for all Input(s) and Output(s)

[krasow]

PR #119 leverages our CUDA.jl tasking system.

However, launch(...) will call Legate.default_alignment(...) on all inputs and outputs.

This is fine for standard elementwise operations like:

a .+ b

However, it is inefficient for stencil computations where inputs are shifted views of the same array.

For example, consider a 6×6 grid split across two tasks, where task 0 owns the top half and task 1 owns the bottom half:

grid   = zeros(6, 6)
center = grid[2:5, 2:5]   # 4×4 interior
south  = grid[3:6, 2:5]   # shifted down by 1 row

With default_alignment, center and south get identical tile boundaries:

# task 0
center_t0 = grid[2:3, 2:5]
south_t0  = grid[3:4, 2:5]

# task 1
center_t1 = grid[4:5, 2:5]
south_t1  = grid[5:6, 2:5]

This is a problem for stencil computations because the shifted views require halo data from neighboring partitions. For a Jacobi stencil using all four neighbors, this creates halo copies for north, south, east, and west every iteration.

Instead, this should be represented with a bloat constraint:

bloat(source=center, bloat=south, low=0, high=1)

This gives each task the necessary overlap in its physical instance:

# task 0
center_t0 = grid[2:3, 2:5]
south_t0  = grid[3:5, 2:5]

# task 1
center_t1 = grid[4:5, 2:5]
south_t1  = grid[4:6, 2:5]

Now the required halo rows are included by construction. The overlap is handled once during partitioning rather than copied every iteration.

For example, with 1000 Jacobi iterations over a 1000×1000 grid and four neighbors:

• default_alignment: pays halo-copy cost 4000 times
• bloat: pays overlap cost once at partitioning time

So we likely need a way for CUDA.jl tasks to specify Legate partitioning constraints other than default_alignment.

See more details here about the various constraints in legate.

[krasow]

claude states a simple solution is to add provenance tracking for slice expressions. If an NDArray was formed by a slice, we have can track its parent and use that information to add a bloat constraint. The offsets are recorded to emit the bloat constraint.

This is the current impl of nda_get_slice

  function nda_get_slice(arr::NDArray{T,N}, slices::Vector{Slice}) where {T,N}
      ptr = ccall((:nda_get_slice, libnda), ...)
      return NDArray(ptr, T, Val(N))   # no parent, no offset recorded
  end

Imagine we have a broadcast expression like:

b[2:end-1, :] = a[1:end-2, :] .* a[3:end,:]

The parent field exists on the NDArray struct but is never set by slicing. So by the time fuse_broadcast_tree! sees a[1:end-2, :] and a[3:end, :], both are opaque NDArrays with no memory of where they came from. If we know a is the parent and their relative offsets, we can add the bloat constraint.

1. Record provenance in nda_get_slice:

 function nda_get_slice(arr::NDArray{T,N}, slices::Vector{Slice}) where {T,N}
     ptr = ccall((:nda_get_slice, libnda), ...)
     offsets = NTuple{N,Int}(s.has_start == 1 ? s.start : 0 for s in slices)
     return NDArray(ptr, T, Val(N), arr, offsets)   # track parent + offset
 end

2. For each pair of inputs that share the same parent, compute the relative offset and emit a bloat constraint:

  for (i, a) in enumerate(input_ndarrays), (j, b) in enumerate(input_ndarrays)
      i >= j && continue
      if same_parent(a, b)
          low, high = relative_halo(a, b)   # derived from their offsets
          Legate.add_constraint(task, Legate.bloat(input_vars[i], input_vars[j], low, high))
      end
  end

[ejmeitz]

Can also look at how CUDA.jl tracks parents. I believe they do something similar.