Merge pull request #676 from JuliaParallel/jps/stencil-reflect

stencils: Add Reflect boundary condition

Author: jpsamaroo

docs/src/stencils.md

@@ -8,7 +8,7 @@ The fundamental structure of a `@stencil` block involves iterating over an impli
 
 ```julia
 using Dagger
-import Dagger: @stencil, Wrap, Pad
+import Dagger: @stencil, Wrap, Pad, Reflect
 
 # Initialize a DArray
 A = zeros(Blocks(2, 2), Int, 4, 4)
@@ -35,6 +35,9 @@ The true power of stencils comes from accessing neighboring elements. The `@neig
 - `boundary_condition`: Defines how to handle accesses beyond the array boundaries. Available conditions are:
     - `Wrap()`: Wraps around to the other side of the array.
     - `Pad(value)`: Pads with a specified `value`.
+    - `Reflect(symmetric)`: Reflects values back into the array at boundaries. The `symmetric` boolean controls whether the edge element is included in the reflection:
+        - `Reflect(true)` (symmetric): Edge element IS repeated. For array `[a,b,c,d]`, extends as `[...,c,b,a,a,b,c,d,d,c,b,...]`.
+        - `Reflect(false)` (mirror): Edge element NOT repeated. For array `[a,b,c,d]`, extends as `[...,d,c,b,a,b,c,d,c,b,a,...]`.
 
 ### Example: Averaging Neighbors with `Wrap`
 
@@ -93,6 +96,59 @@ expected_B_padded = [
 @assert collect(B) == expected_B_padded
 ```
 
+### Example: Smoothing with `Reflect`
+
+The `Reflect` boundary condition mirrors values at the edges, which is useful for operations like smoothing or image processing where you want to avoid artificial discontinuities at boundaries.
+
+#### Symmetric Reflection (`Reflect(true)`)
+
+With symmetric reflection, the edge element is repeated in the reflection:
+
+```julia
+import Dagger: Reflect
+
+# Simple 1D example to illustrate symmetric reflection
+# Array [1, 2, 3, 4] extends as [..., 3, 2, 1, 1, 2, 3, 4, 4, 3, 2, ...]
+#                                         ^edge^         ^edge^
+A = DArray([1, 2, 3, 4], Blocks(2))
+B = zeros(Blocks(2), Int, 4)
+
+Dagger.spawn_datadeps() do
+    @stencil begin
+        B[idx] = sum(@neighbors(A[idx], 1, Reflect(true)))
+    end
+end
+
+# B[1]: indices 0,1,2 -> 0 reflects to 1, so [1,1,2] = 4
+# B[2]: indices 1,2,3 -> all in bounds, [1,2,3] = 6
+# B[3]: indices 2,3,4 -> all in bounds, [2,3,4] = 9
+# B[4]: indices 3,4,5 -> 5 reflects to 4, so [3,4,4] = 11
+@assert collect(B) == [4, 6, 9, 11]
+```
+
+#### Mirror Reflection (`Reflect(false)`)
+
+With mirror reflection, the edge element is NOT repeated:
+
+```julia
+# Array [1, 2, 3, 4] extends as [..., 4, 3, 2, 1, 2, 3, 4, 3, 2, 1, ...]
+#                                            ^edge^   ^edge^
+A = DArray([1, 2, 3, 4], Blocks(2))
+B = zeros(Blocks(2), Int, 4)
+
+Dagger.spawn_datadeps() do
+    @stencil begin
+        B[idx] = sum(@neighbors(A[idx], 1, Reflect(false)))
+    end
+end
+
+# B[1]: indices 0,1,2 -> 0 reflects to 2, so [2,1,2] = 5
+# B[2]: indices 1,2,3 -> all in bounds, [1,2,3] = 6
+# B[3]: indices 2,3,4 -> all in bounds, [2,3,4] = 9
+# B[4]: indices 3,4,5 -> 5 reflects to 3, so [3,4,3] = 10
+@assert collect(B) == [5, 6, 9, 10]
+```
+
 ## Sequential Semantics
 
 Expressions within a `@stencil` block are executed sequentially in terms of their effect on the data. This means that the result of one statement is visible to the subsequent statements, as if they were applied "all at once" across all indices before the next statement begins.

src/stencil.jl

@@ -21,7 +21,6 @@ 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
@@ -72,13 +71,17 @@ function select_neighborhood_chunks(chunks, idx, neigh_dist, boundary)
         new_idx = idx + chunk_offset
 
         if is_past_boundary(size(chunks), new_idx)
+            # Compute which dimensions are actually past boundary
+            boundary_dims = ntuple(N) do d
+                new_idx[d] < 1 || new_idx[d] > size(chunks)[d]
+            end
             if boundary_has_transition(boundary)
                 new_idx = boundary_transition(boundary, new_idx, size(chunks))
             else
                 new_idx = idx
             end
             chunk = chunks[new_idx]
-            push!(accesses, Dagger.@spawn load_boundary_region(boundary, chunk, region_code, neigh_dist))
+            push!(accesses, Dagger.@spawn load_boundary_region(boundary, chunk, region_code, neigh_dist, boundary_dims))
         else
             chunk = chunks[new_idx]
             push!(accesses, Dagger.@spawn load_neighbor_region(chunk, region_code, neigh_dist))
@@ -113,17 +116,23 @@ end
 
 is_past_boundary(size, idx) = any(ntuple(i -> idx[i] < 1 || idx[i] > size[i], length(size)))
 
+"""
+Wrap boundary condition. Non-local accesses wrap around to the other side of the array.
+"""
 struct Wrap end
 boundary_has_transition(::Wrap) = true
 boundary_transition(::Wrap, idx, size) =
     CartesianIndex(ntuple(i -> mod1(idx[i], size[i]), length(size)))
-load_boundary_region(::Wrap, arr, region_code, neigh_dist) = load_neighbor_region(arr, region_code, neigh_dist)
+load_boundary_region(::Wrap, arr, region_code, neigh_dist, boundary_dims) = load_neighbor_region(arr, region_code, neigh_dist)
 
+"""
+Pad boundary condition. Non-local accesses are padded with a specified value.
+"""
 struct Pad{T}
     padval::T
 end
 boundary_has_transition(::Pad) = false
-function load_boundary_region(pad::Pad, arr, region_code::NTuple{N,Int}, neigh_dist) where N
+function load_boundary_region(pad::Pad, arr, region_code::NTuple{N,Int}, neigh_dist, boundary_dims::NTuple{N,Bool}) where N
     # Compute the size of this halo region
     # For dimensions with code 0, use full array size
     # For dimensions with code -1 or +1, use neigh_dist
@@ -134,6 +143,76 @@ function load_boundary_region(pad::Pad, arr, region_code::NTuple{N,Int}, neigh_d
     return move(task_processor(), fill(pad.padval, region_size))
 end
 
+"""
+Reflect boundary condition. Non-local accesses are reflected back into the array.
+If `symm` is true, the reflected values include the nearest center elements.
+If `symm` is false, the reflected values do not include the nearest center elements.
+"""
+struct Reflect{Symmetric} end
+Reflect(symm::Bool) = Reflect{symm}()
+boundary_has_transition(::Reflect) = true
+# Clamp to valid chunk indices - we stay at the boundary chunk
+boundary_transition(::Reflect, idx, size) =
+    CartesianIndex(ntuple(i -> clamp(idx[i], 1, size[i]), length(size)))
+function load_boundary_region(::Reflect{Symm}, arr, region_code::NTuple{N,Int}, neigh_dist, boundary_dims::NTuple{N,Bool}) where {N, Symm}
+    # Only flip region_code for dimensions that are BOTH:
+    # 1. Non-zero in region_code (we're accessing a neighbor in that dimension)
+    # 2. Actually past boundary (boundary_dims[i] is true)
+    # For dimensions not past boundary, keep the original region_code behavior
+    flipped_code = ntuple(N) do i
+        if region_code[i] != 0 && boundary_dims[i]
+            # This dimension needs reflection - flip the code
+            -region_code[i]
+        else
+            # Keep original code (either 0, or not past boundary)
+            region_code[i]
+        end
+    end
+
+    # For non-symmetric (mirror), skip 1 element to exclude the edge
+    # For symmetric, include the edge element (skip = 0)
+    # Only apply skip to dimensions that are being reflected
+    skip = Symm ? 0 : 1
+
+    # Compute region indices
+    start_idx = CartesianIndex(ntuple(N) do i
+        needs_skip = boundary_dims[i] && region_code[i] != 0
+        actual_skip = needs_skip ? skip : 0
+        if flipped_code[i] == -1
+            # Taking from end (high side)
+            lastindex(arr, i) - get_neigh_dist(neigh_dist, i) + 1 - actual_skip
+        elseif flipped_code[i] == +1
+            # Taking from start (low side)
+            firstindex(arr, i) + actual_skip
+        else
+            firstindex(arr, i)
+        end
+    end)
+    stop_idx = CartesianIndex(ntuple(N) do i
+        needs_skip = boundary_dims[i] && region_code[i] != 0
+        actual_skip = needs_skip ? skip : 0
+        if flipped_code[i] == +1
+            firstindex(arr, i) + get_neigh_dist(neigh_dist, i) - 1 + actual_skip
+        elseif flipped_code[i] == -1
+            lastindex(arr, i) - actual_skip
+        else
+            lastindex(arr, i)
+        end
+    end)
+
+    region = move(task_processor(), collect(@view arr[start_idx:stop_idx]))
+
+    # Reverse only along dimensions that are actually being reflected
+    # (both non-zero in region_code AND past boundary)
+    for i in 1:N
+        if region_code[i] != 0 && boundary_dims[i]
+            region = reverse(region, dims=i)
+        end
+    end
+
+    return region
+end
+
 """
     @stencil begin body end
 

test/array/stencil.jl

@@ -1,4 +1,4 @@
-import Dagger: @stencil, Wrap, Pad
+import Dagger: @stencil, Wrap, Pad, Reflect
 
 function test_stencil()
     @testset "Simple assignment" begin
@@ -67,6 +67,99 @@ function test_stencil()
         @test collect(B) == expected_B_pad
     end
 
+    @testset "Reflect boundary (symmetric)" begin
+        # Test symmetric reflection (edge element IS included/repeated)
+        # For A = [1, 2, 3, 4] with Reflect(true):
+        # idx=0 → 1, idx=-1 → 2 (reflection includes edge)
+        # idx=5 → 4, idx=6 → 3 (reflection includes edge)
+        A = DArray([1, 2, 3, 4], Blocks(2))
+        B = zeros(Blocks(2), Int, 4)
+        Dagger.spawn_datadeps() do
+            @stencil begin
+                B[idx] = sum(@neighbors(A[idx], 1, Reflect(true)))
+            end
+        end
+        # B[1]: neighbors at indices 0, 1, 2 -> reflected 0 becomes 1, so [1, 1, 2] = 4
+        # B[2]: neighbors at indices 1, 2, 3 -> [1, 2, 3] = 6
+        # B[3]: neighbors at indices 2, 3, 4 -> [2, 3, 4] = 9
+        # B[4]: neighbors at indices 3, 4, 5 -> reflected 5 becomes 4, so [3, 4, 4] = 11
+        expected_B_symm = [4, 6, 9, 11]
+        @test collect(B) == expected_B_symm
+    end
+
+    @testset "Reflect boundary (mirror)" begin
+        # Test mirror reflection (edge element NOT included/repeated)
+        # For A = [1, 2, 3, 4] with Reflect(false):
+        # idx=0 → 2, idx=-1 → 3 (reflection skips edge)
+        # idx=5 → 3, idx=6 → 2 (reflection skips edge)
+        A = DArray([1, 2, 3, 4], Blocks(2))
+        B = zeros(Blocks(2), Int, 4)
+        Dagger.spawn_datadeps() do
+            @stencil begin
+                B[idx] = sum(@neighbors(A[idx], 1, Reflect(false)))
+            end
+        end
+        # B[1]: neighbors at indices 0, 1, 2 -> reflected 0 becomes 2, so [2, 1, 2] = 5
+        # B[2]: neighbors at indices 1, 2, 3 -> [1, 2, 3] = 6
+        # B[3]: neighbors at indices 2, 3, 4 -> [2, 3, 4] = 9
+        # B[4]: neighbors at indices 3, 4, 5 -> reflected 5 becomes 3, so [3, 4, 3] = 10
+        expected_B_mirror = [5, 6, 9, 10]
+        @test collect(B) == expected_B_mirror
+    end
+
+    @testset "Reflect boundary 2D (symmetric)" begin
+        # Test 2D symmetric reflection with a gradient pattern
+        A = DArray(reshape(1:16, 4, 4), Blocks(2, 2))
+        B = zeros(Blocks(2, 2), Int, 4, 4)
+        Dagger.spawn_datadeps() do
+            @stencil begin
+                B[idx] = sum(@neighbors(A[idx], 1, Reflect(true)))
+            end
+        end
+        # Symmetric: idx < 1 → 1 - idx, idx > size → 2*size + 1 - idx
+        A_collected = collect(A)
+        expected_B_symm = zeros(Int, 4, 4)
+        for i in 1:4, j in 1:4
+            sum_val = 0
+            for di in -1:1, dj in -1:1
+                ni, nj = i + di, j + dj
+                # Apply symmetric reflection logic
+                # For symmetric: idx < 1 → 1 - idx, idx > size → 2*size + 1 - idx
+                ni = ni < 1 ? 1 - ni : (ni > 4 ? 2*4 + 1 - ni : ni)
+                nj = nj < 1 ? 1 - nj : (nj > 4 ? 2*4 + 1 - nj : nj)
+                sum_val += A_collected[ni, nj]
+            end
+            expected_B_symm[i, j] = sum_val
+        end
+        @test collect(B) == expected_B_symm
+    end
+
+    @testset "Reflect boundary 2D (mirror)" begin
+        # Test 2D mirror reflection with a gradient pattern
+        A = DArray(reshape(1:16, 4, 4), Blocks(2, 2))
+        B = zeros(Blocks(2, 2), Int, 4, 4)
+        Dagger.spawn_datadeps() do
+            @stencil begin
+                B[idx] = sum(@neighbors(A[idx], 1, Reflect(false)))
+            end
+        end
+        # Mirror: idx < 1 → 2 - idx, idx > size → 2*size - idx
+        A_collected = collect(A)
+        expected_B_mirror = zeros(Int, 4, 4)
+        for i in 1:4, j in 1:4
+            sum_val = 0
+            for di in -1:1, dj in -1:1
+                ni, nj = i + di, j + dj
+                # Apply mirror reflection logic
+                ni = ni < 1 ? 2 - ni : (ni > 4 ? 2*4 - ni : ni)
+                nj = nj < 1 ? 2 - nj : (nj > 4 ? 2*4 - nj : nj)
+                sum_val += A_collected[ni, nj]
+            end
+            expected_B_mirror[i, j] = sum_val
+        end
+        @test collect(B) == expected_B_mirror
+    end
+
     @testset "Multiple expressions" begin
         A = zeros(Blocks(2, 2), Int, 4, 4)
         B = zeros(Blocks(2, 2), Int, 4, 4)