Add view functionality and tests for DArray chunk slicing

Author: AkhilAkkapelli

docs/src/darray.md

@@ -211,6 +211,78 @@ across the workers in the Julia cluster in a relatively even distribution;
 future operations on a `DArray` may produce a different distribution from the
 one chosen by previous calls.
 
+### DArray Chunk Slicing with `view`
+
+Dagger's `view` function allows you to efficiently create a "view" of a `DArray`'s `Chunk` or a `DTask` that produces a `DArray` `Chunk`. This enables operations on specific parts of your distributed data using standard Julia array slicing, without needing to materialize the entire chunk.
+
+```view(c::Chunk, slices...) & view(c::DTask, slices...)```
+
+These methods create a `view` for a `DArray` `Chunk` object or for a `DTask` that will produce a `DArray` `Chunk`. You specify the desired sub-region using standard Julia array slicing syntax, identical to how you would slice a regular Array.
+
+#### Examples
+
+```julia
+julia> A = rand(64, 64)
+64×64 Matrix{Float64}:
+[...]
+
+julia> DA = DArray(A, Blocks(8,8)) 
+64x64 DMatrix{Float64} with 8x8 partitions of size 8x8:
+[...]
+
+julia> chunk = DA.chunks[1,1] 
+DTask (finished)
+
+julia> view(chunk, :, :) # View the entire 8x8 chunk
+ChunkSlice{2}(Dagger.Chunk(...), (Colon(), Colon()))
+
+julia> view(chunk, 1:4, 1:4) # View the top-left 4x4 sub-region of the chunk
+ChunkSlice{2}(Dagger.Chunk(...), (1:4, 1:4))
+
+julia> view(chunk, 1, :) # View the first row of the chunk
+ChunkSlice{2}(Dagger.Chunk(...), (1, Colon()))
+
+julia> view(chunk, :, 5) # View the fifth column of the chunk
+ChunkSlice{2}(Dagger.Chunk(...), (Colon(), 5))
+
+julia> view(chunk, 1:2:7, 2:2:8) # View with stepped ranges
+ChunkSlice{2}(Dagger.Chunk(...), (1:2:7, 2:2:8))
+```
+
+#### Example Usage: Parallel Row Summation of a DArray using `view`
+This example demonstrates how to sum multiple rows of a `DArray` by using `view` to process individual rows within chunks to get Row Sum Vector.
+
+```julia
+julia> A = DArray(rand(10,1000), Blocks(2,1000))
+10x1000 DMatrix{Float64} with 5x1 partitions of size 2x1000: 
+[...]
+
+# Helper function to sum a single row and store it in a provided array view
+julia> @everywhere function sum_array_row!(row_sum::AbstractArray{Float64}, x::AbstractArray{Float64})
+    row_sum[1] = sum(x)
+end
+
+# Number of rows
+julia> nrows = size(A,1)
+
+# Initialize a zero array in the final row sums
+julia> row_sums = zeros(nrows)
+
+# Spawn tasks to sum each row in parallel using views
+julia> Dagger.spawn_datadeps() do
+           sz = size(A.chunks,1) 
+           nrows_per_chunk = nrows ÷ sz
+           for i in 1:sz
+              for j in 1:nrows_per_chunk
+               Dagger.@spawn sum_array_row!(Out(view(row_sums, (nrows_per_chunk*(i-1)+j):(nrows_per_chunk*(i-1)+j))), In(Dagger.view(BD.chunks[i,1], j:j, :)))
+           end
+       end
+
+# Print the result
+julia> println("Row sum Vector: ", row_sums)
+Row sum Vector: [499.8765, 500.1234, ..., 499.9876]
+```
+
 <!--  -->
 
 ### Explicit Processor Mapping of DArray Blocks

src/memory-spaces.jl

@@ -386,4 +386,56 @@ function will_alias(x_span::MemorySpan, y_span::MemorySpan)
     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
-end
+end
+
+struct ChunkSlice{N} #<: AbstractAliasing
+    chunk::Chunk
+    slices::NTuple{N, Union{Int, AbstractRange{Int}, Colon}}
+end
+
+Base.copyto!(dest::ChunkSlice, src::ChunkSlice) = copyto!(view(unwrap(dest.chunk), dest.slices...), view(unwrap(src.chunk), src.slices...))
+
+@inline function view(c::Chunk, slices...)
+    isa(c.domain, ArrayDomain) || throw(ArgumentError("Chunk must of a DArray (ArrayDomain), got $(typeof(c.domain))"))
+    nd, sz = ndims(c.domain), size(c.domain)
+    nd == length(slices) || throw(DimensionMismatch("Expected $nd slices, got $(length(slices))"))
+    
+    for (i, s) in enumerate(slices)
+        if s isa Int
+            1 ≤ s ≤ sz[i] || throw(ArgumentError("Index $s out of bounds for dimension $i (size $(sz[i]))"))
+        elseif s isa AbstractRange
+            isempty(s) && continue
+            1 ≤ first(s) ≤ last(s) ≤ sz[i] || throw(ArgumentError("Range $s out of bounds for dimension $i (size $(sz[i]))"))
+        elseif s === Colon()
+            continue
+        else
+            throw(ArgumentError("Invalid slice type $(typeof(s)) at dimension $i, Expected Type of Int, AbstractRange, or Colon"))  
+        end
+    end
+
+    return ChunkSlice(c, slices)
+end
+
+view(c::DTask, slices...) = view(fetch(c; raw=true), slices...)
+
+function aliasing(x::ChunkSlice{N}) where N
+    remotecall_fetch(root_worker_id(x.chunk.processor), x.chunk, x.slices) do x, slices
+        x = unwrap(x)
+        v = view(x, slices...)
+        return aliasing(v)
+    end
+end
+
+function move!(dep_mod, to_space::MemorySpace, from_space::MemorySpace, to::ChunkSlice, from::ChunkSlice)
+    to_w = root_worker_id(to_space)
+    remotecall_wait(to_w, dep_mod, to_space, from_space, to, from) do dep_mod, to_space, from_space, to, from
+        to_raw = unwrap(to.chunk)
+        from_w = root_worker_id(from_space)
+        from_raw = to_w == from_w ? unwrap(from.chunk) : remotecall_fetch(unwrap, from_w, from.chunk)
+        from_view = view(from_raw, from.slices...)
+        to_view = view(to_raw, to.slices...)
+        move!(dep_mod, to_space, from_space, to_view, from_view)
+    end
+end
+
+move(from_proc::Processor, to_proc::Processor, slice::ChunkSlice) = view(move(from_proc, to_proc, slice.chunk), slice.slices...)

test/array/allocation.jl

@@ -699,6 +699,44 @@ end
     @test A_v == A[1:8, 1:8]
 end
 
+@testset "Chunk view of DArray" begin
+    A = rand(64, 64)
+    DA = DArray(A, Blocks(8,8))
+    chunk = DA.chunks[1,1]
+
+    @testset "Valid Slices" begin
+        @test view(chunk, :, :)     isa ChunkSlice && view(chunk, 1:8, 1:8)   isa ChunkSlice
+        @test view(chunk, 1:2:7, :) isa ChunkSlice && view(chunk, :, 2:2:8)   isa ChunkSlice
+        @test view(chunk, 1, :)     isa ChunkSlice && view(chunk, :, 1)       isa ChunkSlice
+        @test view(chunk, 3:3, 5:5) isa ChunkSlice && view(chunk, 5:7, 1:2:4) isa ChunkSlice
+        @test view(chunk, 8, 8)     isa ChunkSlice
+        @test view(chunk, 1:0, :)   isa ChunkSlice
+    end
+
+    @testset "Dimension Mismatch" begin
+        @test_throws DimensionMismatch view(chunk, :)
+        @test_throws DimensionMismatch view(chunk, :, :, :)
+    end
+
+    @testset "Int Slice Out of Bounds" begin
+        @test_throws ArgumentError view(chunk, 0, :)
+        @test_throws ArgumentError view(chunk, :, 9)
+        @test_throws ArgumentError view(chunk, 9, 1)
+    end
+
+    @testset "Range Slice Out of Bounds" begin
+        @test_throws ArgumentError view(chunk, 0:5, :)
+        @test_throws ArgumentError view(chunk, 1:8, 5:10)
+        @test_throws ArgumentError view(chunk, 2:2:10, :)
+    end
+
+    @testset "Invalid Slice Types" begin
+        @test_throws DimensionMismatch view(chunk, (1:2, :))
+        @test_throws ArgumentError view(chunk, :, [1, 2])
+    end
+
+end 
+
 @testset "copy/similar" begin
     X1 = ones(Blocks(10, 10), 100, 100)
     X2 = copy(X1)