Merge pull request #674 from JuliaParallel/jps/stencil-tuple-neighborhood

DArray/stencils: Support Tuple for neighborhood distance

Author: jpsamaroo

src/stencil.jl

@@ -4,10 +4,10 @@ const Write = Out
 const ReadWrite = InOut
 
 function validate_neigh_dist(neigh_dist)
-    if !(neigh_dist isa Integer)
-        throw(ArgumentError("Neighborhood distance ($neigh_dist) must be an Integer"))
+    if !(neigh_dist isa Integer) && !(neigh_dist isa Tuple)
+        throw(ArgumentError("Neighborhood distance ($neigh_dist) must be an Integer or Tuple"))
     end
-    if neigh_dist <= 0
+    if any(neigh_dist .<= 0)
         throw(ArgumentError("Neighborhood distance ($neigh_dist) must be greater than 0"))
     end
 end
@@ -18,6 +18,10 @@ function validate_neigh_dist(neigh_dist, size)
     end
 end
 
+get_neigh_dist(neigh_dist::Integer, i::Int) = neigh_dist
+get_neigh_dist(neigh_dist::Tuple, i::Int) = neigh_dist[i]
+
+
 # Load a halo region from a neighboring chunk
 # region_code: N-tuple where each element is -1 (low), 0 (full extent), or +1 (high)
 # For dimensions with code 0, we take the full extent of the array
@@ -27,14 +31,14 @@ function load_neighbor_region(arr, region_code::NTuple{N,Int}, neigh_dist) where
     validate_neigh_dist(neigh_dist, size(arr))
     start_idx = CartesianIndex(ntuple(N) do i
         if region_code[i] == -1
-            lastindex(arr, i) - neigh_dist + 1
+            lastindex(arr, i) - get_neigh_dist(neigh_dist, i) + 1
         else
             firstindex(arr, i)
         end
     end)
     stop_idx = CartesianIndex(ntuple(N) do i
         if region_code[i] == +1
-            firstindex(arr, i) + neigh_dist - 1
+            firstindex(arr, i) + get_neigh_dist(neigh_dist, i) - 1
         else
             lastindex(arr, i)
         end
@@ -88,13 +92,11 @@ function build_halo(neigh_dist, boundary, center, all_halos...)
     N = ndims(center)
     expected_halos = 3^N - 1
     @assert length(all_halos) == expected_halos "Halo mismatch: N=$N expected $expected_halos halos, got $(length(all_halos))"
-    return HaloArray(center, (all_halos...,), ntuple(_->neigh_dist, N))
+    return HaloArray(center, (all_halos...,), ntuple(i->get_neigh_dist(neigh_dist, i), N))
 end
 function load_neighborhood(arr::HaloArray{T,N}, idx) where {T,N}
-    @assert all(arr.halo_width .== arr.halo_width[1])
-    neigh_dist = arr.halo_width[1]
-    start_idx = idx - CartesianIndex(ntuple(_->neigh_dist, ndims(arr)))
-    stop_idx = idx + CartesianIndex(ntuple(_->neigh_dist, ndims(arr)))
+    start_idx = idx - CartesianIndex(ntuple(i->arr.halo_width[i], ndims(arr)))
+    stop_idx = idx + CartesianIndex(ntuple(i->arr.halo_width[i], ndims(arr)))
     return @view arr[start_idx:stop_idx]
 end
 function inner_stencil!(f, output, read_vars)
@@ -126,7 +128,7 @@ function load_boundary_region(pad::Pad, arr, region_code::NTuple{N,Int}, neigh_d
     # For dimensions with code 0, use full array size
     # For dimensions with code -1 or +1, use neigh_dist
     region_size = ntuple(N) do i
-        region_code[i] == 0 ? size(arr, i) : neigh_dist
+        region_code[i] == 0 ? size(arr, i) : get_neigh_dist(neigh_dist, i)
     end
     # FIXME: return Fill(pad.padval, region_size)
     return move(task_processor(), fill(pad.padval, region_size))

test/array/stencil.jl

@@ -137,27 +137,94 @@ function test_stencil()
         end
     end
 
-    @testset "Invalid neighborhood distance" begin
-        A = ones(Blocks(1, 1), Int, 2, 2)
-        B = zeros(Blocks(1, 1), Int, 2, 2)
-        @test_throws_unwrap ArgumentError Dagger.spawn_datadeps() do
-            @stencil begin
-                B[idx] = sum(@neighbors(A[idx], 0, Wrap()))
+    @testset "Tuple neighborhood distance" begin
+        # 1D case: distance (2,)
+        @testset "1D with distance (2,)" begin
+            A = DArray([1, 2, 3, 4, 5, 6], Blocks(2,))
+            B = zeros(Blocks(2,), Int, 6)
+            Dagger.spawn_datadeps() do
+                @stencil begin
+                    B[idx] = sum(@neighbors(A[idx], (2,), Wrap()))
+                end
             end
+            # For each element, neighbors at distance 2 in 1D: [-2, -1, 0, 1, 2]
+            # B[1] neighbors: A[5], A[6], A[1], A[2], A[3] (wrapping) = 5+6+1+2+3 = 17
+            # B[2] neighbors: A[6], A[1], A[2], A[3], A[4] = 6+1+2+3+4 = 16
+            # B[3] neighbors: A[1], A[2], A[3], A[4], A[5] = 1+2+3+4+5 = 15
+            # B[4] neighbors: A[2], A[3], A[4], A[5], A[6] = 2+3+4+5+6 = 20
+            # B[5] neighbors: A[3], A[4], A[5], A[6], A[1] = 3+4+5+6+1 = 19
+            # B[6] neighbors: A[4], A[5], A[6], A[1], A[2] = 4+5+6+1+2 = 18
+            expected_B_1d = [17, 16, 15, 20, 19, 18]
+            @test collect(B) == expected_B_1d
         end
-        @test_throws_unwrap ArgumentError Dagger.spawn_datadeps() do
-            @stencil begin
-                B[idx] = sum(@neighbors(A[idx], -1, Wrap()))
+
+        # 2D case: distance (1, 2) - different per dimension
+        @testset "2D with distance (1, 2)" begin
+            A = DArray(reshape(1:12, 3, 4), Blocks(1, 2))
+            B = zeros(Blocks(1, 2), Int, 3, 4)
+            Dagger.spawn_datadeps() do
+                @stencil begin
+                    B[idx] = sum(@neighbors(A[idx], (1, 2), Wrap()))
+                end
             end
+            # Distance (1, 2) means:
+            # - dimension 1 (rows): offsets -1, 0, 1
+            # - dimension 2 (cols): offsets -2, -1, 0, 1, 2
+            # Total neighborhood size: 3 * 5 = 15 elements
+            expected_B_2d = zeros(Int, 3, 4)
+            for i in 1:3, j in 1:4
+                sum_val = 0
+                for di in -1:1, dj in -2:2
+                    row = mod1(i+di, 3)
+                    col = mod1(j+dj, 4)
+                    sum_val += A[row, col]
+                end
+                expected_B_2d[i, j] = sum_val
+            end
+            @test collect(B) == expected_B_2d
         end
-        @test_throws_unwrap ArgumentError Dagger.spawn_datadeps() do
-            @stencil begin
-                B[idx] = sum(@neighbors(A[idx], 1.5, Wrap()))
+
+        # 3D case: distance (1, 2, 1) - different per dimension
+        @testset "3D with distance (1, 2, 1)" begin
+            # Need chunk sizes >= 2*distance+1 for each dimension
+            # distance (1, 2, 1) requires chunks >= (3, 5, 3)
+            A = DArray(reshape(1:120, 4, 5, 6), Blocks(4, 5, 3))
+            B = zeros(Blocks(4, 5, 3), Int, 4, 5, 6)
+            Dagger.spawn_datadeps() do
+                @stencil begin
+                    B[idx] = sum(@neighbors(A[idx], (1, 2, 1), Wrap()))
+                end
             end
+            # Distance (1, 2, 1) means:
+            # - dimension 1: offsets -1, 0, 1 (3 elements)
+            # - dimension 2: offsets -2, -1, 0, 1, 2 (5 elements)
+            # - dimension 3: offsets -1, 0, 1 (3 elements)
+            # Total neighborhood size: 3 * 5 * 3 = 45 elements
+            expected_B_3d = zeros(Int, 4, 5, 6)
+            for i in 1:4, j in 1:5, k in 1:6
+                sum_val = 0
+                for di in -1:1, dj in -2:2, dk in -1:1
+                    row = mod1(i+di, 4)
+                    col = mod1(j+dj, 5)
+                    depth = mod1(k+dk, 6)
+                    sum_val += A[row, col, depth]
+                end
+                expected_B_3d[i, j, k] = sum_val
+            end
+            @test collect(B) == expected_B_3d
         end
-        @test_throws_unwrap ArgumentError Dagger.spawn_datadeps() do
-            @stencil begin
-                B[idx] = sum(@neighbors(A[idx], 2, Wrap()))
+    end
+
+    @testset "Invalid neighborhood distance" begin
+        A = ones(Blocks(1, 1), Int, 2, 2)
+        B = zeros(Blocks(1, 1), Int, 2, 2)
+        for value in [0, -1, 1.5, 2]
+            for dist in [value, (value,)]
+                @test_throws_unwrap ArgumentError Dagger.spawn_datadeps() do
+                    @stencil begin
+                        B[idx] = sum(@neighbors(A[idx], value, Wrap()))
+                    end
+                end
             end
         end
     end
@@ -196,3 +263,4 @@ end
         end
     end
 end
+