Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor advection stencils to use StencilTest - Part 2 #351

Merged
merged 18 commits into from
Jan 25, 2024
Merged
Show file tree
Hide file tree
Changes from 15 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@

import sys

from gt4py.next.common import GridType
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where

Expand Down Expand Up @@ -665,7 +666,7 @@ def _divide_flux_area_list_stencil_01(
)


@program
@program(grid_type=GridType.UNSTRUCTURED)
def divide_flux_area_list_stencil_01(
famask_int: Field[[EdgeDim, KDim], int32],
p_vn: Field[[EdgeDim, KDim], float],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@

import sys

from gt4py.next.common import GridType
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where

Expand Down Expand Up @@ -174,7 +175,7 @@ def _divide_flux_area_list_stencil_02(
)


@program
@program(grid_type=GridType.UNSTRUCTURED)
def divide_flux_area_list_stencil_02(
famask_int: Field[[EdgeDim, KDim], int32],
p_vn: Field[[EdgeDim, KDim], float],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,9 +11,9 @@
#
# SPDX-License-Identifier: GPL-3.0-or-later

from gt4py.next.common import Field
from gt4py.next.common import GridType
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import int32, where
from gt4py.next.ffront.fbuiltins import Field, int32, where

from icon4py.model.common.dimension import E2C, CellDim, EdgeDim, KDim

Expand Down Expand Up @@ -108,7 +108,7 @@ def _hflux_ffsl_hybrid_stencil_01a(
return p_out_e_hybrid_1a


@program
@program(grid_type=GridType.UNSTRUCTURED)
def hflux_ffsl_hybrid_stencil_01a(
z_lsq_coeff_1: Field[[CellDim, KDim], float],
z_lsq_coeff_2: Field[[CellDim, KDim], float],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
#
# SPDX-License-Identifier: GPL-3.0-or-later

from gt4py.next.common import GridType
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import Field

Expand All @@ -28,7 +29,7 @@ def _hflux_ffsl_hybrid_stencil_02(
return p_out_e_hybrid_2


@program
@program(grid_type=GridType.UNSTRUCTURED)
def hflux_ffsl_hybrid_stencil_02(
p_out_e_hybrid_2: Field[[EdgeDim, KDim], float],
p_mass_flx_e: Field[[EdgeDim, KDim], float],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
#
# SPDX-License-Identifier: GPL-3.0-or-later

from gt4py.next.common import GridType
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import Field, int32, where

Expand Down Expand Up @@ -323,7 +324,7 @@ def _prep_gauss_quadrature_c_list_stencil(
)


@program
@program(grid_type=GridType.UNSTRUCTURED)
def prep_gauss_quadrature_c_list_stencil(
famask_int: Field[[EdgeDim, KDim], int32],
p_coords_dreg_v_1_x: Field[[EdgeDim, KDim], float],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
#
# SPDX-License-Identifier: GPL-3.0-or-later

from gt4py.next.common import GridType
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import ( # noqa: A004 # import gt4py builtin
Field,
Expand Down Expand Up @@ -293,7 +294,7 @@ def _prep_gauss_quadrature_c_stencil(
)


@program
@program(grid_type=GridType.UNSTRUCTURED)
def prep_gauss_quadrature_c_stencil(
p_coords_dreg_v_1_x: Field[[EdgeDim, KDim], float],
p_coords_dreg_v_2_x: Field[[EdgeDim, KDim], float],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
#
# SPDX-License-Identifier: GPL-3.0-or-later

from gt4py.next.common import Field
from gt4py.next.common import Field, GridType
from gt4py.next.ffront.decorator import field_operator, program

from icon4py.model.common.dimension import C2CECEC, C2E2C2E2C, CECECDim, CellDim, KDim
Expand Down Expand Up @@ -176,7 +176,7 @@ def _recon_lsq_cell_c_svd_stencil(
)


@program
@program(grid_type=GridType.UNSTRUCTURED)
def recon_lsq_cell_c_svd_stencil(
p_cc: Field[[CellDim, KDim], float],
lsq_pseudoinv_1: Field[[CECECDim], float],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
#
# SPDX-License-Identifier: GPL-3.0-or-later

from gt4py.next.common import Field
from gt4py.next.common import Field, GridType
from gt4py.next.ffront.decorator import field_operator, program

from icon4py.model.common.dimension import C2CEC, C2E2C, CECDim, CellDim, KDim
Expand Down Expand Up @@ -41,7 +41,7 @@ def _recon_lsq_cell_l_svd_stencil(
return p_coeff_1_dsl, p_coeff_2_dsl, p_coeff_3_dsl


@program
@program(grid_type=GridType.UNSTRUCTURED)
def recon_lsq_cell_l_svd_stencil(
p_cc: Field[[CellDim, KDim], float],
lsq_pseudoinv_1: Field[[CECDim], float],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
#
# SPDX-License-Identifier: GPL-3.0-or-later

from gt4py.next.common import Field
from gt4py.next.common import Field, GridType
from gt4py.next.ffront.decorator import field_operator, program

from icon4py.model.common.dimension import CellDim, KDim, Koff
Expand All @@ -31,7 +31,7 @@ def _step_advection_stencil_01(
return rhodz_ast + k_offset_up_low


@program
@program(grid_type=GridType.UNSTRUCTURED)
def step_advection_stencil_01(
rhodz_ast: Field[[CellDim, KDim], float],
p_mflx_contra_v: Field[[CellDim, KDim], float],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
#
# SPDX-License-Identifier: GPL-3.0-or-later

from gt4py.next.common import Field
from gt4py.next.common import Field, GridType
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import maximum

Expand All @@ -31,7 +31,7 @@ def _step_advection_stencil_02(
)


@program
@program(grid_type=GridType.UNSTRUCTURED)
def step_advection_stencil_02(
p_rhodz_new: Field[[CellDim, KDim], float],
p_mflx_contra_v: Field[[CellDim, KDim], float],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
#
# SPDX-License-Identifier: GPL-3.0-or-later

from gt4py.next.common import GridType
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import Field, maximum

Expand All @@ -27,7 +28,7 @@ def _step_advection_stencil_03(
return p_tracer_new


@program
@program(grid_type=GridType.UNSTRUCTURED)
def step_advection_stencil_03(
p_tracer_now: Field[[CellDim, KDim], float],
p_grf_tend_tracer: Field[[CellDim, KDim], float],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
#
# SPDX-License-Identifier: GPL-3.0-or-later

from gt4py.next.common import GridType
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import Field

Expand All @@ -27,7 +28,7 @@ def _step_advection_stencil_04(
return opt_ddt_tracer_adv


@program
@program(grid_type=GridType.UNSTRUCTURED)
def step_advection_stencil_04(
p_tracer_now: Field[[CellDim, KDim], float],
p_tracer_new: Field[[CellDim, KDim], float],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
#
# SPDX-License-Identifier: GPL-3.0-or-later

from gt4py.next.common import Field
from gt4py.next.common import Field, GridType
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import int32, where

Expand Down Expand Up @@ -114,7 +114,7 @@ def _upwind_hflux_miura3_stencil_01(
return p_out_e_miura3


@program
@program(grid_type=GridType.UNSTRUCTURED)
def upwind_hflux_miura3_stencil_01(
z_lsq_coeff_1: Field[[CellDim, KDim], float],
z_lsq_coeff_2: Field[[CellDim, KDim], float],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
#
# SPDX-License-Identifier: GPL-3.0-or-later

from gt4py.next.common import Field
from gt4py.next.common import Field, GridType
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import int32, where

Expand Down Expand Up @@ -51,7 +51,7 @@ def _upwind_hflux_miura_cycl_stencil_01(
return z_tracer_mflx_dsl


@program
@program(grid_type=GridType.UNSTRUCTURED)
def upwind_hflux_miura_cycl_stencil_01(
z_lsq_coeff_1_dsl: Field[[CellDim, KDim], float],
z_lsq_coeff_2_dsl: Field[[CellDim, KDim], float],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
#
# SPDX-License-Identifier: GPL-3.0-or-later

from gt4py.next.common import GridType
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import Field, int32, neighbor_sum

Expand Down Expand Up @@ -48,7 +49,7 @@ def _upwind_hflux_miura_cycl_stencil_02(
return (z_rhofluxdiv_c_out, z_fluxdiv_c_dsl, z_rho_new_dsl, z_tracer_new_dsl)


@program
@program(grid_type=GridType.UNSTRUCTURED)
def upwind_hflux_miura_cycl_stencil_02(
nsub: int32,
p_mass_flx_e: Field[[EdgeDim, KDim], float],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
#
# SPDX-License-Identifier: GPL-3.0-or-later

from gt4py.next.common import Field
from gt4py.next.common import Field, GridType
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import float64

Expand All @@ -27,7 +27,7 @@ def _upwind_hflux_miura_cycl_stencil_03a(
return p_out_e


@program
@program(grid_type=GridType.UNSTRUCTURED)
def upwind_hflux_miura_cycl_stencil_03a(
z_tracer_mflx_1_dsl: Field[[EdgeDim, KDim], float],
z_tracer_mflx_2_dsl: Field[[EdgeDim, KDim], float],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
#
# SPDX-License-Identifier: GPL-3.0-or-later

from gt4py.next.common import Field
from gt4py.next.common import Field, GridType
from gt4py.next.ffront.decorator import field_operator, program

from icon4py.model.common.dimension import EdgeDim, KDim
Expand All @@ -27,7 +27,7 @@ def _upwind_hflux_miura_cycl_stencil_03b(
return p_out_e


@program
@program(grid_type=GridType.UNSTRUCTURED)
def upwind_hflux_miura_cycl_stencil_03b(
z_tracer_mflx_1_dsl: Field[[EdgeDim, KDim], float],
z_tracer_mflx_2_dsl: Field[[EdgeDim, KDim], float],
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
#
# SPDX-License-Identifier: GPL-3.0-or-later

from gt4py.next.common import Field
from gt4py.next.common import Field, GridType
from gt4py.next.ffront.decorator import field_operator, program
from gt4py.next.ffront.fbuiltins import int32, where

Expand Down Expand Up @@ -39,7 +39,7 @@ def _upwind_hflux_miura_stencil_01(
return p_out_e


@program
@program(grid_type=GridType.UNSTRUCTURED)
def upwind_hflux_miura_stencil_01(
z_lsq_coeff_1: Field[[CellDim, KDim], float],
z_lsq_coeff_2: Field[[CellDim, KDim], float],
Expand Down
Loading