Add masked operators

.buildkite/run_tests.yml

@@ -46,7 +46,7 @@ steps:
  agents:
  queue: "juliagpu"
  rocm: "*"
- rocmgpu: "gfx1101" # select Ludovic's Navi 3 card (ROCm 6.0)
+ rocmgpu: "*" #"gfx1101" # select Ludovic's Navi 3 card (ROCm 6.0)
  timeout_in_minutes: 120
  soft_fail:
  - exit_status: 3

Project.toml

@@ -1,7 +1,7 @@
 name = "Chmy"
 uuid = "33a72cf0-4690-46d7-b987-06506c2248b9"
 authors = ["Ivan Utkin <[email protected]>, Ludovic Raess <[email protected]>, and contributors"]
-version = "0.1.10"
+version = "0.1.12"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/BoundaryConditions/batch.jl

@@ -172,7 +172,7 @@ bc!(arch::Architecture, grid::SG, f_bc::Vararg{FieldAndBC}; kwargs...) = bc!(arc
 end
 
 function bc!(side::Side, dim::Dim, arch::Architecture, grid::SG, batch::FieldBatch)
- worksize = remove_dim(dim, size(grid, Center()) .+ 2)
+ worksize = remove_dim(dim, size(grid, Vertex()) .+ 2)
  bc_kernel!(Architectures.get_backend(arch), 256, worksize)(side, dim, grid, batch.fields, batch.conditions)
  return
 end

src/Chmy.jl

@@ -6,6 +6,8 @@ using KernelAbstractions
 export Dim, Side, Left, Right
 export remove_dim, insert_dim
 
+export Offset
+
 include("macros.jl")
 include("utils.jl")
 

src/Fields/field.jl

@@ -9,7 +9,7 @@ struct Field{T,N,L,H,A} <: AbstractField{T,N,L}
  Field{L,H}(data::AbstractArray{T,N}, dims::NTuple{N,Integer}) where {L,H,T,N} = new{T,N,L,H,typeof(data)}(data, dims)
 end
 
-Base.@assume_effects :foldable halo(::Field{T,N,A,H}) where {T,N,A,H} = H
+halo(::Field{T,N,A,H}) where {T,N,A,H} = H
 
 # AbstractArray interface
 Base.size(f::Field) = f.dims
@@ -58,10 +58,21 @@ Field(arch::Architecture, args...; kwargs...) = Field(Architectures.get_backend(
 
 # set fields
 
-set!(f::Field, other::Field) = (copyto!(interior(f), interior(other)); nothing)
 set!(f::Field, val::Number) = (fill!(interior(f), val); nothing)
 set!(f::Field, A::AbstractArray) = (copyto!(interior(f), A); nothing)
 
+function set!(f::Field, other::AbstractField)
+ dst = interior(f)
+ src = interior(other)
+ backend = KernelAbstractions.get_backend(dst)
+ _set_field!(backend, 256, size(dst))(dst, src)
+end
+
+@kernel inbounds = true function _set_field!(dst, src)
+ I = @index(Global, NTuple)
+ dst[I...] = src[I...]
+end
+
 @kernel inbounds = true function _set_continuous!(dst, grid, loc, fun::F, args...) where {F}
  I = @index(Global, NTuple)
  dst[I...] = fun(coord(grid, loc, I...)..., args...)

src/GridOperators/GridOperators.jl

@@ -7,12 +7,17 @@ export itp, lerp, hlerp
 
 export divg
 
+export AbstractMask, FieldMask, FieldMask1D, FieldMask2D, FieldMask3D, at
+
 using Chmy
 using Chmy.Grids
 using Chmy.Fields
+using Chmy.Architectures
 
 import Chmy.@add_cartesian
 
+using Adapt
+
 import Base: @propagate_inbounds, front
 
 p(::Dim{D}, I::Vararg{Integer,N}) where {D,N} = ntuple(i -> i == D ? I[i] + oneunit(I[i]) : I[i], Val(N))
@@ -27,11 +32,19 @@ m(::Dim{D}, I::Vararg{Integer,N}) where {D,N} = ntuple(i -> i == D ? I[i] - oneu
 @add_cartesian il(loc::L, from::L, dim, I::Vararg{Integer,N}) where {N,L<:Location} = m(dim, I...)
 @add_cartesian ir(loc::L, from::L, dim, I::Vararg{Integer,N}) where {N,L<:Location} = p(dim, I...)
 
-@add_cartesian left(f, loc, from, dim, I::Vararg{Integer,N}) where {N} = f[il(loc, from, dim, I...)...]
-@add_cartesian right(f, loc, from, dim, I::Vararg{Integer,N}) where {N} = f[ir(loc, from, dim, I...)...]
+@add_cartesian left(f, loc::NTuple{N,Location}, from::NTuple{N,Location}, dim::Dim{D}, I::Vararg{Integer,N}) where {N,D} = f[il(loc[D], from[D], dim, I...)...]
+@add_cartesian right(f, loc::NTuple{N,Location}, from::NTuple{N,Location}, dim::Dim{D}, I::Vararg{Integer,N}) where {N,D} = f[ir(loc[D], from[D], dim, I...)...]
+
+flipped(loc::NTuple{N,Location}, ::Dim{D}) where {N,D} = ntuple(dim -> dim == D ? flip(loc[dim]) : loc[dim], Val(N))
 
 include("partial_derivatives.jl")
 include("interpolation.jl")
 include("field_operators.jl")
+include("cartesian_field_operators.jl")
+
+include("masked_operators.jl")
+include("field_mask.jl")
+include("masked_field_operators.jl")
+include("masked_cartesian_field_operators.jl")
 
 end

src/GridOperators/cartesian_field_operators.jl

@@ -0,0 +1,46 @@
+for (dim, coord) in enumerate((:x, :y, :z))
+ _l = Symbol(:left, coord)
+ _r = Symbol(:right, coord)
+ _δ = Symbol(:δ, coord)
+ _∂ = Symbol(:∂, coord)
+
+ @eval begin
+ export $_δ, $_∂, $_l, $_r
+
+ """
+ $($_l)(f, I)
+
+ "left side" of a field (`[1:end-1]`) in $($(string(coord))) direction.
+ """
+ @add_cartesian function $_l(f::AbstractField, I::Vararg{Integer,N}) where {N}
+ left(f, Dim($dim), I...)
+ end
+
+ """
+ $($_r)(f, I)
+
+ "right side" of a field (`[2:end]`) in $($(string(coord))) direction.
+ """
+ @add_cartesian function $_r(f::AbstractField, I::Vararg{Integer,N}) where {N}
+ right(f, Dim($dim), I...)
+ end
+
+ """
+ $($_δ)(f, I)
+
+ Finite difference in $($(string(coord))) direction.
+ """
+ @add_cartesian function $_δ(f::AbstractField, I::Vararg{Integer,N}) where {N}
+ δ(f, Dim($dim), I...)
+ end
+
+ """
+ $($_∂)(f, grid, I)
+
+ Directional partial derivative in $($(string(coord))) direction.
+ """
+ @add_cartesian function $_∂(f::AbstractField, grid, I::Vararg{Integer,N}) where {N}
+ ∂(f, grid, Dim($dim), I...)
+ end
+ end
+end

src/GridOperators/field_mask.jl

@@ -0,0 +1,113 @@
+struct FieldMask1D{T,CF,VF} <: AbstractMask{T,1}
+ c::CF
+ v::VF
+ function FieldMask1D{T}(c::AbstractField{T,1}, v::AbstractField{T,1}) where {T}
+ return new{T,typeof(c),typeof(v)}(c, v)
+ end
+end
+
+function FieldMask1D(arch::Architecture, grid::StructuredGrid{1}, type=eltype(grid); kwargs...)
+ c = Field(arch, grid, Center(), type; kwargs...)
+ v = Field(arch, grid, Vertex(), type; kwargs...)
+ return FieldMask1D{type}(c, v)
+end
+
+@propagate_inbounds at(ω::FieldMask1D, ::Tuple{Center}, ix) = ω.c[ix]
+@propagate_inbounds at(ω::FieldMask1D, ::Tuple{Vertex}, ix) = ω.v[ix]
+
+struct FieldMask2D{T,CCF,VVF,VCF,CVF} <: AbstractMask{T,2}
+ cc::CCF
+ vv::VVF
+ vc::VCF
+ cv::CVF
+ function FieldMask2D{T}(cc::AbstractField{T,2}, vv::AbstractField{T,2}, vc::AbstractField{T,2}, cv::AbstractField{T,2}) where {T}
+ return new{T,typeof(cc),typeof(vv),typeof(vc),typeof(cv)}(cc, vv, vc, cv)
+ end
+end
+
+function FieldMask2D(arch::Architecture, grid::StructuredGrid{2}, type=eltype(grid); kwargs...)
+ cc = Field(arch, grid, Center(), type; kwargs...)
+ vv = Field(arch, grid, Vertex(), type; kwargs...)
+ vc = Field(arch, grid, (Vertex(), Center()), type; kwargs...)
+ cv = Field(arch, grid, (Center(), Vertex()), type; kwargs...)
+ return FieldMask2D{type}(cc, vv, vc, cv)
+end
+
+@propagate_inbounds at(ω::FieldMask2D, ::Tuple{Center,Center}, ix, iy) = ω.cc[ix, iy]
+@propagate_inbounds at(ω::FieldMask2D, ::Tuple{Vertex,Vertex}, ix, iy) = ω.vv[ix, iy]
+@propagate_inbounds at(ω::FieldMask2D, ::Tuple{Vertex,Center}, ix, iy) = ω.vc[ix, iy]
+@propagate_inbounds at(ω::FieldMask2D, ::Tuple{Center,Vertex}, ix, iy) = ω.cv[ix, iy]
+
+struct FieldMask3D{T,CCCF,VVVF,VCCF,CVCF,CCVF,VVCF,VCVF,CVVF} <: AbstractMask{T,3}
+ ccc::CCCF
+ vvv::VVVF
+ vcc::VCCF
+ cvc::CVCF
+ ccv::CCVF
+ vvc::VVCF
+ vcv::VCVF
+ cvv::CVVF
+ function FieldMask3D{T}(ccc::AbstractField{T,3},
+ vvv::AbstractField{T,3},
+ vcc::AbstractField{T,3},
+ cvc::AbstractField{T,3},
+ ccv::AbstractField{T,3},
+ vvc::AbstractField{T,3},
+ vcv::AbstractField{T,3},
+ cvv::AbstractField{T,3}) where {T}
+ return new{T,
+ typeof(ccc),
+ typeof(vvv),
+ typeof(vcc),
+ typeof(cvc),
+ typeof(ccv),
+ typeof(vvc),
+ typeof(vcv),
+ typeof(cvv)}(ccc, vvv, vcc, cvc, ccv, vvc, vcv, cvv)
+ end
+end
+
+function FieldMask3D(arch::Architecture, grid::StructuredGrid{3}, type=eltype(grid); kwargs...)
+ ccc = Field(arch, grid, Center(), type; kwargs...)
+ vvv = Field(arch, grid, Vertex(), type; kwargs...)
+ vcc = Field(arch, grid, (Vertex(), Center(), Center()), type; kwargs...)
+ cvc = Field(arch, grid, (Center(), Vertex(), Center()), type; kwargs...)
+ ccv = Field(arch, grid, (Center(), Center(), Vertex()), type; kwargs...)
+ vvc = Field(arch, grid, (Vertex(), Vertex(), Center()), type; kwargs...)
+ vcv = Field(arch, grid, (Vertex(), Center(), Vertex()), type; kwargs...)
+ cvv = Field(arch, grid, (Center(), Vertex(), Vertex()), type; kwargs...)
+ return FieldMask3D{type}(ccc, vvv, vcc, cvc, ccv, vvc, vcv, cvv)
+end
+
+@propagate_inbounds at(ω::FieldMask3D, ::Tuple{Center,Center,Center}, ix, iy, iz) = ω.ccc[ix, iy, iz]
+@propagate_inbounds at(ω::FieldMask3D, ::Tuple{Vertex,Vertex,Vertex}, ix, iy, iz) = ω.vvv[ix, iy, iz]
+@propagate_inbounds at(ω::FieldMask3D, ::Tuple{Vertex,Center,Center}, ix, iy, iz) = ω.vcc[ix, iy, iz]
+@propagate_inbounds at(ω::FieldMask3D, ::Tuple{Center,Vertex,Center}, ix, iy, iz) = ω.cvc[ix, iy, iz]
+@propagate_inbounds at(ω::FieldMask3D, ::Tuple{Center,Center,Vertex}, ix, iy, iz) = ω.ccv[ix, iy, iz]
+@propagate_inbounds at(ω::FieldMask3D, ::Tuple{Vertex,Vertex,Center}, ix, iy, iz) = ω.vvc[ix, iy, iz]
+@propagate_inbounds at(ω::FieldMask3D, ::Tuple{Vertex,Center,Vertex}, ix, iy, iz) = ω.vcv[ix, iy, iz]
+@propagate_inbounds at(ω::FieldMask3D, ::Tuple{Center,Vertex,Vertex}, ix, iy, iz) = ω.cvv[ix, iy, iz]
+
+FieldMask(arch::Architecture, grid::StructuredGrid{1}, args...; kwargs...) = FieldMask1D(arch, grid, args...; kwargs...)
+FieldMask(arch::Architecture, grid::StructuredGrid{2}, args...; kwargs...) = FieldMask2D(arch, grid, args...; kwargs...)
+FieldMask(arch::Architecture, grid::StructuredGrid{3}, args...; kwargs...) = FieldMask3D(arch, grid, args...; kwargs...)
+
+# Adapt rules
+
+Adapt.adapt_structure(to, f::FieldMask1D{T}) where {T} = FieldMask1D{T}(Adapt.adapt(to, f.c), Adapt.adapt(to, f.v))
+
+function Adapt.adapt_structure(to, f::FieldMask2D{T}) where {T}
+ FieldMask2D{T}(Adapt.adapt(to, f.cc), Adapt.adapt(to, f.vv),
+ Adapt.adapt(to, f.vc), Adapt.adapt(to, f.cv))
+end
+
+function Adapt.adapt_structure(to, f::FieldMask3D{T}) where {T}
+ FieldMask3D{T}(Adapt.adapt(to, f.ccc),
+ Adapt.adapt(to, f.vvv),
+ Adapt.adapt(to, f.vcc),
+ Adapt.adapt(to, f.cvc),
+ Adapt.adapt(to, f.ccv),
+ Adapt.adapt(to, f.vvc),
+ Adapt.adapt(to, f.vcv),
+ Adapt.adapt(to, f.cvv))
+end

src/GridOperators/field_operators.jl

@@ -1,58 +1,26 @@
 # staggered grid operators
-@add_cartesian left(f::AbstractField, dim, I::Vararg{Integer,N}) where {N} = left(f, location(f, dim), flip(location(f, dim)), dim, I...)
-
-@add_cartesian right(f::AbstractField, dim, I::Vararg{Integer,N}) where {N} = right(f, location(f, dim), flip(location(f, dim)), dim, I...)
-
-@add_cartesian δ(f::AbstractField, dim, I::Vararg{Integer,N}) where {N} = δ(f, location(f, dim), flip(location(f, dim)), dim, I...)
-
-@add_cartesian ∂(f::AbstractField, grid, dim, I::Vararg{Integer,N}) where {N} = ∂(f, location(f, dim), flip(location(f, dim)), grid, dim, I...)
-
-# staggered operators on Cartesian grids
-for (dim, coord) in enumerate((:x, :y, :z))
- _l = Symbol(:left, coord)
- _r = Symbol(:right, coord)
- _δ = Symbol(:δ, coord)
- _∂ = Symbol(:∂, coord)
-
- @eval begin
- export $_δ, $_∂, $_l, $_r
-
- """
- $($_l)(f, I)
-
- "left side" of a field (`[1:end-1]`) in $($(string(coord))) direction.
- """
- @add_cartesian function $_l(f::AbstractField, I::Vararg{Integer,N}) where {N}
- left(f, Dim($dim), I...)
- end
-
- """
- $($_r)(f, I)
-
- "right side" of a field (`[2:end]`) in $($(string(coord))) direction.
- """
- @add_cartesian function $_r(f::AbstractField, I::Vararg{Integer,N}) where {N}
- right(f, Dim($dim), I...)
- end
-
- """
- $($_δ)(f, I)
+@add_cartesian function left(f::AbstractField, dim, I::Vararg{Integer,N}) where {N}
+ loc = location(f)
+ from = flipped(loc, dim)
+ left(f, loc, from, dim, I...)
+end
 
- Finite difference in $($(string(coord))) direction.
-  """
-  @add_cartesian function $_δ(f::AbstractField, I::Vararg{Integer,N}) where {N}
-  δ(f, Dim($dim), I...)
- end
+@add_cartesian function right(f::AbstractField, dim, I::Vararg{Integer,N}) where {N}
+ loc = location(f)
+ from = flipped(loc, dim)
+ right(f, loc, from, dim, I...)
+end
 
- """
- $($_∂)(f, grid, I)
+@add_cartesian function δ(f::AbstractField, dim, I::Vararg{Integer,N}) where {N}
+ loc = location(f)
+ from = flipped(loc, dim)
+ return δ(f, loc, from, dim, I...)
+end
 
- Directional partial derivative in $($(string(coord))) direction.
- """
- @add_cartesian function $_∂(f::AbstractField, grid, I::Vararg{Integer,N}) where {N}
- ∂(f, grid, Dim($dim), I...)
- end
- end
+@add_cartesian function ∂(f::AbstractField, grid, dim, I::Vararg{Integer,N}) where {N}
+ loc = location(f)
+ from = flipped(loc, dim)
+ ∂(f, loc, from, grid, dim, I...)
 end
 
 # covariant derivatives

src/GridOperators/masked_cartesian_field_operators.jl

@@ -0,0 +1,46 @@
+for (dim, coord) in enumerate((:x, :y, :z))
+ _l = Symbol(:left, coord)
+ _r = Symbol(:right, coord)
+ _δ = Symbol(:δ, coord)
+ _∂ = Symbol(:∂, coord)
+
+ @eval begin
+ export $_δ, $_∂, $_l, $_r
+
+ """
+ $($_l)(f, ω, I)
+
+ "left side" of a field (`[1:end-1]`) in $($(string(coord))) direction, masked with `ω`.
+ """
+ @add_cartesian function $_l(f::AbstractField, ω::AbstractMask, I::Vararg{Integer,N}) where {N}
+ left(f, ω, Dim($dim), I...)
+ end
+
+ """
+ $($_r)(f, ω, I)
+
+ "right side" of a field (`[2:end]`) in $($(string(coord))) direction, masked with `ω`.
+ """
+ @add_cartesian function $_r(f::AbstractField, ω::AbstractMask, I::Vararg{Integer,N}) where {N}
+ right(f, ω, Dim($dim), I...)
+ end
+
+ """
+ $($_δ)(f, ω, I)
+
+ Finite difference in $($(string(coord))) direction, masked with `ω`.
+ """
+ @add_cartesian function $_δ(f::AbstractField, ω::AbstractMask, I::Vararg{Integer,N}) where {N}
+ δ(f, ω, Dim($dim), I...)
+ end
+
+ """
+ $($_∂)(f, ω, grid, I)
+
+ Directional partial derivative in $($(string(coord))) direction, masked with `ω`.
+ """
+ @add_cartesian function $_∂(f::AbstractField, ω::AbstractMask, grid, I::Vararg{Integer,N}) where {N}
+ ∂(f, ω, grid, Dim($dim), I...)
+ end
+ end
+end