Skip to content

Commit

Permalink
More generalization of tuples (#4049)
Browse files Browse the repository at this point in the history
Co-authored-by: Gregory L. Wagner <[email protected]>
Co-authored-by: Navid C. Constantinou <[email protected]>
  • Loading branch information
3 people authored Jan 20, 2025
1 parent 234bc98 commit 8b4c1ad
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 8 deletions.
17 changes: 10 additions & 7 deletions src/BoundaryConditions/fill_halo_regions_periodic.jl
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ using KernelAbstractions.Extras.LoopInfo: @unroll
@inline parent_size_and_offset(c, dim1, dim2, size, offset) = (parent(c), size, fix_halo_offsets.(offset, c.offsets[[dim1, dim2]]))
@inline parent_size_and_offset(c, dim1, dim2, ::Symbol, offset) = (parent(c), size(parent(c))[[dim1, dim2]], (0, 0))

@inline function parent_size_and_offset(c::NTuple, dim1, dim2, ::Symbol, offset)
@inline function parent_size_and_offset(c::Tuple, dim1, dim2, ::Symbol, offset)
p = parent.(c)
p_size = (minimum([size(t, dim1) for t in p]), minimum([size(t, dim2) for t in p]))
return p, p_size, (0, 0)
Expand Down Expand Up @@ -71,9 +71,10 @@ end
#### Tupled periodic boundary condition
####

@kernel function fill_periodic_west_and_east_halo!(c::NTuple{M}, ::Val{H}, N) where {M, H}
@kernel function fill_periodic_west_and_east_halo!(c::Tuple, ::Val{H}, N) where {H}
j, k = @index(Global, NTuple)
@unroll for n = 1:M
ntuple(Val(length(c))) do n
Base.@_inline_meta
@unroll for i = 1:H
@inbounds begin
c[n][i, j, k] = c[n][N+i, j, k] # west
Expand All @@ -83,9 +84,10 @@ end
end
end

@kernel function fill_periodic_south_and_north_halo!(c::NTuple{M}, ::Val{H}, N) where {M, H}
@kernel function fill_periodic_south_and_north_halo!(c::Tuple, ::Val{H}, N) where {H}
i, k = @index(Global, NTuple)
@unroll for n = 1:M
ntuple(Val(length(c))) do n
Base.@_inline_meta
@unroll for j = 1:H
@inbounds begin
c[n][i, j, k] = c[n][i, N+j, k] # south
Expand All @@ -95,9 +97,10 @@ end
end
end

@kernel function fill_periodic_bottom_and_top_halo!(c::NTuple{M}, ::Val{H}, N) where {M, H}
@kernel function fill_periodic_bottom_and_top_halo!(c::Tuple, ::Val{H}, N) where {H}
i, j = @index(Global, NTuple)
@unroll for n = 1:M
ntuple(Val(length(c))) do n
Base.@_inline_meta
@unroll for k = 1:H
@inbounds begin
c[n][i, j, k] = c[n][i, j, N+k] # top
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -137,4 +137,4 @@ where `c = C[tracer_index]`.
- immersed_∇_dot_qᶜ(i, j, k, grid, c, c_immersed_bc, closure, diffusivities, val_tracer_index, clock, model_fields)
+ biogeochemical_transition(i, j, k, grid, biogeochemistry, val_tracer_name, clock, model_fields)
+ forcing(i, j, k, grid, clock, model_fields))
end
end

0 comments on commit 8b4c1ad

Please sign in to comment.