stencils: Add GPU support

Author: jpsamaroo

Project.toml

@@ -7,7 +7,6 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 DistributedNext = "fab6aee4-877b-4bac-a744-3eca44acbb6f"
-FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
@@ -52,7 +51,6 @@ DataFrames = "1"
 DataStructures = "0.18"
 DistributedNext = "1.0.0"
 Distributions = "0.25"
-FillArrays = "1.11.0"
 GraphViz = "0.2"
 Graphs = "1"
 JSON3 = "1"

src/Dagger.jl

@@ -32,8 +32,6 @@ import TimespanLogging: timespan_start, timespan_finish
 
 import Adapt
 
-import FillArrays: Fill
-
 # Preferences
 import Preferences: @load_preference, @set_preferences!
 

src/stencil.jl

@@ -11,12 +11,13 @@ function load_neighbor_edge(arr, dim, dir, neigh_dist)
         start_idx = CartesianIndex(ntuple(i -> i == dim ? firstindex(arr, i) : firstindex(arr, i), ndims(arr)))
         stop_idx = CartesianIndex(ntuple(i -> i == dim ? (firstindex(arr, i) + neigh_dist - 1) : lastindex(arr, i), ndims(arr)))
     end
-    return collect(@view arr[start_idx:stop_idx])
+    # FIXME: Don't collect
+    return move(thunk_processor(), collect(@view arr[start_idx:stop_idx]))
 end
 function load_neighbor_corner(arr, corner_side, neigh_dist)
     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 collect(@view arr[start_idx:stop_idx])
+    return move(thunk_processor(), collect(@view arr[start_idx:stop_idx]))
 end
 function select_neighborhood_chunks(chunks, idx, neigh_dist, boundary)
     @assert neigh_dist isa Integer && neigh_dist > 0 "Neighborhood distance must be an Integer greater than 0"
@@ -71,19 +72,30 @@ function select_neighborhood_chunks(chunks, idx, neigh_dist, boundary)
     @assert length(accesses) == 1+2*ndims(chunks)+2^ndims(chunks) "Accesses mismatch: $(length(accesses))"
     return accesses
 end
-function build_halo(neigh_dist, boundary, center::Array{T,N}, all_neighbors...) where {T,N}
-    # FIXME: Don't collect views
-    edges = collect.(all_neighbors[1:(2*N)])
-    corners = collect.(all_neighbors[((2^N)+1):end])
+function build_halo(neigh_dist, boundary, center, all_neighbors...)
+    N = ndims(center)
+    edges = all_neighbors[1:(2*N)]
+    corners = all_neighbors[((2^N)+1):end]
     @assert length(edges) == 2*N && length(corners) == 2^N "Halo mismatch: edges=$(length(edges)) corners=$(length(corners))"
-    arr = HaloArray(center, (edges...,), (corners...,), ntuple(_->neigh_dist, N))
-    return arr
+    return HaloArray(center, (edges...,), (corners...,), ntuple(_->neigh_dist, N))
 end
-function load_neighborhood(arr::HaloArray{T,N}, idx, neigh_dist) where {T,N}
+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)))
-    # FIXME: Don't collect HaloArray view
-    return collect(@view arr[start_idx:stop_idx])
+    return @view arr[start_idx:stop_idx]
+end
+function inner_stencil!(f, output, read_vars)
+    processor = thunk_processor()
+    inner_stencil_proc!(processor, f, output, read_vars)
+end
+# Non-KA (for CPUs)
+function inner_stencil_proc!(::ThreadProc, f, output, read_vars)
+    for idx in CartesianIndices(output)
+        f(idx, output, read_vars)
+    end
+    return
 end
 
 is_past_boundary(size, idx) = any(ntuple(i -> idx[i] < 1 || idx[i] > size[i], length(size)))
@@ -108,17 +120,19 @@ function load_boundary_edge(pad::Pad, arr, dim, dir, neigh_dist)
         stop_idx = CartesianIndex(ntuple(i -> i == dim ? (firstindex(arr, i) + neigh_dist - 1) : lastindex(arr, i), ndims(arr)))
     end
     edge_size = ntuple(i -> length(start_idx[i]:stop_idx[i]), ndims(arr))
-    return Fill(pad.padval, edge_size)
+    # FIXME: return Fill(pad.padval, edge_size)
+    return move(thunk_processor(), fill(pad.padval, edge_size))
 end
 function load_boundary_corner(pad::Pad, arr, corner_side, neigh_dist)
     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)))
     corner_size = ntuple(i -> length(start_idx[i]:stop_idx[i]), ndims(arr))
-    return Fill(pad.padval, corner_size)
+    # FIXME: return Fill(pad.padval, corner_size)
+    return move(thunk_processor(), fill(pad.padval, corner_size))
 end
 
 """
-    @stencil idx in range begin body 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
@@ -205,21 +219,25 @@ macro stencil(orig_ex)
         @gensym chunk_idx
 
         # Generate function with transformed body
-        @gensym inner_index_var
+        @gensym inner_vars inner_index_var
         new_inner_ex_body = prewalk(inner_ex) do old_inner_ex
             if @capture(old_inner_ex, read_var_[read_idx_]) && read_idx == write_idx
                 # Direct access
-                return :($read_var[$inner_index_var])
+                if read_var == write_var
+                    return :($write_var[$inner_index_var])
+                else
+                    return :($inner_vars.$read_var[$inner_index_var])
+                end
             elseif @capture(old_inner_ex, @neighbors(read_var_[read_idx_], neigh_dist_, boundary_))
                 # Neighborhood access
-                return :($load_neighborhood($read_var, $inner_index_var, $neigh_dist))
+                return :($load_neighborhood($inner_vars.$read_var, $inner_index_var))
             end
             return old_inner_ex
         end
+        new_inner_f = :(($inner_index_var, $write_var, $inner_vars)->$new_inner_ex_body)
         new_inner_ex = quote
-            for $inner_index_var in CartesianIndices($write_var)
-                $new_inner_ex_body
-            end
+            $inner_vars = (;$(read_vars...))
+            $inner_stencil!($new_inner_f, $write_var, $inner_vars)
         end
         inner_fn = Expr(:->, Expr(:tuple, Expr(:parameters, write_var, read_vars...)), new_inner_ex)
 
@@ -254,7 +272,6 @@ macro stencil(orig_ex)
         end)
     end
 
-    @show final_ex
 
     return esc(final_ex)
 end

src/utils/haloarray.jl

@@ -1,8 +1,8 @@
 # Define the HaloArray type with minimized halo storage
-struct HaloArray{T,N,E,C,A,EA,CA} <: AbstractArray{T,N}
+struct HaloArray{T,N,E,C,A,EAT<:Tuple,CAT<:Tuple} <: AbstractArray{T,N}
     center::A
-    edges::NTuple{E, EA}
-    corners::NTuple{C, CA}
+    edges::EAT
+    corners::CAT
     halo_width::NTuple{N,Int}
 end
 
@@ -17,11 +17,11 @@ function HaloArray{T,N}(center_size::NTuple{N,Int}, halo_width::NTuple{N,Int}) w
     corners = ntuple(2^N) do i
         return Array{T,N}(undef, halo_width)
     end
-    return HaloArray{T,N,2N,2^N}(center, edges, corners, halo_width)
+    return HaloArray(center, edges, corners, halo_width)
 end
 
-HaloArray(center::AT, edges::NTuple{E, EA}, corners::NTuple{C, CA}, halo_width::NTuple{N, Int}) where {T,N,AT<:AbstractArray{T,N},C,E,CA,EA} =
-    HaloArray{T,N,E,C,AT,EA,CA}(center, edges, corners, halo_width)
+HaloArray(center::AT, edges::EAT, corners::CAT, halo_width::NTuple{N, Int}) where {T,N,AT<:AbstractArray{T,N},CAT<:Tuple,EAT<:Tuple} =
+    HaloArray{T,N,length(edges),length(corners),AT,EAT,CAT}(center, edges, corners, halo_width)
 
 Base.size(tile::HaloArray) = size(tile.center) .+ 2 .* tile.halo_width
 function Base.axes(tile::HaloArray{T,N,H}) where {T,N,H}
@@ -57,10 +57,10 @@ function Base.getindex(tile::HaloArray{T,N}, I::Vararg{Int,N}) where {T,N}
     else
         for d in 1:N
             if I[d] < 1
-                halo_idx = (I[1:d-1]..., I[d] + tile.halo_width[d], I[d+1:end]...)
+                halo_idx = ntuple(i->i == d ? I[i] + tile.halo_width[i] : I[i], N)
                 return tile.edges[(2*(d-1))+1][halo_idx...]
             elseif I[d] > size(tile.center, d)
-                halo_idx = (I[1:d-1]..., I[d] - size(tile.center, d), I[d+1:end]...)
+                halo_idx = ntuple(i->i == d ? I[i] - size(tile.center, d) : I[i], N)
                 return tile.edges[(2*(d-1))+2][halo_idx...]
             end
         end
@@ -84,32 +84,13 @@ function Base.setindex!(tile::HaloArray{T,N}, value, I::Vararg{Int,N}) where {T,
         # Edge
         for d in 1:N
             if I[d] < 1
-                halo_idx = (I[1:d-1]..., I[d] + tile.halo_width[d], I[d+1:end]...)
+                halo_idx = ntuple(i->i == d ? I[i] + tile.halo_width[i] : I[i], N)
                 return tile.edges[(2*(d-1))+1][halo_idx...] = value
             elseif I[d] > size(tile.center, d)
-                halo_idx = (I[1:d-1]..., I[d] - size(tile.center, d), I[d+1:end]...)
+                halo_idx = ntuple(i->i == d ? I[i] - size(tile.center, d) : I[i], N)
                 return tile.edges[(2*(d-1))+2][halo_idx...] = value
             end
         end
     end
     error("Index out of bounds")
 end
-
-#=
-# Example usage
-center_size = (3, 5)
-halo_width = (1, 1)
-tile = HaloArray{Float64, 2}(center_size, halo_width)
-
-# Set values in the center and halo
-tile[2, 2] = 1.0
-tile[0, 2] = 2.0  # This should be in an edge
-tile[0, 0] = 3.0  # This should be in a corner
-tile[4, 6] = 4.0  # This should be in a corner
-
-# Get values from the center and halo
-println(tile[2, 2])  # 1.0
-println(tile[0, 2])  # 2.0
-println(tile[0, 0])  # 3.0
-println(tile[4, 6])  # 4.0
-=#