fixup! datadeps: Add at-stencil helper

Author: jpsamaroo

src/stencil.jl

@@ -3,7 +3,20 @@ const Read = In
 const Write = Out
 const ReadWrite = InOut
 
+function validate_neigh_dist(neigh_dist, size)
+    if !(neigh_dist isa Integer)
+        throw(ArgumentError("Neighborhood distance ($neigh_dist) must be an Integer"))
+    end
+    if neigh_dist <= 0
+        throw(ArgumentError("Neighborhood distance ($neigh_dist) must be greater than 0"))
+    end
+    if any(size .< neigh_dist)
+        throw(ArgumentError("Neighborhood distance ($neigh_dist) must not be larger than the chunk size ($size)"))
+    end
+end
+
 function load_neighbor_edge(arr, dim, dir, neigh_dist)
+    validate_neigh_dist(neigh_dist, size(arr))
     if dir == -1
         start_idx = CartesianIndex(ntuple(i -> i == dim ? (lastindex(arr, i) - neigh_dist + 1) : firstindex(arr, i), ndims(arr)))
         stop_idx = CartesianIndex(ntuple(i -> i == dim ? lastindex(arr, i) : lastindex(arr, i), ndims(arr)))
@@ -15,6 +28,7 @@ function load_neighbor_edge(arr, dim, dir, neigh_dist)
     return move(task_processor(), collect(@view arr[start_idx:stop_idx]))
 end
 function load_neighbor_corner(arr, corner_side, neigh_dist)
+    validate_neigh_dist(neigh_dist, size(arr))
     start_idx = CartesianIndex(ntuple(i -> corner_side[i] == 0 ? (lastindex(arr, i) - neigh_dist + 1) : firstindex(arr, i), ndims(arr)))
     stop_idx = CartesianIndex(ntuple(i -> corner_side[i] == 0 ? lastindex(arr, i) : (firstindex(arr, i) + neigh_dist - 1), ndims(arr)))
     return move(task_processor(), collect(@view arr[start_idx:stop_idx]))
@@ -134,9 +148,9 @@ end
 """
     @stencil begin body end
 
-Allows the specification of stencil operations within a `spawn_datadeps`
-region. The `idx` variable is used to iterate over `range`, which must be a
-`DArray`. An example usage may look like:
+Allows the execution of stencil operations within a `spawn_datadeps` region.
+The `idx` variable is used to iterate over one or more `DArray`s. An example
+usage may look like:
 
 ```julia
 import Dagger: @stencil, Wrap
@@ -146,19 +160,19 @@ A[5, 5] = 1
 B = zeros(Blocks(3, 3), Int, 9, 9)
 Dagger.spawn_datadeps() do
     @stencil begin
-        # Sum values of all neighbors with self
-        A[idx] = sum(@neighbors(A[idx], 1, Wrap()))
-        # Decrement all values by 1
-        A[idx] -= 1
-        # Copy A to B
-        B[idx] = A[idx]
+        # Increment all values by 1
+        A[idx] = A[idx] + 1
+        # Sum values of all neighbors with self and write to B
+        B[idx] = sum(@neighbors(A[idx], 1, Wrap()))
+        # Copy B back to A
+        A[idx] = B[idx]
     end
 end
 ```
 
 Each expression within an `@stencil` region that performs an in-place indexing
 expression like `A[idx] = ...` is transformed into a set of tasks that operate
-on each chunk of `A` or any other arrays specified as `A[idx]`, and within each
+on each chunk of `A` or any other arrays specified as `A[idx]`; within each
 task, elements of that chunk of `A` can be accessed. Elements of multiple
 `DArray`s can be accessed, such as `B[idx]`, so long as `B` has the same size,
 shape, and chunk layout as `A`.
@@ -168,18 +182,20 @@ values around `A[idx]`, at a configurable distance (in this case, 1 element
 distance) and with various kinds of boundary conditions (in this case, `Wrap()`
 specifies wrapping behavior on the boundaries). Neighborhoods are computed with
 respect to neighboring chunks as well - if a neighborhood would overflow from
-the current chunk into one or more neighboring chunks, values from those
-neighboring chunks will be included in the neighborhood.
+the current chunk into a neighboring chunk, values from that neighboring chunk
+will be included in the neighborhood.
 
 Note that, while `@stencil` may look like a `for` loop, it does not follow the
 same semantics; in particular, an expression within `@stencil` occurs "all at
 once" (across all indices) before the next expression occurs. This means that
-`A[idx] = sum(@neighbors(A[idx], 1, Wrap()))` will write the sum of
-neighbors for all `idx` values into `A[idx]` before `A[idx] -= 1` decrements
-the values `A` by 1, and that occurs before any of the values are copied to `B`
-in `B[idx] = A[idx]`. Of course, pipelining and other optimizations may still
-occur, so long as they respect the sequential nature of `@stencil` (just like
-with other operations in `spawn_datadeps`).
+`A[idx] = A[idx] + 1` increments the values `A` by 1, which occurs before
+`B[idx] = sum(@neighbors(A[idx], 1, Wrap()))` writes the sum of neighbors for
+all `idx` values into `B[idx]`, and that occurs before any of the values are
+copied to `A` in `A[idx] = B[idx]`. Of course, pipelining and other optimizations
+may still occur, so long as they respect the sequential nature of `@stencil`
+(just like with other operations in `spawn_datadeps`). Due to this behavior,
+expressions like `A[idx] = sum(@neighbors(A[idx], 1, Wrap()))` are not valid,
+as that would currently cause race conditions and lead to undefined behavior.
 """
 macro stencil(orig_ex)
     @assert Meta.isexpr(orig_ex, :block) "Invalid stencil block: $orig_ex"
@@ -200,7 +216,12 @@ macro stencil(orig_ex)
                 push!(accessed_vars, read_var)
                 push!(read_vars, read_var)
             elseif @capture(read_inner_ex, @neighbors(read_var_[read_idx_], neigh_dist_, boundary_))
-                @assert read_idx == write_idx "Neighborhood access must be at the same index as the write: $read_inner_ex"
+                if read_idx != write_idx
+                    throw(ArgumentError("Neighborhood access must be at the same index as the write: $read_inner_ex"))
+                end
+                if write_var == read_var
+                    throw(ArgumentError("Cannot write to the same variable as the neighborhood access: $read_inner_ex"))
+                end
                 push!(accessed_vars, read_var)
                 push!(read_vars, read_var)
                 neighborhoods[read_var] = (neigh_dist, boundary)

test/array/stencil.jl

@@ -118,6 +118,48 @@ function test_stencil()
         expected_B_pad_val = fill(pad_value*5 + 1*4, 2, 2)
         @test collect(B) == expected_B_pad_val
     end
+
+    @testset "Invalid neighborhood distance" begin
+        A = ones(Blocks(1, 1), Int, 2, 2)
+        @test_throws ArgumentError Dagger.spawn_datadeps() do
+            @stencil begin
+                B[idx] = sum(@neighbors(A[idx], 0, Wrap()))
+            end
+        end
+        @test_throws ArgumentError Dagger.spawn_datadeps() do
+            @stencil begin
+                B[idx] = sum(@neighbors(A[idx], -1, Wrap()))
+            end
+        end
+        @test_throws ArgumentError Dagger.spawn_datadeps() do
+            @stencil begin
+                B[idx] = sum(@neighbors(A[idx], 1.5, Wrap()))
+            end
+        end
+        @test_throws ArgumentError Dagger.spawn_datadeps() do
+            @stencil begin
+                B[idx] = sum(@neighbors(A[idx], 2, Wrap()))
+            end
+        end
+    end
+
+    @testset "Invalid neighborhood access of written variable" begin
+        A = ones(Blocks(1, 1), Int, 2, 2)
+        @test_throws ArgumentError @eval Dagger.spawn_datadeps() do
+            @stencil begin
+                A[idx] = sum(@neighbors(A[idx], 1, Wrap()))
+            end
+        end
+    end
+
+    @testset "Invalid update expression" begin
+        A = ones(Blocks(1, 1), Int, 2, 2)
+        @test_throws ArgumentError @eval Dagger.spawn_datadeps() do
+            @stencil begin
+                A[idx] += 1
+            end
+        end
+    end
 end
 
 @testset "CPU" begin