diff --git a/.github/workflows/ClimaCoreMakie.yml b/.github/workflows/ClimaCoreMakie.yml index e43a210935..1712a00d5a 100644 --- a/.github/workflows/ClimaCoreMakie.yml +++ b/.github/workflows/ClimaCoreMakie.yml @@ -17,7 +17,7 @@ jobs: timeout-minutes: 45 steps: - name: Checkout - uses: actions/checkout@v6 + uses: actions/checkout@v7 - uses: julia-actions/setup-julia@v3 with: version: '1.10' diff --git a/.github/workflows/ClimaCorePlots.yml b/.github/workflows/ClimaCorePlots.yml index c20ab7618e..510b218d73 100644 --- a/.github/workflows/ClimaCorePlots.yml +++ b/.github/workflows/ClimaCorePlots.yml @@ -17,7 +17,7 @@ jobs: timeout-minutes: 30 steps: - name: Checkout - uses: actions/checkout@v6 + uses: actions/checkout@v7 - uses: julia-actions/setup-julia@v3 with: version: '1.10' diff --git a/.github/workflows/ClimaCoreSpectra.yml b/.github/workflows/ClimaCoreSpectra.yml index 90e3cdf422..35a803d449 100644 --- a/.github/workflows/ClimaCoreSpectra.yml +++ b/.github/workflows/ClimaCoreSpectra.yml @@ -17,7 +17,7 @@ jobs: timeout-minutes: 30 steps: - name: Checkout - uses: actions/checkout@v6 + uses: actions/checkout@v7 - uses: julia-actions/setup-julia@v3 with: version: '1.10' diff --git a/.github/workflows/ClimaCoreTempestRemap.yml b/.github/workflows/ClimaCoreTempestRemap.yml index 6e0d45a5de..fbc6791726 100644 --- a/.github/workflows/ClimaCoreTempestRemap.yml +++ b/.github/workflows/ClimaCoreTempestRemap.yml @@ -18,7 +18,7 @@ jobs: timeout-minutes: 30 steps: - name: Checkout - uses: actions/checkout@v6 + uses: actions/checkout@v7 - uses: julia-actions/setup-julia@v3 with: version: '1.10' diff --git a/.github/workflows/ClimaCoreVTK.yml b/.github/workflows/ClimaCoreVTK.yml index 7fe64d0c56..7b058cd7ee 100644 --- a/.github/workflows/ClimaCoreVTK.yml +++ b/.github/workflows/ClimaCoreVTK.yml @@ -17,7 +17,7 @@ jobs: timeout-minutes: 30 steps: - name: Checkout - uses: actions/checkout@v6 + uses: actions/checkout@v7 - uses: julia-actions/setup-julia@v3 with: version: '1.10' diff --git a/.github/workflows/DocCleanup.yml b/.github/workflows/DocCleanup.yml index 4fe76e4df1..1a04614094 100644 --- a/.github/workflows/DocCleanup.yml +++ b/.github/workflows/DocCleanup.yml @@ -9,7 +9,7 @@ jobs: runs-on: ubuntu-24.04 steps: - name: Checkout gh-pages branch - uses: actions/checkout@v6 + uses: actions/checkout@v7 with: ref: gh-pages diff --git a/.github/workflows/Downstream.yml b/.github/workflows/Downstream.yml index 2b414e005a..8211ce57df 100644 --- a/.github/workflows/Downstream.yml +++ b/.github/workflows/Downstream.yml @@ -35,13 +35,13 @@ jobs: - '1.10' - '1.11' steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - uses: julia-actions/setup-julia@v3 with: version: ${{ matrix.version }} - uses: julia-actions/cache@v3 - uses: julia-actions/julia-buildpkg@v1 - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 with: repository: 'CliMA/${{ matrix.package }}' path: ${{ matrix.package }} diff --git a/.github/workflows/Invalidations.yml b/.github/workflows/Invalidations.yml index e370a2fa5a..14d8aca5e0 100644 --- a/.github/workflows/Invalidations.yml +++ b/.github/workflows/Invalidations.yml @@ -19,12 +19,12 @@ jobs: - uses: julia-actions/setup-julia@v3 with: version: '1.10' - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - uses: julia-actions/julia-buildpkg@v1 - uses: julia-actions/julia-invalidations@v1 id: invs_pr - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 with: ref: ${{ github.event.repository.default_branch }} - uses: julia-actions/julia-buildpkg@v1 diff --git a/.github/workflows/UnitTests.yml b/.github/workflows/UnitTests.yml index fff9d8e721..af1138f95d 100644 --- a/.github/workflows/UnitTests.yml +++ b/.github/workflows/UnitTests.yml @@ -40,7 +40,7 @@ jobs: steps: - name: Checkout - uses: actions/checkout@v6 + uses: actions/checkout@v7 - uses: julia-actions/setup-julia@v3 with: version: ${{ matrix.version }} diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 1351485bb9..264ac45a9b 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -11,7 +11,7 @@ jobs: docbuild: runs-on: ubuntu-24.04 steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - uses: julia-actions/setup-julia@v3 with: version: '1.10' diff --git a/.github/workflows/downgrade.yml b/.github/workflows/downgrade.yml index 3bf677902e..0eb251f7ae 100644 --- a/.github/workflows/downgrade.yml +++ b/.github/workflows/downgrade.yml @@ -22,7 +22,7 @@ jobs: matrix: version: ['1.10', '1.11'] steps: - - uses: actions/checkout@v6 + - uses: actions/checkout@v7 - uses: julia-actions/setup-julia@latest with: version: ${{ matrix.version }} diff --git a/.github/workflows/update_dev_guides.yml b/.github/workflows/update_dev_guides.yml index da0e7d8c57..b697ca9b89 100644 --- a/.github/workflows/update_dev_guides.yml +++ b/.github/workflows/update_dev_guides.yml @@ -20,7 +20,7 @@ jobs: update-subtree: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 with: fetch-depth: 0 # Full history is required for subtree pull diff --git a/ext/cuda/column_matrix_helpers.jl b/ext/cuda/column_matrix_helpers.jl index f832769a08..26223d5413 100644 --- a/ext/cuda/column_matrix_helpers.jl +++ b/ext/cuda/column_matrix_helpers.jl @@ -11,7 +11,7 @@ Base.@propagate_inbounds function row_mul_mat!( ) where {P} prod_eltype = P v = threadIdx().x - i = threadIdx().y + block_col_idx = threadIdx().y mat1_eltype = typeof(mat1_row) mat2_eltype = eltype(matrix2) ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) @@ -28,7 +28,7 @@ Base.@propagate_inbounds function row_mul_mat!( if ld2 <= pd - mat1_row_d <= ud2 && (0i32 < v + mat1_row_d + half <= CUDA.blockDim().x) @inbounds mat1_row[mat1_row_d] * - matrix2[v + mat1_row_d + half + (i - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] + matrix2[v + mat1_row_d + half + (block_col_idx - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] else zero_entry end @@ -47,7 +47,7 @@ Base.@propagate_inbounds function row_mul_mat!( ) where {P} prod_eltype = P v = threadIdx().x - i = threadIdx().y + block_col_idx = threadIdx().y mat1_eltype = typeof(mat1_row) mat2_eltype = eltype(matrix2) ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) @@ -64,7 +64,7 @@ Base.@propagate_inbounds function row_mul_mat!( if ld2 <= pd - mat1_row_d <= ud2 && (0i32 < v + mat1_row_d - half < CUDA.blockDim().x) @inbounds mat1_row[mat1_row_d] * - matrix2[v + mat1_row_d - half + (i - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] + matrix2[v + mat1_row_d - half + (block_col_idx - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] else zero_entry end @@ -83,7 +83,7 @@ Base.@propagate_inbounds function row_mul_mat!( ) where {P} prod_eltype = P v = threadIdx().x - i = threadIdx().y + block_col_idx = threadIdx().y mat1_eltype = typeof(mat1_row) mat2_eltype = eltype(matrix2) ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) @@ -100,7 +100,7 @@ Base.@propagate_inbounds function row_mul_mat!( if ld2 <= pd - mat1_row_d <= ud2 && (0i32 < v + mat1_row_d <= CUDA.blockDim().x - 1i32) @inbounds mat1_row[mat1_row_d] * - matrix2[v + mat1_row_d + (i - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] + matrix2[v + mat1_row_d + (block_col_idx - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] else zero_entry end @@ -119,7 +119,7 @@ Base.@propagate_inbounds function row_mul_mat!( ) where {P} prod_eltype = P v = threadIdx().x - i = threadIdx().y + block_col_idx = threadIdx().y mat1_eltype = typeof(mat1_row) mat2_eltype = eltype(matrix2) ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) @@ -137,7 +137,7 @@ Base.@propagate_inbounds function row_mul_mat!( UnrolledUtilities.unrolled_mapreduce(+, (ld1:ud1...,)) do mat1_row_d if ld2 <= pd - mat1_row_d <= ud2 && (0i32 < v + mat1_row_d <= CUDA.blockDim().x) @inbounds mat1_row[mat1_row_d] * - matrix2[v + mat1_row_d + (i - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] + matrix2[v + mat1_row_d + (block_col_idx - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] else zero_entry end @@ -156,7 +156,7 @@ Base.@propagate_inbounds function row_mul_mat!( ) where {P} prod_eltype = P v = threadIdx().x - i = threadIdx().y + block_col_idx = threadIdx().y mat1_eltype = typeof(mat1_row) mat2_eltype = eltype(matrix2) ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) @@ -173,7 +173,7 @@ Base.@propagate_inbounds function row_mul_mat!( if ld2 <= pd - mat1_row_d <= ud2 && (0i32 < v + mat1_row_d + half <= CUDA.blockDim().x) @inbounds mat1_row[mat1_row_d] * - matrix2[v + mat1_row_d + half + (i - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] + matrix2[v + mat1_row_d + half + (block_col_idx - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] else zero_entry end @@ -192,7 +192,7 @@ Base.@propagate_inbounds function row_mul_mat!( ) where {P} prod_eltype = P v = threadIdx().x - i = threadIdx().y + block_col_idx = threadIdx().y mat1_eltype = typeof(mat1_row) mat2_eltype = eltype(matrix2) ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) @@ -209,7 +209,7 @@ Base.@propagate_inbounds function row_mul_mat!( if ld2 <= pd - mat1_row_d <= ud2 && (0i32 < v + mat1_row_d - half < CUDA.blockDim().x) @inbounds mat1_row[mat1_row_d] * - matrix2[v + mat1_row_d - half + (i - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] + matrix2[v + mat1_row_d - half + (block_col_idx - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] else zero_entry end @@ -228,7 +228,7 @@ Base.@propagate_inbounds function row_mul_mat!( ) where {P} prod_eltype = P v = threadIdx().x - i = threadIdx().y + block_col_idx = threadIdx().y mat1_eltype = typeof(mat1_row) mat2_eltype = eltype(matrix2) ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) @@ -244,7 +244,7 @@ Base.@propagate_inbounds function row_mul_mat!( UnrolledUtilities.unrolled_mapreduce(+, (ld1:ud1...,)) do mat1_row_d if ld2 <= pd - mat1_row_d <= ud2 && (0i32 < v + mat1_row_d <= CUDA.blockDim().x) @inbounds mat1_row[mat1_row_d] * - matrix2[v + mat1_row_d + (i - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] + matrix2[v + mat1_row_d + (block_col_idx - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] else zero_entry end @@ -263,7 +263,7 @@ Base.@propagate_inbounds function row_mul_mat!( ) where {P} prod_eltype = P v = threadIdx().x - i = threadIdx().y + block_col_idx = threadIdx().y mat1_eltype = typeof(mat1_row) mat2_eltype = eltype(matrix2) ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) @@ -279,7 +279,7 @@ Base.@propagate_inbounds function row_mul_mat!( UnrolledUtilities.unrolled_mapreduce(+, (ld1:ud1...,)) do mat1_row_d if ld2 <= pd - mat1_row_d <= ud2 && (0i32 < v + mat1_row_d < CUDA.blockDim().x) @inbounds mat1_row[mat1_row_d] * - matrix2[v + mat1_row_d + (i - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] + matrix2[v + mat1_row_d + (block_col_idx - 1i32) * CUDA.blockDim().x][pd - mat1_row_d] else zero_entry end @@ -301,7 +301,7 @@ Base.@propagate_inbounds function row_mul_vec!( ) where {P} prod_eltype = P v = threadIdx().x - i = threadIdx().y + block_col_idx = threadIdx().y mat1_eltype = typeof(mat1_row) mat2_eltype = eltype(matrix2) ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) @@ -316,7 +316,7 @@ Base.@propagate_inbounds function row_mul_vec!( if (0i32 < v + mat1_row_d + half <= CUDA.blockDim().x) @inbounds outer_or_mul( mat1_row[mat1_row_d], - matrix2[v + mat1_row_d + half + (i - 1i32) * CUDA.blockDim().x], + matrix2[v + mat1_row_d + half + (block_col_idx - 1i32) * CUDA.blockDim().x], ) else zero_entry @@ -332,7 +332,7 @@ Base.@propagate_inbounds function row_mul_vec!( ) where {P} prod_eltype = P v = threadIdx().x - i = threadIdx().y + block_col_idx = threadIdx().y mat1_eltype = typeof(mat1_row) mat2_eltype = eltype(matrix2) ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) @@ -347,7 +347,7 @@ Base.@propagate_inbounds function row_mul_vec!( if (0i32 < v + mat1_row_d - half < CUDA.blockDim().x) @inbounds outer_or_mul( mat1_row[mat1_row_d], - matrix2[v + mat1_row_d - half + (i - 1i32) * CUDA.blockDim().x], + matrix2[v + mat1_row_d - half + (block_col_idx - 1i32) * CUDA.blockDim().x], ) else zero_entry @@ -363,7 +363,7 @@ Base.@propagate_inbounds function row_mul_vec!( ) where {P} prod_eltype = P v = threadIdx().x - i = threadIdx().y + block_col_idx = threadIdx().y mat1_eltype = typeof(mat1_row) mat2_eltype = eltype(matrix2) ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) @@ -378,7 +378,7 @@ Base.@propagate_inbounds function row_mul_vec!( if (0i32 < v + mat1_row_d <= CUDA.blockDim().x - 1i32) @inbounds outer_or_mul( mat1_row[mat1_row_d], - matrix2[v + mat1_row_d + (i - 1i32) * CUDA.blockDim().x], + matrix2[v + mat1_row_d + (block_col_idx - 1i32) * CUDA.blockDim().x], ) else zero_entry @@ -394,7 +394,7 @@ Base.@propagate_inbounds function row_mul_vec!( ) where {P} prod_eltype = P v = threadIdx().x - i = threadIdx().y + block_col_idx = threadIdx().y mat1_eltype = typeof(mat1_row) mat2_eltype = eltype(matrix2) ld1, ud1 = MatrixFields.outer_diagonals(mat1_eltype) @@ -409,7 +409,7 @@ Base.@propagate_inbounds function row_mul_vec!( if (0i32 < v + mat1_row_d <= CUDA.blockDim().x) @inbounds outer_or_mul( mat1_row[mat1_row_d], - matrix2[v + mat1_row_d + (i - 1i32) * CUDA.blockDim().x], + matrix2[v + mat1_row_d + (block_col_idx - 1i32) * CUDA.blockDim().x], ) else zero_entry diff --git a/ext/cuda/operators_fd_eager.jl b/ext/cuda/operators_fd_eager.jl index 0e65c4792a..431cc90752 100644 --- a/ext/cuda/operators_fd_eager.jl +++ b/ext/cuda/operators_fd_eager.jl @@ -152,14 +152,25 @@ expression at the given index, and then copies the result into `out`. Base.@propagate_inbounds function eager_copyto_stencil_kernel!( out, bc::BC, + cart_inds, + mask, space, ) where {BC} - i = threadIdx().y - j = blockIdx().y v = threadIdx().x - h = blockIdx().z + col_idx = threadIdx().y + (blockIdx().x - 1) * blockDim().y + (i, j, h) = if mask isa NoMask + col_idx > length(cart_inds) && return nothing + @inbounds cart_inds[col_idx].I + else + (; i_map, j_map, h_map) = mask + @inbounds col_idx > length(i_map) && return nothing + @inbounds i = i_map[col_idx] + @inbounds j = j_map[col_idx] + @inbounds h = h_map[col_idx] + (i, j, h) + end hidx = (i, j, h) - val = @inbounds @inline calc_level_val(bc, space) + val = @inbounds @inline calc_level_val(bc, hidx, space) if space.staggering isa ClimaCore.Grids.CellFace @inbounds @inline setidx!(space, out, v - half, hidx, val) else @@ -179,14 +190,11 @@ Call `calc_level_val` on all the arguments of `bc`, and then apply the function """ Base.@propagate_inbounds function calc_level_val( bc::BC, + hidx, space, ) where {BC <: Base.Broadcast.Broadcasted} - i = threadIdx().y - j = blockIdx().y - v = threadIdx().x - h = blockIdx().z resolved_args = @inbounds @inline UnrolledUtilities.unrolled_map( - Base.Fix2(reconstruct_space_and_call_calc_level_val, space), + Base.Fix2(reconstruct_space_and_call_calc_level_val, (hidx, space)), bc.args, ) return @inline @inbounds bc.f(resolved_args...) @@ -201,15 +209,19 @@ us to use Base.Fix2. """ Base.@propagate_inbounds reconstruct_space_and_call_calc_level_val( arg::A, - space::S, + space_idx_tpl::S, ) where { A <: Union{Base.Broadcast.Broadcasted{<:AbstractFieldStyle}, StencilBroadcasted, Field}, S, -} = @inbounds @inline calc_level_val(arg, reconstruct_placeholder_space(axes(arg), space)) +} = @inbounds @inline calc_level_val( + arg, + space_idx_tpl[1], + reconstruct_placeholder_space(axes(arg), space_idx_tpl[2]), +) Base.@propagate_inbounds reconstruct_space_and_call_calc_level_val( arg::A, - space::S, -) where {A, S} = @inbounds @inline calc_level_val(arg, space) + space_idx_tpl::S, +) where {A, S} = @inbounds @inline calc_level_val(arg, space_idx_tpl[1], space_idx_tpl[2]) """ calc_level_val(val::T, space) @@ -217,10 +229,10 @@ Base.@propagate_inbounds reconstruct_space_and_call_calc_level_val( If `val` is not a `Broadcasted`, `StencilBroadcasted`, or `Field`, just return `val`. If it is a `Ref`, return `val[]`. If it is a one element tuple, return the element. """ -Base.@propagate_inbounds calc_level_val(val::T, space) where {T <: Ref} = val[] -Base.@propagate_inbounds calc_level_val(val::T, space) where {V, T <: Tuple{V}} = +Base.@propagate_inbounds calc_level_val(val::T, hidx, space) where {T <: Ref} = val[] +Base.@propagate_inbounds calc_level_val(val::T, hidx, space) where {V, T <: Tuple{V}} = first(val) -Base.@propagate_inbounds calc_level_val(arg::S, space) where {S} = arg +Base.@propagate_inbounds calc_level_val(arg::S, hidx, space) where {S} = arg """ calc_level_val(bc::StencilBroadcasted{<:Any, <: MultiplyColumnwiseBandMatrixField}, space) @@ -230,6 +242,7 @@ and then perform the multiplication. """ Base.@propagate_inbounds function calc_level_val( bc::BC, + hidx, space, ) where { S, @@ -238,17 +251,17 @@ Base.@propagate_inbounds function calc_level_val( } if check_if_fits_in_shmem(bc.args[2i32]) v = threadIdx().x - i = threadIdx().y + block_col_idx = threadIdx().y mat1_space = reconstruct_placeholder_space(axes(bc.args[1i32]), space) mat2_space = reconstruct_placeholder_space(axes(bc.args[2i32]), space) - mat2_row = calc_level_val(bc.args[2i32], mat2_space) - mat1_row = calc_level_val(bc.args[1i32], mat1_space) + mat2_row = calc_level_val(bc.args[2i32], hidx, mat2_space) + mat1_row = calc_level_val(bc.args[1i32], hidx, mat1_space) # project before placing in shared memory to avoid projecting multiple times mat2_row_converted = - @inbounds @inline project_row2_for_mul(mat1_row, mat2_row, mat2_space) + @inbounds @inline project_row2_for_mul(mat1_row, mat2_row, hidx, mat2_space) # It should be possible to use static shared memory here, but it allocates new shared memory # for each layer of recursion CUDA.sync_threads() @@ -258,7 +271,7 @@ Base.@propagate_inbounds function calc_level_val( typeof(mat2_row_converted), CUDA.blockDim().x * CUDA.blockDim().y, ) - @inbounds mat2[v + (i - 1) * CUDA.blockDim().x] = mat2_row_converted + @inbounds mat2[v + (block_col_idx - 1) * CUDA.blockDim().x] = mat2_row_converted CUDA.sync_threads() # if the output is on centers, the CUDA.blockDim().xth thread can just return 0 mat1_space.staggering isa Spaces.CellCenter && v == CUDA.blockDim().x && @@ -301,11 +314,7 @@ Base.@propagate_inbounds function calc_level_val( end else # values that won't fit in shmmem should just call getidx - i = threadIdx().y - j = blockIdx().y v = threadIdx().x - h = blockIdx().z - hidx = (i, j, h) if space.staggering isa Spaces.CellCenter v == CUDA.blockDim().x && return @inline @inbounds new(eltype(bc)) end @@ -324,13 +333,10 @@ will error if the operator is eagerly evaluated at the boundaries. """ Base.@propagate_inbounds function calc_level_val( bc::BC, + hidx, space, ) where {S, Op <: Operators.LinVanLeerC2F, BC <: StencilBroadcasted{S, Op}} - i = threadIdx().y - j = blockIdx().y v = threadIdx().x - h = blockIdx().z - hidx = (i, j, h) if v == 1i32 || v == CUDA.blockDim().x return zero(eltype(bc)) end @@ -346,13 +352,10 @@ affine BCs or values that won't fit in shmmem. """ Base.@propagate_inbounds function calc_level_val( bc::BC, + hidx, space, ) where {BC <: StencilBroadcasted} - i = threadIdx().y - j = blockIdx().y v = threadIdx().x - h = blockIdx().z - hidx = (i, j, h) if space.staggering isa Spaces.CellCenter v == CUDA.blockDim().x && return @inline @inbounds new(eltype(bc)) end @@ -369,13 +372,12 @@ When the staggering of `space` is `CellCenter`, the thread with `v == CUDA.block """ Base.@propagate_inbounds function calc_level_val( arg::F, + hidx, space, ) where {F <: Field} data = field_values(arg) - i = threadIdx().y - j = blockIdx().y v = threadIdx().x - h = blockIdx().z + (i, j, h) = hidx if space isa Union{Spaces.ExtrudedFiniteDifferenceSpace, Spaces.FiniteDifferenceSpace} && space.staggering isa Spaces.CellCenter @@ -391,6 +393,7 @@ Return the correct row of the operator matrix for the current thread """ Base.@propagate_inbounds function calc_level_val( bc::BC, + hidx, space, ) where { S, @@ -399,7 +402,7 @@ Base.@propagate_inbounds function calc_level_val( } op = bc.op.op args = bc.args - val = @inbounds @inline get_op_row(op, args, space) + val = @inbounds @inline get_op_row(op, args, hidx, space) return val end @@ -409,13 +412,9 @@ end Get the correct row of the operator matrix for the current thread, taking into account boundary conditions. """ -Base.@propagate_inbounds function get_op_row(op, args, space) +Base.@propagate_inbounds function get_op_row(op, args, hidx, space) FT = Spaces.undertype(space) - i = threadIdx().y - j = blockIdx().y v = threadIdx().x - h = blockIdx().z - hidx = (i, j, h) outputs_to_face = space.staggering isa ClimaCore.Grids.CellFace row_type = @inbounds @inline op_matrix_row_type(op, FT, args[1:(end - 1)]...) @@ -468,17 +467,14 @@ end Projects `mat2_row` onto the correct axis for multiplication with `mat1_row` if necessary, and returns the projected row. """ -Base.@propagate_inbounds function project_row2_for_mul(mat1_row, mat2_row, space) +Base.@propagate_inbounds function project_row2_for_mul(mat1_row, mat2_row, hidx, space) mat1_et = mat1_row isa BandMatrixRow ? eltype(mat1_row) : typeof(mat1_row) project_onto = ClimaCore.Geometry.recursively_find_dual_axes_for_projection(mat1_et) isnothing(project_onto) && return mat2_row v = threadIdx().x - i = threadIdx().y - j = blockIdx().y - v = threadIdx().x - h = blockIdx().z - hidx = (i, j, h) + project_onto = + ClimaCore.Geometry.recursively_find_dual_axes_for_projection(mat1_et) if space.staggering isa Spaces.CellCenter && v == CUDA.blockDim().x lg = new(Spaces.local_geometry_type(typeof(space))) else diff --git a/ext/cuda/operators_finite_difference.jl b/ext/cuda/operators_finite_difference.jl index d8f49ffedd..d534467630 100644 --- a/ext/cuda/operators_finite_difference.jl +++ b/ext/cuda/operators_finite_difference.jl @@ -86,21 +86,34 @@ function Base.copyto!( # this block config is better for VIJFH. It is only used when the total number of # threads in a block is between 32 and 256 to avoid underutilization of the GPU and # errors due to too many registers used when the block size is too large. - if !Topologies.isperiodic(space) && mask isa NoMask && - 32 <= n_face_levels * Ni <= 256 + # TODO: auto reduce max reg usage when needed because of high res columns + if !Topologies.isperiodic(space) && n_face_levels ≤ 256 + # 32 <= n_face_levels * Ni <= 256 + n_columns = mask isa NoMask ? Ni * Nj * Nh : mask.N[1] + # 108 is the number of SMs in an A100. TODO: get this value from CUDA.jl to better optimize for different GPUs + threads_dim_y = n_columns > 256 * 108 ? div(256, n_face_levels) : 1 + block_dim_x = div(n_columns, threads_dim_y, RoundUp) op_matrix_bc = replace_fd_ops(bc′) + cart_inds = if mask isa NoMask + CartesianIndices(map(Base.OneTo, (Ni, Nj, Nh))) + else + nothing + end args = ( strip_space(out, space), strip_space(op_matrix_bc, space), + cart_inds, + mask, axes(out), ) + auto_launch!( eager_copyto_stencil_kernel!, args; - threads_s = (n_face_levels, Ni, 1), - blocks_s = (1, Nj, Nh), + threads_s = (n_face_levels, threads_dim_y, 1), + blocks_s = (block_dim_x, 1, 1), always_inline = true, - shmem = n_face_levels * Ni * 9 * 4, # see `check_if_fits_in_shmem` for how this is calculated + shmem = n_face_levels * threads_dim_y * 9 * 4, # see `check_if_fits_in_shmem` for how this is calculated ) return out end diff --git a/src/MatrixFields/matrix_multiplication.jl b/src/MatrixFields/matrix_multiplication.jl index 4c436b117c..f4bd2e344b 100644 --- a/src/MatrixFields/matrix_multiplication.jl +++ b/src/MatrixFields/matrix_multiplication.jl @@ -269,17 +269,6 @@ function Operators.return_eltype( matrix1, arg, ) - # this is needed to support the divergence operator matrix applied to a Tensor field - if ( - matrix1 isa Operators.StencilBroadcasted && matrix1.op isa FDOperatorMatrix && - !(eltype(arg) <: BandMatrixRow) - ) - if matrix1.op.op isa OneArgFDOperator - return Operators.return_eltype(matrix1.op.op, arg) - else - return Operators.return_eltype(matrix1.op.op, matrix1.args[1], arg) - end - end et_mat1 = eltype(matrix1) et_arg = eltype(arg) et_mat1 <: BandMatrixRow || error( diff --git a/src/MatrixFields/operator_matrices.jl b/src/MatrixFields/operator_matrices.jl index e445e2cdcf..2d986b48c8 100644 --- a/src/MatrixFields/operator_matrices.jl +++ b/src/MatrixFields/operator_matrices.jl @@ -115,6 +115,250 @@ Base.Broadcast.broadcasted( Fields.local_geometry_field(operator_input_space(op_matrix.op, axes(arg))), ) +# TODO:specify these are only for setvalue +modifies_output( + op, + boundary_condition::BC, +) where {BC <: Union{ + Operators.SetGradient, + Operators.SetDivergence, + Operators.SetCurl, + }} = true +modifies_output( + op::Operators.InterpolationOperator, + boundary_condition::Operators.SetValue, +) = true +modifies_output( + op::Union{Operators.GradientF2C, Operators.DivergenceF2C}, + boundary_condition::Operators.SetValue, +) = false +modifies_output(op, boundary_condition::Operators.SetValue) = true +modifies_output(op, boundary_condition) = false + + +modifies_input( + op::Union{Operators.GradientF2C, Operators.DivergenceF2C}, + boundary_condition::Operators.SetValue, +) = true +modifies_input(op, boundary_condition) = false + + +function Operators.StencilBroadcasted{Style}( + op::OneArgFDOperator, + args::Args, + axes::Spaces.AbstractSpace, + work::Work = nothing, +) where {Style, Args, Work} + # can have 0, 1, or 2 boundary conds + if op isa Operators.SetBoundaryOperator + return Operators.StencilBroadcasted{ + Style, + typeof(op), + Args, + typeof(axes), + Work, + }( + op, + args, + axes, + work, + ) + end + if length(op.bcs) == 0 + maybe_adjointed = + op isa Operators.GradientOperator ? + Base.Broadcast.broadcasted(adjoint, args[1]) : args[1] + new_args = ( + Base.Broadcast.broadcasted( + FDOperatorMatrix(op), + Fields.local_geometry_field(operator_input_space(op, axes)), + ), maybe_adjointed) + raw_mul_bc = Operators.StencilBroadcasted{ + Style, + MultiplyColumnwiseBandMatrixField, + typeof(new_args), + typeof(axes), + Work, + }( + MultiplyColumnwiseBandMatrixField(), + new_args, + axes, + work, + ) + return op isa Operators.DivergenceOperator ? + Base.Broadcast.broadcasted(adjoint, raw_mul_bc) : raw_mul_bc + end + has_two_bcs = length(op.bcs) == 2 + remove_bc1 = modifies_input(op, op.bcs[1]) || modifies_output(op, op.bcs[1]) + if has_two_bcs + remove_bc2 = modifies_input(op, op.bcs[2]) || modifies_output(op, op.bcs[2]) + new_bcs = + remove_bc1 && remove_bc2 ? NamedTuple{}() : + remove_bc1 ? NamedTuple{(keys(op.bcs)[2],)}((op.bcs[2],)) : + remove_bc2 ? NamedTuple{(keys(op.bcs)[1],)}((op.bcs[1],)) : + op.bcs + else + new_bcs = remove_bc1 ? NamedTuple{}() : op.bcs + end + op_mat = Base.Broadcast.broadcasted( + FDOperatorMatrix(Base.typename(typeof(op)).wrapper(new_bcs)), + Fields.local_geometry_field(operator_input_space(op, axes)), + ) + # Base.typename(typeof(bc.op)).wrapper() + if modifies_input(op, op.bcs[1]) || (has_two_bcs && modifies_input(op, op.bcs[2])) + if modifies_input(op, op.bcs[1]) + if has_two_bcs && modifies_input(op, op.bcs[2]) + inner_op = Operators.SetBoundaryOperator(op.bcs) + else + inner_op = + Operators.SetBoundaryOperator( + NamedTuple{(keys(op.bcs)[1],)}((op.bcs[1],)), + ) + end + else + inner_op = + Operators.SetBoundaryOperator(NamedTuple{(keys(op.bcs)[2],)}((op.bcs[2],))) + end + wrapped_inner_arg = Operators.StencilBroadcasted{ + Style, + typeof(inner_op), + Args, + typeof(Base.axes(args[1])), + Nothing, + }( + inner_op, + args, + Base.axes(args[1]), + nothing, + ) + else + wrapped_inner_arg = args[1] + end + maybe_adjointed_wrapped_inner = + op isa Operators.GradientOperator ? + Base.Broadcast.broadcasted(adjoint, wrapped_inner_arg) : wrapped_inner_arg + new_args = (op_mat, wrapped_inner_arg) + new_broadcasted_op = Operators.StencilBroadcasted{ + Style, + MultiplyColumnwiseBandMatrixField, + typeof(new_args), + typeof(axes), + Work, + }( + MultiplyColumnwiseBandMatrixField(), + new_args, + axes, + work, + ) + + if modifies_output(op, op.bcs[1]) || (has_two_bcs && modifies_output(op, op.bcs[2])) + if modifies_output(op, op.bcs[1]) + if has_two_bcs && modifies_output(op, op.bcs[2]) + outer_op = Operators.SetBoundaryOperator(op.bcs) + else + outer_op = + Operators.SetBoundaryOperator( + NamedTuple{(keys(op.bcs)[1],)}((op.bcs[1],)), + ) + end + else + outer_op = + Operators.SetBoundaryOperator(NamedTuple{(keys(op.bcs)[2],)}((op.bcs[2],))) + end + raw_mul_with_setbcs = Operators.StencilBroadcasted{ + Style, + typeof(outer_op), + Tuple{typeof(new_broadcasted_op)}, + typeof(axes), + Work, + }( + outer_op, + (new_broadcasted_op,), + axes, + work, + ) + else + raw_mul_with_setbcs = new_broadcasted_op + end + return op isa Operators.DivergenceOperator ? + Base.Broadcast.broadcasted(adjoint, raw_mul_with_setbcs) : raw_mul_with_setbcs +end + + +function Operators.StencilBroadcasted{Style}( + op::TwoArgFDOperator, + args::Args, + axes::Spaces.AbstractSpace, + work::Work = nothing, +) where {Style, Args, Work} + # can have 0, 1, or 2 boundary conds + if op isa Operators.WeightedInterpolateC2F && length(op.bcs) > 0 && + any(bc -> bc isa Operators.SetValue, values(op.bcs)) + has_two_bcs = length(op.bcs) == 2 + remove_bc1 = modifies_output(op, op.bcs[1]) + if has_two_bcs + remove_bc2 = modifies_output(op, op.bcs[2]) + new_bcs = + remove_bc1 && remove_bc2 ? NamedTuple{}() : + remove_bc1 ? NamedTuple{(keys(op.bcs)[2],)}((op.bcs[2],)) : + remove_bc2 ? NamedTuple{(keys(op.bcs)[1],)}((op.bcs[1],)) : + op.bcs + outer_op_bcs = + remove_bc1 && remove_bc2 ? op.bcs : + remove_bc1 ? NamedTuple{(keys(op.bcs)[1],)}((op.bcs[1],)) : + remove_bc2 ? NamedTuple{(keys(op.bcs)[2],)}((op.bcs[2],)) : + NamedTuple{}() + else + new_bcs = remove_bc1 ? NamedTuple{}() : op.bcs + outer_op_bcs = remove_bc1 ? op.bcs : NamedTuple{}() + end + new_op = Base.typename(typeof(op)).wrapper(new_bcs) + outer_op = + length(outer_op_bcs) > 0 ? Operators.SetBoundaryOperator(outer_op_bcs) : nothing + else + new_op = op + outer_op = nothing + end + + + new_args = ( + Base.Broadcast.broadcasted( + FDOperatorMatrix(new_op), + args[1], + ), args[2]) + result_without_setvalue_applied = Operators.StencilBroadcasted{ + Style, + MultiplyColumnwiseBandMatrixField, + typeof(new_args), + typeof(axes), + Work, + }( + MultiplyColumnwiseBandMatrixField(), + new_args, + axes, + work, + ) + + if isnothing(outer_op) + return result_without_setvalue_applied + else + Operators.StencilBroadcasted{ + Style, + typeof(outer_op), + Tuple{typeof(result_without_setvalue_applied)}, + typeof(axes), + Work, + }( + outer_op, + (result_without_setvalue_applied,), + axes, + work, + ) + end +end + + + """ operator_matrix(op) @@ -315,18 +559,18 @@ Operators.stencil_left_boundary( op_matrix::FDOperatorMatrix, ::Operators.NullBoundaryCondition, space, - _, - _, + idx, + hidx, args..., -) = zero(Operators.return_eltype(op_matrix, args...)) +) = Operators.stencil_interior(op_matrix, space, idx, hidx, args...) Operators.stencil_right_boundary( op_matrix::FDOperatorMatrix, ::Operators.NullBoundaryCondition, space, - _, - _, + idx, + hidx, args..., -) = zero(Operators.return_eltype(op_matrix, args...)) +) = Operators.stencil_interior(op_matrix, space, idx, hidx, args...) ################################################################################ @@ -377,16 +621,6 @@ op_matrix_interior_row( ::Union{Operators.InterpolateC2F, Operators.InterpolateF2C}, ::Type{FT}, ) where {FT} = BidiagonalMatrixRow(FT(1), FT(1)) / 2 -op_matrix_first_row( - ::Operators.InterpolateC2F, - ::Operators.SetValue, - ::Type{FT}, -) where {FT} = UpperDiagonalMatrixRow(FT(0)) -op_matrix_last_row( - ::Operators.InterpolateC2F, - ::Operators.SetValue, - ::Type{FT}, -) where {FT} = LowerDiagonalMatrixRow(FT(0)) op_matrix_first_row( ::Operators.InterpolateC2F, ::Operators.Extrapolate, @@ -402,31 +636,11 @@ op_matrix_interior_row( ::Union{Operators.LeftBiasedC2F, Operators.LeftBiasedF2C}, ::Type{FT}, ) where {FT} = LowerDiagonalMatrixRow(FT(1)) -op_matrix_first_row( - ::Operators.LeftBiasedC2F, - ::Operators.SetValue, - ::Type{FT}, -) where {FT} = LowerEmptyMatrixRow() -op_matrix_first_row( - ::Operators.LeftBiasedF2C, - ::Operators.SetValue, - ::Type{FT}, -) where {FT} = LowerDiagonalMatrixRow(FT(0)) op_matrix_interior_row( ::Union{Operators.RightBiasedC2F, Operators.RightBiasedF2C}, ::Type{FT}, ) where {FT} = UpperDiagonalMatrixRow(FT(1)) -op_matrix_last_row( - ::Operators.RightBiasedC2F, - ::Operators.SetValue, - ::Type{FT}, -) where {FT} = UpperEmptyMatrixRow() -op_matrix_last_row( - ::Operators.RightBiasedF2C, - ::Operators.SetValue, - ::Type{FT}, -) where {FT} = UpperDiagonalMatrixRow(FT(0)) op_matrix_row_type( ::Operators.WeightedInterpolationOperator, @@ -445,16 +659,6 @@ Base.@propagate_inbounds function op_matrix_interior_row( denominator = w⁻ + w⁺ return BidiagonalMatrixRow(w⁻ / denominator, w⁺ / denominator) end -op_matrix_first_row( - ::Operators.WeightedInterpolateC2F, - ::Operators.SetValue, - ::Type{FT}, -) where {FT} = UpperDiagonalMatrixRow(FT(0)) -op_matrix_last_row( - ::Operators.WeightedInterpolateC2F, - ::Operators.SetValue, - ::Type{FT}, -) where {FT} = LowerDiagonalMatrixRow(FT(0)) op_matrix_first_row( ::Operators.WeightedInterpolateC2F, ::Operators.Extrapolate, @@ -578,42 +782,12 @@ op_matrix_row_type(::Operators.AdvectionOperator, ::Type{FT}, _) where {FT} = op_matrix_interior_row(::Operators.SetBoundaryOperator, ::Type{FT}) where {FT} = DiagonalMatrixRow(FT(1)) -op_matrix_first_row( - ::Operators.SetBoundaryOperator, - ::Operators.SetValue, - ::Type{FT}, -) where {FT} = DiagonalMatrixRow(FT(0)) -op_matrix_last_row( - ::Operators.SetBoundaryOperator, - ::Operators.SetValue, - ::Type{FT}, -) where {FT} = DiagonalMatrixRow(FT(0)) op_matrix_row_type(op::Operators.GradientOperator, ::Type{FT}) where {FT} = uses_extrapolate(op) ? QuaddiagonalMatrixRow{C3{FT}} : BidiagonalMatrixRow{C3{FT}} op_matrix_interior_row(::Operators.GradientOperator, ::Type{FT}) where {FT} = BidiagonalMatrixRow(-C3(FT(1)), C3(FT(1))) -op_matrix_first_row( - ::Operators.GradientC2F, - ::Operators.SetGradient, - ::Type{FT}, -) where {FT} = UpperDiagonalMatrixRow(C3(FT(0))) -op_matrix_last_row( - ::Operators.GradientC2F, - ::Operators.SetGradient, - ::Type{FT}, -) where {FT} = LowerDiagonalMatrixRow(C3(FT(0))) -op_matrix_first_row( - ::Operators.GradientF2C, - ::Operators.SetValue, - ::Type{FT}, -) where {FT} = BidiagonalMatrixRow(C3(FT(0)), C3(FT(1))) -op_matrix_last_row( - ::Operators.GradientF2C, - ::Operators.SetValue, - ::Type{FT}, -) where {FT} = BidiagonalMatrixRow(-C3(FT(1)), C3(FT(0))) op_matrix_first_row( ::Operators.GradientF2C, ::Operators.Extrapolate, @@ -639,50 +813,6 @@ Base.@propagate_inbounds function op_matrix_interior_row( J⁺ = Geometry.LocalGeometry(space, idx + half, hidx).J return BidiagonalMatrixRow(-C3(J⁻)', C3(J⁺)') * invJ end -op_matrix_first_row( - ::Operators.DivergenceC2F, - ::Operators.SetDivergence, - ::Type{FT}, -) where {FT} = UpperDiagonalMatrixRow(C3(FT(0))') -op_matrix_last_row( - ::Operators.DivergenceC2F, - ::Operators.SetDivergence, - ::Type{FT}, -) where {FT} = LowerDiagonalMatrixRow(C3(FT(0))') -Base.@propagate_inbounds function op_matrix_first_row( - ::Operators.DivergenceF2C, - ::Operators.SetValue, - space, - idx, - hidx, -) - FT = Spaces.undertype(space) - invJ = Geometry.LocalGeometry(space, idx, hidx).invJ - J⁺ = Geometry.LocalGeometry(space, idx + half, hidx).J - return BidiagonalMatrixRow(C3(FT(0))', C3(J⁺)') * invJ -end -Base.@propagate_inbounds function op_matrix_last_row( - ::Operators.DivergenceF2C, - ::Operators.SetValue, - space, - idx, - hidx, -) - FT = Spaces.undertype(space) - invJ = Geometry.LocalGeometry(space, idx, hidx).invJ - J⁻ = Geometry.LocalGeometry(space, idx - half, hidx).J - return BidiagonalMatrixRow(-C3(J⁻)', C3(FT(0))') * invJ -end -op_matrix_first_row( - ::Operators.DivergenceF2C, - ::Operators.SetDivergence, - ::Type{FT}, -) where {FT} = BidiagonalMatrixRow(C3(FT(0))', C3(FT(0))') -op_matrix_last_row( - ::Operators.DivergenceF2C, - ::Operators.SetDivergence, - ::Type{FT}, -) where {FT} = BidiagonalMatrixRow(C3(FT(0))', C3(FT(0))') Base.@propagate_inbounds function op_matrix_first_row( ::Operators.DivergenceF2C, ::Operators.Extrapolate, @@ -723,13 +853,3 @@ Base.@propagate_inbounds function op_matrix_interior_row( invJ = Geometry.LocalGeometry(space, idx, hidx).invJ return BidiagonalMatrixRow(-εⁱʲ, εⁱʲ) * invJ end -op_matrix_first_row( - ::Operators.CurlC2F, - ::Operators.SetCurl, - ::Type{FT}, -) where {FT} = UpperDiagonalMatrixRow(zero(CT12_CT12{FT})) -op_matrix_last_row( - ::Operators.CurlC2F, - ::Operators.SetCurl, - ::Type{FT}, -) where {FT} = LowerDiagonalMatrixRow(zero(CT12_CT12{FT})) diff --git a/src/Operators/finitedifference.jl b/src/Operators/finitedifference.jl index 1ba1bd2645..96ae93ed99 100644 --- a/src/Operators/finitedifference.jl +++ b/src/Operators/finitedifference.jl @@ -391,17 +391,6 @@ return_space(::InterpolateF2C, space::AllFaceFiniteDifferenceSpace) = Spaces.space(space, Spaces.CellCenter()) stencil_interior_width(::InterpolateF2C, arg) = ((-half, half),) -Base.@propagate_inbounds function stencil_interior( - ::InterpolateF2C, - space, - idx, - hidx, - arg, -) - a⁺ = getidx(space, arg, idx + half, hidx) - a⁻ = getidx(space, arg, idx - half, hidx) - (a⁺ + a⁻) / 2 -end boundary_width(::InterpolateF2C, ::AbstractBoundaryCondition) = 0 @@ -445,66 +434,9 @@ return_space(::InterpolateC2F, space::AllCenterFiniteDifferenceSpace) = Spaces.space(space, Spaces.CellFace()) stencil_interior_width(::InterpolateC2F, arg) = ((-half, half),) -Base.@propagate_inbounds function stencil_interior( - ::InterpolateC2F, - space, - idx, - hidx, - arg, -) - a⁺ = getidx(space, arg, idx + half, hidx) - a⁻ = getidx(space, arg, idx - half, hidx) - (a⁺ + a⁻) / 2 -end -boundary_width(::InterpolateC2F, ::AbstractBoundaryCondition) = 1 -Base.@propagate_inbounds function stencil_left_boundary( - ::InterpolateC2F, - bc::SetValue, - space, - idx, - hidx, - arg, -) - @assert idx == left_face_boundary_idx(space) - getidx(space, bc.val, nothing, hidx) -end -Base.@propagate_inbounds function stencil_right_boundary( - ::InterpolateC2F, - bc::SetValue, - space, - idx, - hidx, - arg, -) - @assert idx == right_face_boundary_idx(space) - getidx(space, bc.val, nothing, hidx) -end +boundary_width(::InterpolateC2F, ::AbstractBoundaryCondition) = 1 -Base.@propagate_inbounds function stencil_left_boundary( - ::InterpolateC2F, - bc::Extrapolate, - space, - idx, - hidx, - arg, -) - @assert idx == left_face_boundary_idx(space) - a⁺ = getidx(space, arg, idx + half, hidx) - a⁺ -end -Base.@propagate_inbounds function stencil_right_boundary( - ::InterpolateC2F, - bc::Extrapolate, - space, - idx, - hidx, - arg, -) - @assert idx == right_face_boundary_idx(space) - a⁻ = getidx(space, arg, idx - half, hidx) - a⁻ -end """ L = LeftBiasedC2F(;boundaries) @@ -534,13 +466,6 @@ return_space(::LeftBiasedC2F, space::AllCenterFiniteDifferenceSpace) = Spaces.space(space, Spaces.CellFace()) stencil_interior_width(::LeftBiasedC2F, arg) = ((-half, -half),) -Base.@propagate_inbounds stencil_interior( - ::LeftBiasedC2F, - space, - idx, - hidx, - arg, -) = getidx(space, arg, idx - half, hidx) left_interior_idx( space::AbstractSpace, @@ -555,18 +480,6 @@ right_interior_idx( arg, ) = right_idx(space) -Base.@propagate_inbounds function stencil_left_boundary( - ::LeftBiasedC2F, - bc::SetValue, - space, - idx, - hidx, - arg, -) - @assert idx == left_face_boundary_idx(space) - getidx(space, bc.val, nothing, hidx) -end - """ L = LeftBiasedF2C(;boundaries) L.(x) @@ -790,13 +703,6 @@ return_space(::RightBiasedC2F, space::AllCenterFiniteDifferenceSpace) = Spaces.space(space, Spaces.CellFace()) stencil_interior_width(::RightBiasedC2F, arg) = ((half, half),) -Base.@propagate_inbounds stencil_interior( - ::RightBiasedC2F, - space, - idx, - hidx, - arg, -) = getidx(space, arg, idx + half, hidx) left_interior_idx( space::AbstractSpace, @@ -811,18 +717,6 @@ right_interior_idx( arg, ) = right_idx(space) - 1 -Base.@propagate_inbounds function stencil_right_boundary( - ::RightBiasedC2F, - bc::SetValue, - space, - idx, - hidx, - arg, -) - @assert idx == right_face_boundary_idx(space) - getidx(space, bc.val, nothing, hidx) -end - """ R = RightBiasedF2C(;boundaries) R.(x) @@ -851,13 +745,6 @@ return_space(::RightBiasedF2C, space::AllFaceFiniteDifferenceSpace) = Spaces.space(space, Spaces.CellCenter()) stencil_interior_width(::RightBiasedF2C, arg) = ((half, half),) -Base.@propagate_inbounds stencil_interior( - ::RightBiasedF2C, - space, - idx, - hidx, - arg, -) = getidx(space, arg, idx + half, hidx) left_interior_idx( space::AbstractSpace, @@ -874,18 +761,6 @@ right_interior_idx( right_interior_idx(space::AbstractSpace, ::RightBiasedF2C, ::SetValue, arg) = right_idx(space) - 1 -Base.@propagate_inbounds function stencil_right_boundary( - ::RightBiasedF2C, - bc::SetValue, - space, - idx, - hidx, - arg, -) - @assert idx == right_center_boundary_idx(space) - getidx(space, bc.val, nothing, hidx) -end - """ R = RightBiased3rdOrderC2F(;boundaries) @@ -1218,51 +1093,8 @@ upwind_biased_product(v, a⁻, a⁺) = ((v + abs(v)) * a⁻ + (v - abs(v)) * a stencil_interior_width(::UpwindBiasedProductC2F, velocity, arg) = ((0, 0), (-half, half)) -Base.@propagate_inbounds function stencil_interior( - ::UpwindBiasedProductC2F, - space, - idx, - hidx, - velocity, - arg, -) - a⁻ = stencil_interior(LeftBiasedC2F(), space, idx, hidx, arg) - a⁺ = stencil_interior(RightBiasedC2F(), space, idx, hidx, arg) - vᶠ = Geometry.contravariant3( - getidx(space, velocity, idx, hidx), - Geometry.LocalGeometry(space, idx, hidx), - ) - return Geometry.Contravariant3Vector(upwind_biased_product(vᶠ, a⁻, a⁺)) -end - boundary_width(::UpwindBiasedProductC2F, ::AbstractBoundaryCondition) = 1 -Base.@propagate_inbounds function stencil_left_boundary( - op::UpwindBiasedProductC2F, - ::Extrapolate, - space, - idx, - hidx, - velocity, - arg, -) - @assert idx == left_face_boundary_idx(space) - stencil_interior(op, space, idx + 1, hidx, velocity, arg) -end - -Base.@propagate_inbounds function stencil_right_boundary( - op::UpwindBiasedProductC2F, - ::Extrapolate, - space, - idx, - hidx, - velocity, - arg, -) - @assert idx == right_face_boundary_idx(space) - stencil_interior(op, space, idx - 1, hidx, velocity, arg) -end - """ LinVanLeerC2F @@ -1537,118 +1369,14 @@ return_space( arg_space::AllCenterFiniteDifferenceSpace, ) = velocity_space -upwind_3rdorder_biased_product(v, a⁻, a⁻⁻, a⁺, a⁺⁺) = - ( - v * (7 * (a⁺ + a⁻) - (a⁺⁺ + a⁻⁻)) - - abs(v) * (3 * (a⁺ - a⁻) - (a⁺⁺ - a⁻⁻)) - ) / 12 stencil_interior_width(::Upwind3rdOrderBiasedProductC2F, velocity, arg) = ((0, 0), (-half - 1, half + 1)) -Base.@propagate_inbounds function stencil_interior( - ::Upwind3rdOrderBiasedProductC2F, - space, - idx, - hidx, - velocity, - arg, -) - a⁻ = getidx(space, arg, idx - half, hidx) - a⁻⁻ = getidx(space, arg, idx - half - 1, hidx) - a⁺ = getidx(space, arg, idx + half, hidx) - a⁺⁺ = getidx(space, arg, idx + half + 1, hidx) - vᶠ = Geometry.contravariant3( - getidx(space, velocity, idx, hidx), - Geometry.LocalGeometry(space, idx, hidx), - ) - return Geometry.Contravariant3Vector( - upwind_3rdorder_biased_product(vᶠ, a⁻, a⁻⁻, a⁺, a⁺⁺), - ) -end boundary_width(::Upwind3rdOrderBiasedProductC2F, ::AbstractBoundaryCondition) = 2 -Base.@propagate_inbounds function stencil_left_boundary( - ::Upwind3rdOrderBiasedProductC2F, - bc::FirstOrderOneSided, - space, - idx, - hidx, - velocity, - arg, -) - @assert idx <= left_face_boundary_idx(space) + 1 - v = Geometry.contravariant3( - getidx(space, velocity, idx, hidx), - Geometry.LocalGeometry(space, idx, hidx), - ) - a⁻ = stencil_interior(LeftBiasedC2F(), space, idx, hidx, arg) - a⁺ = stencil_interior(RightBiased3rdOrderC2F(), space, idx, hidx, arg) - return Geometry.Contravariant3Vector(upwind_biased_product(v, a⁻, a⁺)) -end - -Base.@propagate_inbounds function stencil_right_boundary( - ::Upwind3rdOrderBiasedProductC2F, - bc::FirstOrderOneSided, - space, - idx, - hidx, - velocity, - arg, -) - @assert idx >= right_face_boundary_idx(space) - 1 - v = Geometry.contravariant3( - getidx(space, velocity, idx, hidx), - Geometry.LocalGeometry(space, idx, hidx), - ) - a⁻ = stencil_interior(LeftBiased3rdOrderC2F(), space, idx, hidx, arg) - a⁺ = stencil_interior(RightBiasedC2F(), space, idx, hidx, arg) - return Geometry.Contravariant3Vector(upwind_biased_product(v, a⁻, a⁺)) - -end - -Base.@propagate_inbounds function stencil_left_boundary( - ::Upwind3rdOrderBiasedProductC2F, - bc::ThirdOrderOneSided, - space, - idx, - hidx, - velocity, - arg, -) - @assert idx <= left_face_boundary_idx(space) + 1 - - vᶠ = Geometry.contravariant3( - getidx(space, velocity, idx, hidx), - Geometry.LocalGeometry(space, idx, hidx), - ) - a = stencil_interior(RightBiased3rdOrderC2F(), space, idx, hidx, arg) - - return Geometry.Contravariant3Vector(vᶠ * a) -end - -Base.@propagate_inbounds function stencil_right_boundary( - ::Upwind3rdOrderBiasedProductC2F, - bc::ThirdOrderOneSided, - space, - idx, - hidx, - velocity, - arg, -) - @assert idx <= right_face_boundary_idx(space) - 1 - - vᶠ = Geometry.contravariant3( - getidx(space, velocity, idx, hidx), - Geometry.LocalGeometry(space, idx, hidx), - ) - a = stencil_interior(LeftBiased3rdOrderC2F(), space, idx, hidx, arg) - - return Geometry.Contravariant3Vector(vᶠ * a) -end - """ U = FCTBorisBook(;boundaries) U.(v, x) @@ -2120,14 +1848,18 @@ abstract type BoundaryOperator <: FiniteDifferenceOperator end """ SetBoundaryOperator(;boundaries...) -This operator only modifies the values at the boundary: +This operator only modifies the values at the boundary faces, or the center cells adjacent to the boundary faces: - [`SetValue(val)`](@ref): set the value to be `val` on the boundary. """ struct SetBoundaryOperator{BCS} <: BoundaryOperator bcs::BCS function SetBoundaryOperator(; kwargs...) - assert_valid_bcs("SetBoundaryOperator", kwargs, (SetValue,)) + assert_valid_bcs( + "SetBoundaryOperator", + kwargs, + (SetValue, SetGradient, SetCurl, SetDivergence), + ) new{typeof(NamedTuple(kwargs))}(NamedTuple(kwargs)) end SetBoundaryOperator(bcs) = SetBoundaryOperator(; bcs...) @@ -2144,29 +1876,58 @@ Base.@propagate_inbounds stencil_interior( arg, ) = getidx(space, arg, idx, hidx) -boundary_width(::SetBoundaryOperator, ::AbstractBoundaryCondition) = 0 -boundary_width(::SetBoundaryOperator, ::SetValue) = 1 +boundary_width(::SetBoundaryOperator, ::AbstractBoundaryCondition) = 1 Base.@propagate_inbounds function stencil_left_boundary( ::SetBoundaryOperator, - bc::SetValue, + bc::Union{SetValue, SetGradient, SetCurl, SetDivergence}, space, idx, hidx, arg, ) - @assert idx == left_face_boundary_idx(space) - getidx(space, bc.val, nothing, hidx) + @assert idx == left_idx(space) + val = getidx(space, bc.val, nothing, hidx) + if bc isa SetGradient + # val = Geometry.gradient_result_type(Val((3,)), eltype(val))(val) + return Geometry.project( + Geometry.Covariant3Axis(), + val, + Geometry.LocalGeometry(space, idx, hidx), + ) + elseif bc isa SetCurl + return Geometry.project( + Geometry.Contravariant123Axis(), + val, + Geometry.LocalGeometry(space, idx, hidx), + ) + end + return val end Base.@propagate_inbounds function stencil_right_boundary( ::SetBoundaryOperator, - bc::SetValue, + bc::Union{SetValue, SetGradient, SetCurl, SetDivergence}, space, idx, hidx, arg, ) - @assert idx == right_face_boundary_idx(space) - getidx(space, bc.val, nothing, hidx) + @assert idx == right_idx(space) + val = getidx(space, bc.val, nothing, hidx) + if bc isa SetGradient + # val = Geometry.gradient_result_type(Val((3,)), eltype(val))(val) + return Geometry.project( + Geometry.Covariant3Axis(), + val, + Geometry.LocalGeometry(space, idx, hidx), + ) + elseif bc isa SetCurl + return Geometry.project( + Geometry.Contravariant123Axis(), + val, + Geometry.LocalGeometry(space, idx, hidx), + ) + end + return val end @@ -2217,83 +1978,11 @@ return_space(::GradientF2C, space::AllFaceFiniteDifferenceSpace) = Spaces.space(space, Spaces.CellCenter()) stencil_interior_width(::GradientF2C, arg) = ((-half, half),) -Base.@propagate_inbounds function stencil_interior( - ::GradientF2C, - space, - idx, - hidx, - arg, -) - Geometry.Covariant3Vector(1) ⊗ ( - getidx(space, arg, idx + half, hidx) - - getidx(space, arg, idx - half, hidx) - ) -end boundary_width(::GradientF2C, ::AbstractBoundaryCondition) = 0 boundary_width(::GradientF2C, ::SetValue) = 1 -Base.@propagate_inbounds function stencil_left_boundary( - ::GradientF2C, - bc::SetValue, - space, - idx, - hidx, - arg, -) - @assert idx == left_center_boundary_idx(space) - Geometry.Covariant3Vector(1) ⊗ ( - getidx(space, arg, idx + half, hidx) - - getidx(space, bc.val, nothing, hidx) - ) -end -Base.@propagate_inbounds function stencil_right_boundary( - ::GradientF2C, - bc::SetValue, - space, - idx, - hidx, - arg, -) - @assert idx == right_center_boundary_idx(space) - Geometry.Covariant3Vector(1) ⊗ ( - getidx(space, bc.val, nothing, hidx) - - getidx(space, arg, idx - half, hidx) - ) -end - boundary_width(::GradientF2C, ::Extrapolate) = 1 -Base.@propagate_inbounds function stencil_left_boundary( - op::GradientF2C, - ::Extrapolate, - space, - idx, - hidx, - arg, -) - @assert idx == left_center_boundary_idx(space) - Geometry.project( - Geometry.Covariant3Axis(), - stencil_interior(op, space, idx + 1, hidx, arg), - Geometry.LocalGeometry(space, idx, hidx), - ) -end -Base.@propagate_inbounds function stencil_right_boundary( - op::GradientF2C, - ::Extrapolate, - space, - idx, - hidx, - arg, -) - @assert idx == right_center_boundary_idx(space) - Geometry.project( - Geometry.Covariant3Axis(), - stencil_interior(op, space, idx - 1, hidx, arg), - Geometry.LocalGeometry(space, idx, hidx), - ) -end - """ G = GradientC2F(;boundaryname=boundarycondition...) G.(x) @@ -2324,55 +2013,8 @@ return_space(::GradientC2F, space::AllCenterFiniteDifferenceSpace) = Spaces.space(space, Spaces.CellFace()) stencil_interior_width(::GradientC2F, arg) = ((-half, half),) -Base.@propagate_inbounds function stencil_interior( - ::GradientC2F, - space, - idx, - hidx, - arg, -) - Geometry.Covariant3Vector(1) ⊗ ( - getidx(space, arg, idx + half, hidx) - - getidx(space, arg, idx - half, hidx) - ) -end - boundary_width(::GradientC2F, ::AbstractBoundaryCondition) = 1 -# left / right SetGradient boundary conditions -Base.@propagate_inbounds function stencil_left_boundary( - ::GradientC2F, - bc::SetGradient, - space, - idx, - hidx, - arg, -) - @assert idx == left_face_boundary_idx(space) - # imposed flux boundary condition at left most face - Geometry.project( - Geometry.Covariant3Axis(), - getidx(space, bc.val, nothing, hidx), - Geometry.LocalGeometry(space, idx, hidx), - ) -end -Base.@propagate_inbounds function stencil_right_boundary( - ::GradientC2F, - bc::SetGradient, - space, - idx, - hidx, - arg, -) - @assert idx == right_face_boundary_idx(space) - # imposed flux boundary condition at right most face - Geometry.project( - Geometry.Covariant3Axis(), - getidx(space, bc.val, nothing, hidx), - Geometry.LocalGeometry(space, idx, hidx), - ) -end - """ UG = UpwindBiasedGradient() UG.(v, θ) @@ -2522,115 +2164,8 @@ return_space(::DivergenceF2C, space::AllFaceFiniteDifferenceSpace) = Spaces.space(space, Spaces.CellCenter()) stencil_interior_width(::DivergenceF2C, arg) = ((-half, half),) -Base.@propagate_inbounds function stencil_interior( - ::DivergenceF2C, - space, - idx, - hidx, - arg, -) - local_geometry = Geometry.LocalGeometry(space, idx, hidx) - Ju³₊ = Geometry.Jcontravariant3( - getidx(space, arg, idx + half, hidx), - Geometry.LocalGeometry(space, idx + half, hidx), - ) - Ju³₋ = Geometry.Jcontravariant3( - getidx(space, arg, idx - half, hidx), - Geometry.LocalGeometry(space, idx - half, hidx), - ) - (Ju³₊ - Ju³₋) * local_geometry.invJ -end - boundary_width(::DivergenceF2C, ::AbstractBoundaryCondition) = 0 boundary_width(::DivergenceF2C, ::SetValue) = 1 -Base.@propagate_inbounds function stencil_left_boundary( - ::DivergenceF2C, - bc::SetValue, - space, - idx, - hidx, - arg, -) - @assert idx == left_center_boundary_idx(space) - local_geometry = Geometry.LocalGeometry(space, idx, hidx) - Ju³₊ = Geometry.Jcontravariant3( - getidx(space, arg, idx + half, hidx), - Geometry.LocalGeometry(space, idx + half, hidx), - ) - Ju³₋ = Geometry.Jcontravariant3( - getidx(space, bc.val, nothing, hidx), - Geometry.LocalGeometry(space, idx - half, hidx), - ) - (Ju³₊ - Ju³₋) * local_geometry.invJ -end -Base.@propagate_inbounds function stencil_right_boundary( - ::DivergenceF2C, - bc::SetValue, - space, - idx, - hidx, - arg, -) - @assert idx == right_center_boundary_idx(space) - local_geometry = Geometry.LocalGeometry(space, idx, hidx) - Ju³₊ = Geometry.Jcontravariant3( - getidx(space, bc.val, nothing, hidx), - Geometry.LocalGeometry(space, idx + half, hidx), - ) - Ju³₋ = Geometry.Jcontravariant3( - getidx(space, arg, idx - half, hidx), - Geometry.LocalGeometry(space, idx - half, hidx), - ) - (Ju³₊ - Ju³₋) * local_geometry.invJ -end - -boundary_width(::DivergenceF2C, ::SetDivergence) = 1 -Base.@propagate_inbounds function stencil_left_boundary( - ::DivergenceF2C, - bc::SetDivergence, - space, - idx, - hidx, - arg, -) - @assert idx == left_center_boundary_idx(space) - getidx(space, bc.val, nothing, hidx) -end -Base.@propagate_inbounds function stencil_right_boundary( - ::DivergenceF2C, - bc::SetDivergence, - space, - idx, - hidx, - arg, -) - @assert idx == right_center_boundary_idx(space) - getidx(space, bc.val, nothing, hidx) -end - -boundary_width(::DivergenceF2C, ::Extrapolate) = 1 -Base.@propagate_inbounds function stencil_left_boundary( - op::DivergenceF2C, - ::Extrapolate, - space, - idx, - hidx, - arg, -) - @assert idx == left_center_boundary_idx(space) - stencil_interior(op, space, idx + 1, hidx, arg) -end -Base.@propagate_inbounds function stencil_right_boundary( - op::DivergenceF2C, - ::Extrapolate, - space, - idx, - hidx, - arg, -) - @assert idx == right_center_boundary_idx(space) - stencil_interior(op, space, idx - 1, hidx, arg) -end # Extend `adapt_structure` for all boundary conditions containing a `val` field. function Adapt.adapt_structure(to, bc::AbstractBoundaryCondition) @@ -2689,51 +2224,8 @@ return_space(::DivergenceC2F, space::AllCenterFiniteDifferenceSpace) = Spaces.space(space, Spaces.CellFace()) stencil_interior_width(::DivergenceC2F, arg) = ((-half, half),) -Base.@propagate_inbounds function stencil_interior( - ::DivergenceC2F, - space, - idx, - hidx, - arg, -) - local_geometry = Geometry.LocalGeometry(space, idx, hidx) - Ju³₊ = Geometry.Jcontravariant3( - getidx(space, arg, idx + half, hidx), - Geometry.LocalGeometry(space, idx + half, hidx), - ) - Ju³₋ = Geometry.Jcontravariant3( - getidx(space, arg, idx - half, hidx), - Geometry.LocalGeometry(space, idx - half, hidx), - ) - (Ju³₊ - Ju³₋) * local_geometry.invJ -end boundary_width(::DivergenceC2F, ::AbstractBoundaryCondition) = 1 -# left / right SetDivergence boundary conditions -Base.@propagate_inbounds function stencil_left_boundary( - ::DivergenceC2F, - bc::SetDivergence, - space, - idx, - hidx, - arg, -) - @assert idx == left_face_boundary_idx(space) - # imposed flux boundary condition at left most face - getidx(space, bc.val, nothing, hidx) -end -Base.@propagate_inbounds function stencil_right_boundary( - ::DivergenceC2F, - bc::SetDivergence, - space, - idx, - hidx, - arg, -) - @assert idx == right_face_boundary_idx(space) - # imposed flux boundary condition at right most face - getidx(space, bc.val, nothing, hidx) -end abstract type CurlFiniteDifferenceOperator <: FiniteDifferenceOperator end @@ -2783,58 +2275,11 @@ end return_space(::CurlC2F, space::AllCenterFiniteDifferenceSpace) = Spaces.space(space, Spaces.CellFace()) -function fd3_curl(u₊::CV, u₋::CV, invJ) where {CV <: Geometry.CovariantVector} - Geometry.Contravariant123Vector( - -(u₊.u₂ - u₋.u₂) * invJ, - (u₊.u₁ - u₋.u₁) * invJ, - zero(eltype(invJ)), - ) -end stencil_interior_width(::CurlC2F, arg) = ((-half, half),) -Base.@propagate_inbounds function stencil_interior( - ::CurlC2F, - space, - idx, - hidx, - arg, -) - u₊ = getidx(space, arg, idx + half, hidx) - u₋ = getidx(space, arg, idx - half, hidx) - local_geometry = Geometry.LocalGeometry(space, idx, hidx) - return fd3_curl(u₊, u₋, local_geometry.invJ) -end boundary_width(::CurlC2F, ::AbstractBoundaryCondition) = 1 -# Project the user-supplied curl value onto the full Contravariant123 axis so -# the boundary output matches `return_eltype` (uniformly Contravariant123Vector). -Base.@propagate_inbounds function stencil_left_boundary( - ::CurlC2F, - bc::SetCurl, - space, - idx, - hidx, - arg, -) - return Geometry.project( - Geometry.Contravariant123Axis(), - getidx(space, bc.val, nothing, hidx), - ) -end -Base.@propagate_inbounds function stencil_right_boundary( - ::CurlC2F, - bc::SetCurl, - space, - idx, - hidx, - arg, -) - return Geometry.project( - Geometry.Contravariant123Axis(), - getidx(space, bc.val, nothing, hidx), - ) -end # code for figuring out boundary widths @@ -2980,10 +2425,13 @@ end @inline function should_call_left_boundary(idx, space, op, args...) Topologies.isperiodic(space) && return false loc = left_boundary_window(space) + boundary_condition = Operators.get_boundary(op, loc) + op isa SetBoundaryOperator && boundary_condition isa NullBoundaryCondition && + return false return idx < Operators.left_interior_idx( space, op, - Operators.get_boundary(op, loc), + boundary_condition, args..., ) end @@ -2991,6 +2439,9 @@ end @inline function should_call_right_boundary(idx, space, op, args...) Topologies.isperiodic(space) && return false loc = right_boundary_window(space) + boundary_condition = Operators.get_boundary(op, loc) + op isa SetBoundaryOperator && boundary_condition isa NullBoundaryCondition && + return false return idx > Operators.right_interior_idx( space, op, @@ -3029,7 +2480,6 @@ Base.@propagate_inbounds function getidx( return call_bc_f(bc.f, space, idx, hidx, args...) end op = bc.op - if should_call_left_boundary(idx, space, bc.op, bc.args...) stencil_left_boundary( op, diff --git a/test/Operators/finitedifference/unit_column.jl b/test/Operators/finitedifference/unit_column.jl index 30b996cfcd..787c02126c 100644 --- a/test/Operators/finitedifference/unit_column.jl +++ b/test/Operators/finitedifference/unit_column.jl @@ -167,8 +167,8 @@ end # 3) we set boundaries on both: 2nd should take precedence I = Operators.InterpolateC2F( - left = Operators.SetValue(FT(NaN)), - right = Operators.SetValue(FT(NaN)), + left = Operators.SetValue(FT(1)), + right = Operators.SetValue(FT(1)), ) ∂ = Operators.GradientF2C( left = Operators.SetValue(FT(0)),