stencils: Add tests and documentation (#619)

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>

Author: jpsamaroo

docs/src/index.md

@@ -361,6 +361,39 @@ DA = rand(Blocks(32, 32), 256, 128)
 collect(DA) # returns a `Matrix{Float64}`
 ```
 
+-----
+
+## Quickstart: Stencil Operations
+
+Dagger's `@stencil` macro allows for easy specification of stencil operations on `DArray`s, often used in simulations and image processing. These operations typically involve updating an element based on the values of its neighbors.
+
+For more details: [Stencil Operations](@ref)
+
+### Applying a Simple Stencil
+
+Here's how to apply a stencil that averages each element with its immediate neighbors, using a `Wrap` boundary condition (where neighbor access at the array edges wrap around).
+
+```julia
+using Dagger
+import Dagger: @stencil, Wrap
+
+# Create a 5x5 DArray, partitioned into 2x2 blocks
+A = rand(Blocks(2, 2), 5, 5)
+B = zeros(Blocks(2,2), 5, 5)
+
+Dagger.spawn_datadeps() do
+    @stencil begin
+        # For each element in A, calculate the sum of its 3x3 neighborhood
+        # (including itself) and store the average in B.
+        # Values outside the array bounds are determined by Wrap().
+        B[idx] = sum(@neighbors(A[idx], 1, Wrap())) / 9.0
+    end
+end
+
+# B now contains the averaged values.
+```
+In this example, `idx` refers to the coordinates of each element being processed. `@neighbors(A[idx], 1, Wrap())` fetches the 3x3 neighborhood around `A[idx]`. The `1` indicates a neighborhood distance of 1 from the central element, and `Wrap()` specifies the boundary behavior.
+
 ## Quickstart: Datadeps
 
 Datadeps is a feature in Dagger.jl that facilitates parallelism control within designated regions, allowing tasks to write to their arguments while ensuring dependencies are respected.

docs/src/stencils.jl

@@ -0,0 +1,183 @@
+# Stencil Operations
+
+The `@stencil` macro in Dagger.jl provides a convenient way to perform stencil computations on `DArray`s. It operates within a `Dagger.spawn_datadeps()` block and allows you to define operations that apply to each element of a `DArray`, potentially accessing values from each element's neighbors.
+
+## Basic Usage
+
+The fundamental structure of a `@stencil` block involves iterating over an implicit index, named `idx` in the following example , which represents the coordinates of an element in the processed `DArray`s.
+
+```julia
+using Dagger
+import Dagger: @stencil, Wrap, Pad
+
+# Initialize a DArray
+A = zeros(Blocks(2, 2), Int, 4, 4)
+
+Dagger.spawn_datadeps() do
+    @stencil begin
+        A[idx] = 1 # Assign 1 to every element of A
+    end
+end
+
+@assert all(collect(A) .== 1)
+```
+
+In this example, `A[idx] = 1` is executed for each chunk of `A`. The `idx` variable corresponds to the indices within each chunk.
+
+## Neighborhood Access with `@neighbors`
+
+The true power of stencils comes from accessing neighboring elements. The `@neighbors` macro facilitates this.
+
+`@neighbors(array[idx], distance, boundary_condition)`
+
+- `array[idx]`: The array and current index from which to find neighbors.
+- `distance`: An integer specifying the extent of the neighborhood (e.g., `1` for a 3x3 neighborhood in 2D).
+- `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`.
+
+### Example: Averaging Neighbors with `Wrap`
+
+```julia
+import Dagger: Wrap
+
+# Initialize a DArray
+A = ones(Blocks(1, 1), Int, 3, 3)
+A[2,2] = 10 # Central element has a different value
+B = zeros(Blocks(1, 1), Float64, 3, 3)
+
+Dagger.spawn_datadeps() do
+    @stencil begin
+        # Calculate the average of the 3x3 neighborhood (including the center)
+        B[idx] = sum(@neighbors(A[idx], 1, Wrap())) / 9.0
+    end
+end
+
+# Manually calculate expected B for verification
+expected_B = zeros(Float64, 3, 3)
+A_collected = collect(A)
+for r in 1:3, c in 1:3
+    local_sum = 0.0
+    for dr in -1:1, dc in -1:1
+        nr, nc = mod1(r+dr, 3), mod1(c+dc, 3)
+        local_sum += A_collected[nr, nc]
+    end
+    expected_B[r,c] = local_sum / 9.0
+end
+
+@assert collect(B) ≈ expected_B
+```
+
+### Example: Convolution with `Pad`
+
+```julia
+import Pad
+
+# Initialize a DArray
+A = ones(Blocks(2, 2), Int, 4, 4)
+B = zeros(Blocks(2, 2), Int, 4, 4)
+
+Dagger.spawn_datadeps() do
+    @stencil begin
+        B[idx] = sum(@neighbors(A[idx], 1, Pad(0))) # Pad with 0
+    end
+end
+
+# Expected result for a 3x3 sum filter with zero padding
+expected_B_padded = [
+    4 6 6 4;
+    6 9 9 6;
+    6 9 9 6;
+    4 6 6 4
+]
+@assert collect(B) == expected_B_padded
+```
+
+## 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.
+
+```julia
+A = zeros(Blocks(2, 2), Int, 4, 4)
+B = zeros(Blocks(2, 2), Int, 4, 4)
+
+Dagger.spawn_datadeps() do
+    @stencil begin
+        A[idx] = 1  # First, A is initialized
+        B[idx] = A[idx] * 2       # Then, B is computed using the new values of A
+    end
+end
+
+expected_A = [1 for r in 1:4, c in 1:4]
+expected_B_seq = expected_A .* 2
+
+@assert collect(A) == expected_A
+@assert collect(B) == expected_B_seq
+```
+
+## Operations on Multiple `DArray`s
+
+You can read from and write to multiple `DArray`s within a single `@stencil` block, provided they have compatible chunk structures.
+
+```julia
+A = ones(Blocks(1, 1), Int, 2, 2)
+B = DArray(fill(3, 2, 2), Blocks(1, 1))
+C = zeros(Blocks(1, 1), Int, 2, 2)
+
+Dagger.spawn_datadeps() do
+    @stencil begin
+        C[idx] = A[idx] + B[idx]
+    end
+end
+@assert all(collect(C) .== 4)
+```
+
+## Example: Game of Life
+
+The following demonstrates a more complex example: Conway's Game of Life.
+
+```julia
+# Ensure Plots and other necessary packages are available for the example
+using Plots
+
+N = 27 # Size of one dimension of a tile
+nt = 3 # Number of tiles in each dimension (results in nt x nt grid of tiles)
+niters = 10 # Number of iterations for the animation
+
+tiles = zeros(Blocks(N, N), Bool, N*nt, N*nt)
+outputs = zeros(Blocks(N, N), Bool, N*nt, N*nt)
+
+# Create a fun initial state (e.g., a glider and some random noise)
+tiles[13, 14] = true
+tiles[14, 14] = true
+tiles[15, 14] = true
+tiles[15, 15] = true
+tiles[14, 16] = true
+# Add some random noise in one of the tiles
+@view(tiles[(2N+1):3N, (2N+1):3N]) .= rand(Bool, N, N)
+
+
+
+anim = @animate for _ in 1:niters
+    Dagger.spawn_datadeps() do
+        @stencil begin
+            outputs[idx] = begin
+                nhood = @neighbors(tiles[idx], 1, Wrap())
+                neighs = sum(nhood) - tiles[idx] # Sum neighborhood, but subtract own value
+                if tiles[idx] && neighs < 2
+                    0 # Dies of underpopulation
+                elseif tiles[idx] && neighs > 3
+                    0 # Dies of overpopulation
+                elseif !tiles[idx] && neighs == 3
+                    1 # Becomes alive by reproduction
+                else
+                    tiles[idx] # Keeps its prior value
+                end
+            end
+            tiles[idx] = outputs[idx] # Update tiles for the next iteration
+        end
+    end
+    heatmap(Int.(collect(outputs))) # Generate a heatmap visualization
+end
+path = mp4(anim; fps=5, show_msg=true).filename # Create an animation of the heatmaps over time
+```

docs/src/stencils.md

@@ -0,0 +1,120 @@
+import Dagger: @stencil, Wrap, Pad
+
+@testset "@stencil" begin
+    @testset "Simple assignment" begin
+        A = zeros(Blocks(2, 2), Int, 4, 4)
+        Dagger.spawn_datadeps() do
+            @stencil begin
+                A[idx] = 1
+            end
+        end
+        @test all(collect(A) .== 1)
+    end
+
+    @testset "Wrap boundary" begin
+        A = zeros(Blocks(2, 2), Int, 4, 4)
+        A[1,1] = 10
+        B = zeros(Blocks(2, 2), Int, 4, 4)
+        Dagger.spawn_datadeps() do
+            @stencil begin
+                B[idx] = sum(@neighbors(A[idx], 1, Wrap()))
+            end
+        end
+        # Expected result after convolution with wrap around
+        # Corner element (1,1) will sum its 3 neighbors + itself (10) + 5 wrapped around neighbors
+        # For A[1,1], neighbors are A[4,4], A[4,1], A[4,2], A[1,4], A[1,2], A[2,4], A[2,1], A[2,2]
+        # Since only A[1,1] is 10 and others are 0, sum for B[1,1] will be 10 (A[1,1])
+        # Sum for B[1,2] will be A[1,1] = 10
+        # Sum for B[2,1] will be A[1,1] = 10
+        # Sum for B[2,2] will be A[1,1] = 10
+        # Sum for B[4,4] will be A[1,1] = 10
+        # ... and so on for elements that wrap around to include A[1,1]
+        expected_B_calc = zeros(Int, 4, 4)
+        for i in 1:4, j in 1:4
+            sum_val = 0
+            for ni in -1:1, nj in -1:1
+                # Apply wrap around logic for neighbors
+                row = mod1(i+ni, 4)
+                col = mod1(j+nj, 4)
+                if row == 1 && col == 1 # Check if the wrapped neighbor is A[1,1]
+                    sum_val += 10
+                end
+            end
+            expected_B_calc[i,j] = sum_val
+        end
+        @test collect(B) == expected_B_calc
+    end
+
+    @testset "Pad boundary" begin
+        A = DArray(ones(Int, 4, 4), Blocks(2, 2))
+        B = DArray(zeros(Int, 4, 4), Blocks(2, 2))
+        Dagger.spawn_datadeps() do
+            @stencil begin
+                B[idx] = sum(@neighbors(A[idx], 1, Pad(0)))
+            end
+        end
+        # Expected result after convolution with zero padding
+        # Inner elements (e.g., B[2,2]) will sum 9 (3x3 neighborhood of 1s)
+        # Edge elements (e.g., B[1,2]) will sum 6 (2x3 neighborhood of 1s, 3 zeros from padding)
+        # Corner elements (e.g., B[1,1]) will sum 4 (2x2 neighborhood of 1s, 5 zeros from padding)
+        expected_B_pad = [
+            4 6 6 4;
+            6 9 9 6;
+            6 9 9 6;
+            4 6 6 4
+        ]
+        @test collect(B) == expected_B_pad
+    end
+
+    @testset "Multiple expressions" begin
+        A = zeros(Blocks(2, 2), Int, 4, 4)
+        B = zeros(Blocks(2, 2), Int, 4, 4)
+        Dagger.spawn_datadeps() do
+            @stencil begin
+                A[idx] = 1
+                B[idx] = A[idx] * 2
+            end
+        end
+        expected_A_multi = [1 for r in 1:4, c in 1:4]
+        expected_B_multi = expected_A_multi .* 2
+        @test collect(A) == expected_A_multi
+        @test collect(B) == expected_B_multi
+    end
+
+    @testset "Multiple DArrays" begin
+        A = ones(Blocks(2, 2), Int, 4, 4)
+        B = DArray(fill(2, 4, 4), Blocks(2, 2))
+        C = zeros(Blocks(2, 2), Int, 4, 4)
+        Dagger.spawn_datadeps() do
+            @stencil begin
+                C[idx] = A[idx] + B[idx]
+            end
+        end
+        @test all(collect(C) .== 3)
+    end
+
+    @testset "Pad boundary with non-zero value" begin
+        A = ones(Blocks(1, 1), Int, 2, 2) # Simpler 2x2 case
+        B = zeros(Blocks(1, 1), Int, 2, 2)
+        pad_value = 5
+        Dagger.spawn_datadeps() do
+            @stencil begin
+                B[idx] = sum(@neighbors(A[idx], 1, Pad(pad_value)))
+            end
+        end
+        # For A = [1 1; 1 1] and Pad(5)
+        # B[1,1] neighbors considering a 3x3 neighborhood around A[1,1]:
+        # P P P
+        # P A11 A12
+        # P A21 A22
+        # Values:
+        # 5 5 5
+        # 5 1 1
+        # 5 1 1
+        # Sum = 5*5 (for the padded values) + 1*4 (for the actual values from A) = 25 + 4 = 29.
+        # This logic applies to all elements in B because the array A is small (2x2) and the neighborhood is 1.
+        # Every element's 3x3 neighborhood will include 5 padded values and the 4 values of A.
+        expected_B_pad_val = fill(pad_value*5 + 1*4, 2, 2)
+        @test collect(B) == expected_B_pad_val
+    end
+end

test/array/stencil.jl

@@ -21,6 +21,7 @@ tests = [
     ("Array - LinearAlgebra - Cholesky", "array/linalg/cholesky.jl"),
     ("Array - LinearAlgebra - LU", "array/linalg/lu.jl"),
     ("Array - Random", "array/random.jl"),
+    ("Array - Stencils", "array/stencil.jl"),
     ("Caching", "cache.jl"),
     ("Disk Caching", "diskcaching.jl"),
     ("File IO", "file-io.jl"),