From 1e00c05fdaf6ad6153475f29f6109ec663893251 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Thu, 24 Nov 2022 09:04:58 +0100 Subject: [PATCH 01/49] First version of fused diffusion stencils --- .../fused_mo_nh_diffusion_stencil_02_03.py | 47 ++++++++++ .../fused_mo_nh_diffusion_stencil_04_05_06.py | 82 ++++++++++++++++++ ...sed_mo_nh_diffusion_stencil_07_08_09_10.py | 86 +++++++++++++++++++ .../fused_mo_nh_diffusion_stencil_11_12.py | 40 +++++++++ .../fused_mo_nh_diffusion_stencil_13_14.py | 67 +++++++++++++++ .../mo_nh_diffusion_stencil_10.py | 8 +- 6 files changed, 326 insertions(+), 4 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py new file mode 100644 index 0000000000..8c111b5d4d --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py @@ -0,0 +1,47 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, neighbor_sum + +from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim, Koff + +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import _mo_nh_diffusion_stencil_02 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import _mo_nh_diffusion_stencil_03 + +@field_operator +def _fused_mo_nh_diffusion_stencil_02_03( + kh_smag_ec: Field[[EdgeDim, KDim], float], + vn: Field[[EdgeDim, KDim], float], + e_bln_c_s: Field[[CellDim, C2EDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + diff_multfac_smag: Field[[KDim], float], + wgtfac_c: Field[[CellDim, KDim], float], +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + kh_c, div = _mo_nh_diffusion_stencil_02(kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag) + div_ic, hdef_ic = _mo_nh_diffusion_stencil_03(div, kh_c, wgtfac_c) + return div_ic, hdef_ic + + +@program +def fused_mo_nh_diffusion_stencil_02_03( + kh_smag_ec: Field[[EdgeDim, KDim], float], + vn: Field[[EdgeDim, KDim], float], + e_bln_c_s: Field[[CellDim, C2EDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + diff_multfac_smag: Field[[KDim], float], + wgtfac_c: Field[[CellDim, KDim], float], + div_ic: Field[[CellDim, KDim], float], + hdef_ic: Field[[CellDim, KDim], float], +): + _fused_mo_nh_diffusion_stencil_02_03(kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag, wgtfac_c, out=(div_ic, hdef_ic)) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py new file mode 100644 index 0000000000..1438e07c10 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py @@ -0,0 +1,82 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, neighbor_sum, maximum, where, int32 + +from icon4py.common.dimension import ( + E2C2V, + E2ECV, + ECVDim, + EdgeDim, + KDim, + VertexDim, +) + +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_04 import _mo_nh_diffusion_stencil_04 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_05 import _mo_nh_diffusion_stencil_05 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_06 import _mo_nh_diffusion_stencil_06 + +@field_operator +def _fused_mo_nh_diffusion_stencil_04_05_06( + u_vert: Field[[VertexDim, KDim], float], + v_vert: Field[[VertexDim, KDim], float], + primal_normal_vert_v1: Field[[ECVDim], float], + primal_normal_vert_v2: Field[[ECVDim], float], + z_nabla2_e: Field[[EdgeDim, KDim], float], + inv_vert_vert_length: Field[[EdgeDim], float], + inv_primal_edge_length: Field[[EdgeDim], float], + area_edge: Field[[EdgeDim], float], + kh_smag_e: Field[[EdgeDim, KDim], float], + diff_multfac_vn: Field[[KDim], float], + nudgecoeff_e: Field[[EdgeDim], float], + vn: Field[[EdgeDim, KDim], float], + horz_idx: Field[[EdgeDim], int32], + nudgezone_diff: float, + fac_bdydiff_v: float, + start_2nd_nudge_line_idx_e: int32, +) -> Field[[EdgeDim, KDim], float]: + + z_nabla4_e2 = _mo_nh_diffusion_stencil_04(u_vert, v_vert, primal_normal_vert_v1, primal_normal_vert_v2, z_nabla2_e, inv_vert_vert_length, inv_primal_edge_length) + + vn = where( + horz_idx >= start_2nd_nudge_line_idx_e, + _mo_nh_diffusion_stencil_05(area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, nudgecoeff_e, vn, nudgezone_diff), + _mo_nh_diffusion_stencil_06(z_nabla2_e, area_edge, vn, fac_bdydiff_v) + ) + + return vn + + +@program +def fused_mo_nh_diffusion_stencil_04_05_06( + u_vert: Field[[VertexDim, KDim], float], + v_vert: Field[[VertexDim, KDim], float], + primal_normal_vert_v1: Field[[ECVDim], float], + primal_normal_vert_v2: Field[[ECVDim], float], + z_nabla2_e: Field[[EdgeDim, KDim], float], + inv_vert_vert_length: Field[[EdgeDim], float], + inv_primal_edge_length: Field[[EdgeDim], float], + area_edge: Field[[EdgeDim], float], + kh_smag_e: Field[[EdgeDim, KDim], float], + diff_multfac_vn: Field[[KDim], float], + nudgecoeff_e: Field[[EdgeDim], float], + vn: Field[[EdgeDim, KDim], float], + horz_idx: Field[[EdgeDim], int32], + nudgezone_diff: float, + fac_bdydiff_v: float, + start_2nd_nudge_line_idx_e: int32, +): + _fused_mo_nh_diffusion_stencil_04_05_06(u_vert, v_vert, primal_normal_vert_v1, primal_normal_vert_v2, z_nabla2_e, inv_vert_vert_length, + inv_primal_edge_length, area_edge, kh_smag_e, diff_multfac_vn, nudgecoeff_e, vn, horz_idx, nudgezone_diff, fac_bdydiff_v, + start_2nd_nudge_line_idx_e, out=vn) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py new file mode 100644 index 0000000000..a1e305ae55 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py @@ -0,0 +1,86 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, neighbor_sum, where, int32, broadcast + +from icon4py.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim + +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_07 import _mo_nh_diffusion_stencil_07 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_08 import _mo_nh_diffusion_stencil_08 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_09 import _mo_nh_diffusion_stencil_09 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_10 import _mo_nh_diffusion_stencil_10 + +@field_operator +def _fused_mo_nh_diffusion_stencil_07_08_09_10( + area: Field[[CellDim], float], + geofac_n2s: Field[[CellDim, C2E2CODim], float], + geofac_grg_x: Field[[CellDim, C2E2CODim], float], + geofac_grg_y: Field[[CellDim, C2E2CODim], float], + w_old: Field[[CellDim, KDim], float], + w: Field[[CellDim, KDim], float], + dwdx: Field[[CellDim, KDim], float], + dwdy: Field[[CellDim, KDim], float], + diff_multfac_w: float, + diff_multfac_n2w: Field[[KDim], float], + vert_idx: Field[[KDim], int32], + horz_idx: Field[[CellDim], int32], + nrdmax: int32, + interior_idx: int32, + halo_idx: int32, +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + + vert_idx = broadcast(vert_idx, (CellDim, KDim)) + + dwdx, dwdy = where( + vert_idx > int32(0), + _mo_nh_diffusion_stencil_08(w_old, geofac_grg_x, geofac_grg_y), + (dwdx, dwdy) + ) + + z_nabla2_c = _mo_nh_diffusion_stencil_07(w_old, geofac_n2s) + + w = where( + (horz_idx >= interior_idx) & (horz_idx < halo_idx), + _mo_nh_diffusion_stencil_09(area, z_nabla2_c, geofac_n2s, w_old, diff_multfac_w), + w_old, + ) + + w = where( + (vert_idx > int32(0)) & (vert_idx < nrdmax) & (horz_idx >= interior_idx) & (horz_idx < halo_idx), + _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, area, z_nabla2_c), + w, + ) + + return w, dwdx, dwdy + + +@program +def fused_mo_nh_diffusion_stencil_07_08_09_10( + area: Field[[CellDim], float], + geofac_n2s: Field[[CellDim, C2E2CODim], float], + geofac_grg_x: Field[[CellDim, C2E2CODim], float], + geofac_grg_y: Field[[CellDim, C2E2CODim], float], + w_old: Field[[CellDim, KDim], float], + w: Field[[CellDim, KDim], float], + dwdx: Field[[CellDim, KDim], float], + dwdy: Field[[CellDim, KDim], float], + diff_multfac_w: float, + diff_multfac_n2w: Field[[KDim], float], + vert_idx: Field[[KDim], int32], + horz_idx: Field[[CellDim], int32], + nrdmax: int32, + interior_idx: int32, + halo_idx: int32, +): + _fused_mo_nh_diffusion_stencil_07_08_09_10(area, geofac_n2s, geofac_grg_x, geofac_grg_y, w_old, w, dwdx, dwdy, diff_multfac_w, diff_multfac_n2w, vert_idx, horz_idx, nrdmax, interior_idx, halo_idx, out = (w, dwdx, dwdy)) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py new file mode 100644 index 0000000000..02fd005304 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py @@ -0,0 +1,40 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, where, neighbor_sum, max_over, maximum + +from icon4py.common.dimension import E2C, E2CDim, C2E2C, C2E2CDim, CellDim, EdgeDim, KDim + +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_11 import _mo_nh_diffusion_stencil_11 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_12 import _mo_nh_diffusion_stencil_12 + +@field_operator +def _fused_mo_nh_diffusion_stencil_11_12( + theta_v: Field[[CellDim, KDim], float], + theta_ref_mc: Field[[CellDim, KDim], float], + thresh_tdiff: float, + kh_smag_e: Field[[EdgeDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + enh_diffu_3d = _mo_nh_diffusion_stencil_11(theta_v, theta_ref_mc, thresh_tdiff) + kh_smag_e = _mo_nh_diffusion_stencil_12(kh_smag_e, enh_diffu_3d) + return kh_smag_e + +@program +def fused_mo_nh_diffusion_stencil_11_12( + theta_v: Field[[CellDim, KDim], float], + theta_ref_mc: Field[[CellDim, KDim], float], + thresh_tdiff: float, + kh_smag_e: Field[[EdgeDim, KDim], float], +): + _fused_mo_nh_diffusion_stencil_11_12(theta_v, theta_ref_mc, thresh_tdiff, kh_smag_e, out=kh_smag_e) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py new file mode 100644 index 0000000000..40faa10b54 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py @@ -0,0 +1,67 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field, neighbor_sum + +from icon4py.common.dimension import ( + E2C, + E2CDim, + C2E2C, + C2E2CDim, + C2CE, + C2E, + C2EDim, + CEDim, + CellDim, + EdgeDim, + KDim, +) + +@field_operator +def _mo_nh_diffusion_stencil_13( + kh_smag_e: Field[[EdgeDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], + theta_v: Field[[CellDim, KDim], float], +) -> Field[[EdgeDim, KDim], float]: + z_nabla2_e = kh_smag_e * inv_dual_edge_length * (theta_v(E2C[1]) - theta_v(E2C[0])) + return z_nabla2_e + +@field_operator +def _mo_nh_diffusion_stencil_14( + z_nabla2_e: Field[[EdgeDim, KDim], float], + geofac_div: Field[[CEDim], float], +) -> Field[[CellDim, KDim], float]: + z_temp = neighbor_sum(z_nabla2_e(C2E) * geofac_div(C2CE), axis=C2EDim) + return z_temp + +@field_operator +def _fused_mo_nh_diffusion_stencil_13_14( + kh_smag_e: Field[[EdgeDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], + theta_v: Field[[CellDim, KDim], float], + geofac_div: Field[[CEDim], float], +) -> Field[[CellDim, KDim], float]: + z_nabla2_e = _mo_nh_diffusion_stencil_13(kh_smag_e, inv_dual_edge_length, theta_v) + z_temp = _mo_nh_diffusion_stencil_14(z_nabla2_e, geofac_div) + return z_temp + +@program +def fused_mo_nh_diffusion_stencil_13_14( + kh_smag_e: Field[[EdgeDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], + theta_v: Field[[CellDim, KDim], float], + geofac_div: Field[[CEDim], float], + z_temp: Field[[CellDim, KDim], float], +): + _fused_mo_nh_diffusion_stencil_13_14(kh_smag_e, inv_dual_edge_length, theta_v, geofac_div, out=z_temp) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py index e439f742ea..fed277e852 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py @@ -21,10 +21,10 @@ def _mo_nh_diffusion_stencil_10( w: Field[[CellDim, KDim], float], diff_multfac_n2w: Field[[KDim], float], - cell_area: Field[[CellDim], float], + area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - w = w + diff_multfac_n2w * (cell_area * z_nabla2_c) + w = w + diff_multfac_n2w * (area * z_nabla2_c) return w @@ -32,7 +32,7 @@ def _mo_nh_diffusion_stencil_10( def mo_nh_diffusion_stencil_10( w: Field[[CellDim, KDim], float], diff_multfac_n2w: Field[[KDim], float], - cell_area: Field[[CellDim], float], + area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], ): - _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, cell_area, z_nabla2_c, out=w) + _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, area, z_nabla2_c, out=w) From 2abd8b93738c6193255b9b42b3c242b72474f7b4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Tue, 29 Nov 2022 12:30:45 +0100 Subject: [PATCH 02/49] Added prototype of fused_mo_nh_diffusion_stencil_01_02_03_rbf Currently this stencil does not work since it would require elements of a tuple of output variables of a field_operator to have different types. --- ...ed_mo_nh_diffusion_stencil_01_02_03_rbf.py | 92 +++++++++++++++++++ 1 file changed, 92 insertions(+) create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py new file mode 100644 index 0000000000..ae2b30777c --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py @@ -0,0 +1,92 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from functional.ffront.decorator import field_operator, program +from functional.ffront.fbuiltins import Field + +from icon4py.common.dimension import ECVDim, C2EDim, V2EDim, CellDim, EdgeDim, VertexDim, KDim + +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_01 import _mo_nh_diffusion_stencil_01 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import _mo_nh_diffusion_stencil_02 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import _mo_nh_diffusion_stencil_03 +from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import _mo_intp_rbf_rbf_vec_interpol_vertex + +@field_operator +def _fused_mo_nh_diffusion_stencil_01_02_03_rbf( + diff_multfac_smag: Field[[KDim], float], + tangent_orientation: Field[[EdgeDim], float], + inv_primal_edge_length: Field[[EdgeDim], float], + inv_vert_vert_length: Field[[EdgeDim], float], + u_vert_old: Field[[VertexDim, KDim], float], + v_vert_old: Field[[VertexDim, KDim], float], + primal_normal_vert_x: Field[[ECVDim], float], + primal_normal_vert_y: Field[[ECVDim], float], + dual_normal_vert_x: Field[[ECVDim], float], + dual_normal_vert_y: Field[[ECVDim], float], + vn: Field[[EdgeDim, KDim], float], + smag_limit: Field[[KDim], float], + smag_offset: float, + e_bln_c_s: Field[[CellDim, C2EDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + wgtfac_c: Field[[CellDim, KDim], float], + ptr_coeff_1: Field[[VertexDim, V2EDim], float], + ptr_coeff_2: Field[[VertexDim, V2EDim], float], +) -> tuple[ + Field[[EdgeDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[VertexDim, KDim], float], + Field[[VertexDim, KDim], float], + ]: + + kh_smag_e, kh_smag_ec, z_nabla2_e = _mo_nh_diffusion_stencil_01( diff_multfac_smag, + tangent_orientation, inv_primal_edge_length, inv_vert_vert_length, u_vert_old, v_vert_old, + primal_normal_vert_x, primal_normal_vert_y, dual_normal_vert_x, dual_normal_vert_y, + vn, smag_limit, smag_offset) + + kh_c, div = _mo_nh_diffusion_stencil_02(kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag) + + div_ic, hdef_ic = _mo_nh_diffusion_stencil_03(div, kh_c, wgtfac_c) + + u_vert, v_vert = _mo_intp_rbf_rbf_vec_interpol_vertex(z_nabla2_e, ptr_coeff_1, ptr_coeff_2) + + return kh_smag_e, div_ic, hdef_ic, u_vert, v_vert + + +@program +def fused_mo_nh_diffusion_stencil_01_02_03_rbf( + diff_multfac_smag: Field[[KDim], float], + tangent_orientation: Field[[EdgeDim], float], + inv_primal_edge_length: Field[[EdgeDim], float], + inv_vert_vert_length: Field[[EdgeDim], float], + u_vert_old: Field[[VertexDim, KDim], float], + v_vert_old: Field[[VertexDim, KDim], float], + primal_normal_vert_x: Field[[ECVDim], float], + primal_normal_vert_y: Field[[ECVDim], float], + dual_normal_vert_x: Field[[ECVDim], float], + dual_normal_vert_y: Field[[ECVDim], float], + vn: Field[[EdgeDim, KDim], float], + smag_limit: Field[[KDim], float], + smag_offset: float, + e_bln_c_s: Field[[CellDim, C2EDim], float], + geofac_div: Field[[CellDim, C2EDim], float], + wgtfac_c: Field[[CellDim, KDim], float], + ptr_coeff_1: Field[[VertexDim, V2EDim], float], + ptr_coeff_2: Field[[VertexDim, V2EDim], float], + kh_smag_e: Field[[EdgeDim, KDim], float], + div_ic: Field[[CellDim, KDim], float], + hdef_ic: Field[[CellDim, KDim], float], + u_vert: Field[[VertexDim, KDim], float], + v_vert: Field[[VertexDim, KDim], float], +): + _fused_mo_nh_diffusion_stencil_01_02_03_rbf(diff_multfac_smag, tangent_orientation, inv_primal_edge_length, inv_vert_vert_length, u_vert_old, v_vert_old, primal_normal_vert_x, primal_normal_vert_y, dual_normal_vert_x, dual_normal_vert_y, vn, smag_limit, smag_offset, e_bln_c_s, geofac_div, wgtfac_c, ptr_coeff_1, ptr_coeff_2, out=(kh_smag_e, div_ic, hdef_ic, u_vert, v_vert)) From ee4e21e7ce44e43212cecb61cc27ed89b4c065e2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Wed, 7 Dec 2022 15:19:32 +0100 Subject: [PATCH 03/49] Removed duplicate stencil definition from fused stencil 13 and 14 --- ...used_mo_nh_diffusion_stencil_07_08_09_10.py | 2 +- .../fused_mo_nh_diffusion_stencil_13_14.py | 18 ++---------------- 2 files changed, 3 insertions(+), 17 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py index a1e305ae55..a97a75d21f 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py @@ -40,7 +40,7 @@ def _fused_mo_nh_diffusion_stencil_07_08_09_10( halo_idx: int32, ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - vert_idx = broadcast(vert_idx, (CellDim, KDim)) + vert_idx = broadcast(vert_idx, (CellDim, KDim)) dwdx, dwdy = where( vert_idx > int32(0), diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py index 40faa10b54..8d20920085 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py @@ -28,22 +28,8 @@ KDim, ) -@field_operator -def _mo_nh_diffusion_stencil_13( - kh_smag_e: Field[[EdgeDim, KDim], float], - inv_dual_edge_length: Field[[EdgeDim], float], - theta_v: Field[[CellDim, KDim], float], -) -> Field[[EdgeDim, KDim], float]: - z_nabla2_e = kh_smag_e * inv_dual_edge_length * (theta_v(E2C[1]) - theta_v(E2C[0])) - return z_nabla2_e - -@field_operator -def _mo_nh_diffusion_stencil_14( - z_nabla2_e: Field[[EdgeDim, KDim], float], - geofac_div: Field[[CEDim], float], -) -> Field[[CellDim, KDim], float]: - z_temp = neighbor_sum(z_nabla2_e(C2E) * geofac_div(C2CE), axis=C2EDim) - return z_temp +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_13 import _mo_nh_diffusion_stencil_13 +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_14 import _mo_nh_diffusion_stencil_14 @field_operator def _fused_mo_nh_diffusion_stencil_13_14( From 381be77f2e55c5f1138929b8623a7b0d2e98d060 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Fri, 20 Jan 2023 14:07:29 +0100 Subject: [PATCH 04/49] Cleanup --- .../fused_mo_nh_diffusion_stencil_02_03.py | 2 +- .../fused_mo_nh_diffusion_stencil_04_05_06.py | 11 ++--------- .../fused_mo_nh_diffusion_stencil_07_08_09_10.py | 4 ++-- .../fused_mo_nh_diffusion_stencil_11_12.py | 4 ++-- .../fused_mo_nh_diffusion_stencil_13_14.py | 16 ++-------------- .../atm_dyn_iconam/mo_nh_diffusion_stencil_10.py | 8 ++++---- 6 files changed, 13 insertions(+), 32 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py index 8c111b5d4d..f034e5a62f 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py @@ -14,7 +14,7 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field, neighbor_sum -from icon4py.common.dimension import C2E, C2EDim, CellDim, EdgeDim, KDim, Koff +from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import _mo_nh_diffusion_stencil_02 from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import _mo_nh_diffusion_stencil_03 diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py index 1438e07c10..b74c52b4e5 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py @@ -12,16 +12,9 @@ # SPDX-License-Identifier: GPL-3.0-or-later from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, neighbor_sum, maximum, where, int32 +from functional.ffront.fbuiltins import Field, where, int32 -from icon4py.common.dimension import ( - E2C2V, - E2ECV, - ECVDim, - EdgeDim, - KDim, - VertexDim, -) +from icon4py.common.dimension import ECVDim, EdgeDim, KDim, VertexDim from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_04 import _mo_nh_diffusion_stencil_04 from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_05 import _mo_nh_diffusion_stencil_05 diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py index a97a75d21f..0bbcce058a 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py @@ -12,9 +12,9 @@ # SPDX-License-Identifier: GPL-3.0-or-later from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, neighbor_sum, where, int32, broadcast +from functional.ffront.fbuiltins import Field, where, int32, broadcast -from icon4py.common.dimension import C2E2CO, C2E2CODim, CellDim, KDim +from icon4py.common.dimension import C2E2CODim, CellDim, KDim from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_07 import _mo_nh_diffusion_stencil_07 from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_08 import _mo_nh_diffusion_stencil_08 diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py index 02fd005304..cda940f529 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py @@ -12,9 +12,9 @@ # SPDX-License-Identifier: GPL-3.0-or-later from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, where, neighbor_sum, max_over, maximum +from functional.ffront.fbuiltins import Field -from icon4py.common.dimension import E2C, E2CDim, C2E2C, C2E2CDim, CellDim, EdgeDim, KDim +from icon4py.common.dimension import CellDim, EdgeDim, KDim from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_11 import _mo_nh_diffusion_stencil_11 from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_12 import _mo_nh_diffusion_stencil_12 diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py index 8d20920085..ae317277c1 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py @@ -12,21 +12,9 @@ # SPDX-License-Identifier: GPL-3.0-or-later from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, neighbor_sum +from functional.ffront.fbuiltins import Field -from icon4py.common.dimension import ( - E2C, - E2CDim, - C2E2C, - C2E2CDim, - C2CE, - C2E, - C2EDim, - CEDim, - CellDim, - EdgeDim, - KDim, -) +from icon4py.common.dimension import CEDim, CellDim, EdgeDim, KDim from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_13 import _mo_nh_diffusion_stencil_13 from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_14 import _mo_nh_diffusion_stencil_14 diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py index fed277e852..e439f742ea 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_10.py @@ -21,10 +21,10 @@ def _mo_nh_diffusion_stencil_10( w: Field[[CellDim, KDim], float], diff_multfac_n2w: Field[[KDim], float], - area: Field[[CellDim], float], + cell_area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], ) -> Field[[CellDim, KDim], float]: - w = w + diff_multfac_n2w * (area * z_nabla2_c) + w = w + diff_multfac_n2w * (cell_area * z_nabla2_c) return w @@ -32,7 +32,7 @@ def _mo_nh_diffusion_stencil_10( def mo_nh_diffusion_stencil_10( w: Field[[CellDim, KDim], float], diff_multfac_n2w: Field[[KDim], float], - area: Field[[CellDim], float], + cell_area: Field[[CellDim], float], z_nabla2_c: Field[[CellDim, KDim], float], ): - _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, area, z_nabla2_c, out=w) + _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, cell_area, z_nabla2_c, out=w) From f5ad668dd24bdb55aa338fc6dfe50822bdb482b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Mon, 23 Jan 2023 01:26:50 +0100 Subject: [PATCH 05/49] Cleanup --- ...ed_mo_nh_diffusion_stencil_01_02_03_rbf.py | 87 +++++++++++++++---- .../fused_mo_nh_diffusion_stencil_02_03.py | 22 ++++- .../fused_mo_nh_diffusion_stencil_04_05_06.py | 59 ++++++++++--- ...sed_mo_nh_diffusion_stencil_07_08_09_10.py | 54 +++++++++--- .../fused_mo_nh_diffusion_stencil_11_12.py | 13 ++- .../fused_mo_nh_diffusion_stencil_13_14.py | 13 ++- 6 files changed, 200 insertions(+), 48 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py index ae2b30777c..fb3ecaf12d 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py @@ -14,12 +14,28 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field -from icon4py.common.dimension import ECVDim, C2EDim, V2EDim, CellDim, EdgeDim, VertexDim, KDim +from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( + _mo_intp_rbf_rbf_vec_interpol_vertex, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_01 import ( + _mo_nh_diffusion_stencil_01, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import ( + _mo_nh_diffusion_stencil_02, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import ( + _mo_nh_diffusion_stencil_03, +) +from icon4py.common.dimension import ( + C2EDim, + CellDim, + ECVDim, + EdgeDim, + KDim, + V2EDim, + VertexDim, +) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_01 import _mo_nh_diffusion_stencil_01 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import _mo_nh_diffusion_stencil_02 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import _mo_nh_diffusion_stencil_03 -from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import _mo_intp_rbf_rbf_vec_interpol_vertex @field_operator def _fused_mo_nh_diffusion_stencil_01_02_03_rbf( @@ -42,23 +58,38 @@ def _fused_mo_nh_diffusion_stencil_01_02_03_rbf( ptr_coeff_1: Field[[VertexDim, V2EDim], float], ptr_coeff_2: Field[[VertexDim, V2EDim], float], ) -> tuple[ - Field[[EdgeDim, KDim], float], - Field[[CellDim, KDim], float], - Field[[CellDim, KDim], float], - Field[[VertexDim, KDim], float], - Field[[VertexDim, KDim], float], - ]: + Field[[EdgeDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[VertexDim, KDim], float], + Field[[VertexDim, KDim], float], +]: - kh_smag_e, kh_smag_ec, z_nabla2_e = _mo_nh_diffusion_stencil_01( diff_multfac_smag, - tangent_orientation, inv_primal_edge_length, inv_vert_vert_length, u_vert_old, v_vert_old, - primal_normal_vert_x, primal_normal_vert_y, dual_normal_vert_x, dual_normal_vert_y, - vn, smag_limit, smag_offset) + kh_smag_e, kh_smag_ec, z_nabla2_e = _mo_nh_diffusion_stencil_01( + diff_multfac_smag, + tangent_orientation, + inv_primal_edge_length, + inv_vert_vert_length, + u_vert_old, + v_vert_old, + primal_normal_vert_x, + primal_normal_vert_y, + dual_normal_vert_x, + dual_normal_vert_y, + vn, + smag_limit, + smag_offset, + ) - kh_c, div = _mo_nh_diffusion_stencil_02(kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag) + kh_c, div = _mo_nh_diffusion_stencil_02( + kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag + ) div_ic, hdef_ic = _mo_nh_diffusion_stencil_03(div, kh_c, wgtfac_c) - u_vert, v_vert = _mo_intp_rbf_rbf_vec_interpol_vertex(z_nabla2_e, ptr_coeff_1, ptr_coeff_2) + u_vert, v_vert = _mo_intp_rbf_rbf_vec_interpol_vertex( + z_nabla2_e, ptr_coeff_1, ptr_coeff_2 + ) return kh_smag_e, div_ic, hdef_ic, u_vert, v_vert @@ -89,4 +120,24 @@ def fused_mo_nh_diffusion_stencil_01_02_03_rbf( u_vert: Field[[VertexDim, KDim], float], v_vert: Field[[VertexDim, KDim], float], ): - _fused_mo_nh_diffusion_stencil_01_02_03_rbf(diff_multfac_smag, tangent_orientation, inv_primal_edge_length, inv_vert_vert_length, u_vert_old, v_vert_old, primal_normal_vert_x, primal_normal_vert_y, dual_normal_vert_x, dual_normal_vert_y, vn, smag_limit, smag_offset, e_bln_c_s, geofac_div, wgtfac_c, ptr_coeff_1, ptr_coeff_2, out=(kh_smag_e, div_ic, hdef_ic, u_vert, v_vert)) + _fused_mo_nh_diffusion_stencil_01_02_03_rbf( + diff_multfac_smag, + tangent_orientation, + inv_primal_edge_length, + inv_vert_vert_length, + u_vert_old, + v_vert_old, + primal_normal_vert_x, + primal_normal_vert_y, + dual_normal_vert_x, + dual_normal_vert_y, + vn, + smag_limit, + smag_offset, + e_bln_c_s, + geofac_div, + wgtfac_c, + ptr_coeff_1, + ptr_coeff_2, + out=(kh_smag_e, div_ic, hdef_ic, u_vert, v_vert), + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py index f034e5a62f..c3a392b226 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py @@ -14,10 +14,14 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field, neighbor_sum +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import ( + _mo_nh_diffusion_stencil_02, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import ( + _mo_nh_diffusion_stencil_03, +) from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import _mo_nh_diffusion_stencil_02 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import _mo_nh_diffusion_stencil_03 @field_operator def _fused_mo_nh_diffusion_stencil_02_03( @@ -28,7 +32,9 @@ def _fused_mo_nh_diffusion_stencil_02_03( diff_multfac_smag: Field[[KDim], float], wgtfac_c: Field[[CellDim, KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - kh_c, div = _mo_nh_diffusion_stencil_02(kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag) + kh_c, div = _mo_nh_diffusion_stencil_02( + kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag + ) div_ic, hdef_ic = _mo_nh_diffusion_stencil_03(div, kh_c, wgtfac_c) return div_ic, hdef_ic @@ -44,4 +50,12 @@ def fused_mo_nh_diffusion_stencil_02_03( div_ic: Field[[CellDim, KDim], float], hdef_ic: Field[[CellDim, KDim], float], ): - _fused_mo_nh_diffusion_stencil_02_03(kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag, wgtfac_c, out=(div_ic, hdef_ic)) + _fused_mo_nh_diffusion_stencil_02_03( + kh_smag_ec, + vn, + e_bln_c_s, + geofac_div, + diff_multfac_smag, + wgtfac_c, + out=(div_ic, hdef_ic), + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py index b74c52b4e5..58e3345016 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py @@ -12,13 +12,19 @@ # SPDX-License-Identifier: GPL-3.0-or-later from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, where, int32 +from functional.ffront.fbuiltins import Field, int32, where +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_04 import ( + _mo_nh_diffusion_stencil_04, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_05 import ( + _mo_nh_diffusion_stencil_05, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_06 import ( + _mo_nh_diffusion_stencil_06, +) from icon4py.common.dimension import ECVDim, EdgeDim, KDim, VertexDim -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_04 import _mo_nh_diffusion_stencil_04 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_05 import _mo_nh_diffusion_stencil_05 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_06 import _mo_nh_diffusion_stencil_06 @field_operator def _fused_mo_nh_diffusion_stencil_04_05_06( @@ -40,12 +46,29 @@ def _fused_mo_nh_diffusion_stencil_04_05_06( start_2nd_nudge_line_idx_e: int32, ) -> Field[[EdgeDim, KDim], float]: - z_nabla4_e2 = _mo_nh_diffusion_stencil_04(u_vert, v_vert, primal_normal_vert_v1, primal_normal_vert_v2, z_nabla2_e, inv_vert_vert_length, inv_primal_edge_length) + z_nabla4_e2 = _mo_nh_diffusion_stencil_04( + u_vert, + v_vert, + primal_normal_vert_v1, + primal_normal_vert_v2, + z_nabla2_e, + inv_vert_vert_length, + inv_primal_edge_length, + ) vn = where( horz_idx >= start_2nd_nudge_line_idx_e, - _mo_nh_diffusion_stencil_05(area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, nudgecoeff_e, vn, nudgezone_diff), - _mo_nh_diffusion_stencil_06(z_nabla2_e, area_edge, vn, fac_bdydiff_v) + _mo_nh_diffusion_stencil_05( + area_edge, + kh_smag_e, + z_nabla2_e, + z_nabla4_e2, + diff_multfac_vn, + nudgecoeff_e, + vn, + nudgezone_diff, + ), + _mo_nh_diffusion_stencil_06(z_nabla2_e, area_edge, vn, fac_bdydiff_v), ) return vn @@ -70,6 +93,22 @@ def fused_mo_nh_diffusion_stencil_04_05_06( fac_bdydiff_v: float, start_2nd_nudge_line_idx_e: int32, ): - _fused_mo_nh_diffusion_stencil_04_05_06(u_vert, v_vert, primal_normal_vert_v1, primal_normal_vert_v2, z_nabla2_e, inv_vert_vert_length, - inv_primal_edge_length, area_edge, kh_smag_e, diff_multfac_vn, nudgecoeff_e, vn, horz_idx, nudgezone_diff, fac_bdydiff_v, - start_2nd_nudge_line_idx_e, out=vn) + _fused_mo_nh_diffusion_stencil_04_05_06( + u_vert, + v_vert, + primal_normal_vert_v1, + primal_normal_vert_v2, + z_nabla2_e, + inv_vert_vert_length, + inv_primal_edge_length, + area_edge, + kh_smag_e, + diff_multfac_vn, + nudgecoeff_e, + vn, + horz_idx, + nudgezone_diff, + fac_bdydiff_v, + start_2nd_nudge_line_idx_e, + out=vn, + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py index 0bbcce058a..a2a611a42d 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py @@ -12,14 +12,22 @@ # SPDX-License-Identifier: GPL-3.0-or-later from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, where, int32, broadcast +from functional.ffront.fbuiltins import Field, broadcast, int32, where +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_07 import ( + _mo_nh_diffusion_stencil_07, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_08 import ( + _mo_nh_diffusion_stencil_08, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_09 import ( + _mo_nh_diffusion_stencil_09, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_10 import ( + _mo_nh_diffusion_stencil_10, +) from icon4py.common.dimension import C2E2CODim, CellDim, KDim -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_07 import _mo_nh_diffusion_stencil_07 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_08 import _mo_nh_diffusion_stencil_08 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_09 import _mo_nh_diffusion_stencil_09 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_10 import _mo_nh_diffusion_stencil_10 @field_operator def _fused_mo_nh_diffusion_stencil_07_08_09_10( @@ -38,26 +46,35 @@ def _fused_mo_nh_diffusion_stencil_07_08_09_10( nrdmax: int32, interior_idx: int32, halo_idx: int32, -) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: +) -> tuple[ + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], + Field[[CellDim, KDim], float], +]: vert_idx = broadcast(vert_idx, (CellDim, KDim)) dwdx, dwdy = where( vert_idx > int32(0), _mo_nh_diffusion_stencil_08(w_old, geofac_grg_x, geofac_grg_y), - (dwdx, dwdy) + (dwdx, dwdy), ) z_nabla2_c = _mo_nh_diffusion_stencil_07(w_old, geofac_n2s) w = where( (horz_idx >= interior_idx) & (horz_idx < halo_idx), - _mo_nh_diffusion_stencil_09(area, z_nabla2_c, geofac_n2s, w_old, diff_multfac_w), + _mo_nh_diffusion_stencil_09( + area, z_nabla2_c, geofac_n2s, w_old, diff_multfac_w + ), w_old, ) w = where( - (vert_idx > int32(0)) & (vert_idx < nrdmax) & (horz_idx >= interior_idx) & (horz_idx < halo_idx), + (vert_idx > int32(0)) + & (vert_idx < nrdmax) + & (horz_idx >= interior_idx) + & (horz_idx < halo_idx), _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, area, z_nabla2_c), w, ) @@ -83,4 +100,21 @@ def fused_mo_nh_diffusion_stencil_07_08_09_10( interior_idx: int32, halo_idx: int32, ): - _fused_mo_nh_diffusion_stencil_07_08_09_10(area, geofac_n2s, geofac_grg_x, geofac_grg_y, w_old, w, dwdx, dwdy, diff_multfac_w, diff_multfac_n2w, vert_idx, horz_idx, nrdmax, interior_idx, halo_idx, out = (w, dwdx, dwdy)) + _fused_mo_nh_diffusion_stencil_07_08_09_10( + area, + geofac_n2s, + geofac_grg_x, + geofac_grg_y, + w_old, + w, + dwdx, + dwdy, + diff_multfac_w, + diff_multfac_n2w, + vert_idx, + horz_idx, + nrdmax, + interior_idx, + halo_idx, + out=(w, dwdx, dwdy), + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py index cda940f529..375f8031d3 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py @@ -14,10 +14,14 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_11 import ( + _mo_nh_diffusion_stencil_11, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_12 import ( + _mo_nh_diffusion_stencil_12, +) from icon4py.common.dimension import CellDim, EdgeDim, KDim -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_11 import _mo_nh_diffusion_stencil_11 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_12 import _mo_nh_diffusion_stencil_12 @field_operator def _fused_mo_nh_diffusion_stencil_11_12( @@ -30,6 +34,7 @@ def _fused_mo_nh_diffusion_stencil_11_12( kh_smag_e = _mo_nh_diffusion_stencil_12(kh_smag_e, enh_diffu_3d) return kh_smag_e + @program def fused_mo_nh_diffusion_stencil_11_12( theta_v: Field[[CellDim, KDim], float], @@ -37,4 +42,6 @@ def fused_mo_nh_diffusion_stencil_11_12( thresh_tdiff: float, kh_smag_e: Field[[EdgeDim, KDim], float], ): - _fused_mo_nh_diffusion_stencil_11_12(theta_v, theta_ref_mc, thresh_tdiff, kh_smag_e, out=kh_smag_e) + _fused_mo_nh_diffusion_stencil_11_12( + theta_v, theta_ref_mc, thresh_tdiff, kh_smag_e, out=kh_smag_e + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py index ae317277c1..0941bbd774 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py @@ -14,10 +14,14 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_13 import ( + _mo_nh_diffusion_stencil_13, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_14 import ( + _mo_nh_diffusion_stencil_14, +) from icon4py.common.dimension import CEDim, CellDim, EdgeDim, KDim -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_13 import _mo_nh_diffusion_stencil_13 -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_14 import _mo_nh_diffusion_stencil_14 @field_operator def _fused_mo_nh_diffusion_stencil_13_14( @@ -30,6 +34,7 @@ def _fused_mo_nh_diffusion_stencil_13_14( z_temp = _mo_nh_diffusion_stencil_14(z_nabla2_e, geofac_div) return z_temp + @program def fused_mo_nh_diffusion_stencil_13_14( kh_smag_e: Field[[EdgeDim, KDim], float], @@ -38,4 +43,6 @@ def fused_mo_nh_diffusion_stencil_13_14( geofac_div: Field[[CEDim], float], z_temp: Field[[CellDim, KDim], float], ): - _fused_mo_nh_diffusion_stencil_13_14(kh_smag_e, inv_dual_edge_length, theta_v, geofac_div, out=z_temp) + _fused_mo_nh_diffusion_stencil_13_14( + kh_smag_e, inv_dual_edge_length, theta_v, geofac_div, out=z_temp + ) From ae10ca27736694b17dd982d49b5ef601e32e6bb9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Mon, 23 Jan 2023 10:22:55 +0100 Subject: [PATCH 06/49] Change includes, remove 01_02_03_rbf from tox tests --- ..._02_03_rbf.py => fused_mo_nh_diffusion_stencil_01_02_03_rbf} | 0 .../atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py | 2 +- 2 files changed, 1 insertion(+), 1 deletion(-) rename atm_dyn_iconam/src/icon4py/atm_dyn_iconam/{fused_mo_nh_diffusion_stencil_01_02_03_rbf.py => fused_mo_nh_diffusion_stencil_01_02_03_rbf} (100%) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf similarity index 100% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf.py rename to atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py index c3a392b226..6b32d5a67e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, neighbor_sum +from functional.ffront.fbuiltins import Field from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import ( _mo_nh_diffusion_stencil_02, From f30f1b086517267118af9fdf425d422a04ae94be Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 1 Feb 2023 08:29:55 +0100 Subject: [PATCH 07/49] use renamed diffusion stencils in fused stencils --- ...fused_mo_nh_diffusion_stencil_01_02_03_rbf | 18 +++++------ .../fused_mo_nh_diffusion_stencil_02_03.py | 12 ++++---- .../fused_mo_nh_diffusion_stencil_04_05_06.py | 20 ++++++------- ...sed_mo_nh_diffusion_stencil_07_08_09_10.py | 30 +++++++++---------- .../fused_mo_nh_diffusion_stencil_11_12.py | 16 ++++++---- .../fused_mo_nh_diffusion_stencil_13_14.py | 12 ++++---- 6 files changed, 54 insertions(+), 54 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf index fb3ecaf12d..e33e884f8e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf @@ -17,14 +17,14 @@ from functional.ffront.fbuiltins import Field from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( _mo_intp_rbf_rbf_vec_interpol_vertex, ) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_01 import ( - _mo_nh_diffusion_stencil_01, +from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( + _calculate_nabla2_and_smag_coefficients_for_vn, ) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import ( - _mo_nh_diffusion_stencil_02, +from icon4py.atm_dyn_iconam.temporary_fields_for_turbulance_diagnostics import ( + _temporary_fields_for_turbulance_diagnostics, ) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import ( - _mo_nh_diffusion_stencil_03, +from icon4py.atm_dyn_iconam.calculate_diagnostics_for_turbulance import ( + _calculate_diagnostics_for_turbulance, ) from icon4py.common.dimension import ( C2EDim, @@ -65,7 +65,7 @@ def _fused_mo_nh_diffusion_stencil_01_02_03_rbf( Field[[VertexDim, KDim], float], ]: - kh_smag_e, kh_smag_ec, z_nabla2_e = _mo_nh_diffusion_stencil_01( + kh_smag_e, kh_smag_ec, z_nabla2_e = _calculate_nabla2_and_smag_coefficients_for_vn( diff_multfac_smag, tangent_orientation, inv_primal_edge_length, @@ -81,11 +81,11 @@ def _fused_mo_nh_diffusion_stencil_01_02_03_rbf( smag_offset, ) - kh_c, div = _mo_nh_diffusion_stencil_02( + kh_c, div = _temporary_fields_for_turbulance_diagnostics( kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag ) - div_ic, hdef_ic = _mo_nh_diffusion_stencil_03(div, kh_c, wgtfac_c) + div_ic, hdef_ic = _calculate_diagnostics_for_turbulance(div, kh_c, wgtfac_c) u_vert, v_vert = _mo_intp_rbf_rbf_vec_interpol_vertex( z_nabla2_e, ptr_coeff_1, ptr_coeff_2 diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py index 6b32d5a67e..2d74f63e9c 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py @@ -14,11 +14,11 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_02 import ( - _mo_nh_diffusion_stencil_02, +from icon4py.atm_dyn_iconam.calculate_diagnostics_for_turbulance import ( + _calculate_diagnostics_for_turbulance, ) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_03 import ( - _mo_nh_diffusion_stencil_03, +from icon4py.atm_dyn_iconam.temporary_fields_for_turbulance_diagnostics import ( + _temporary_fields_for_turbulance_diagnostics, ) from icon4py.common.dimension import C2EDim, CellDim, EdgeDim, KDim @@ -32,10 +32,10 @@ def _fused_mo_nh_diffusion_stencil_02_03( diff_multfac_smag: Field[[KDim], float], wgtfac_c: Field[[CellDim, KDim], float], ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: - kh_c, div = _mo_nh_diffusion_stencil_02( + kh_c, div = _temporary_fields_for_turbulance_diagnostics( kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag ) - div_ic, hdef_ic = _mo_nh_diffusion_stencil_03(div, kh_c, wgtfac_c) + div_ic, hdef_ic = _calculate_diagnostics_for_turbulance(div, kh_c, wgtfac_c) return div_ic, hdef_ic diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py index 58e3345016..a2a00a427e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py @@ -14,15 +14,13 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field, int32, where -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_04 import ( - _mo_nh_diffusion_stencil_04, +from icon4py.atm_dyn_iconam.apply_nabla2_and_nabla4_to_vn import ( + _apply_nabla2_and_nabla4_to_vn, ) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_05 import ( - _mo_nh_diffusion_stencil_05, -) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_06 import ( - _mo_nh_diffusion_stencil_06, +from icon4py.atm_dyn_iconam.apply_nabla2_to_vn_in_lateral_boundary import ( + _apply_nabla2_to_vn_in_lateral_boundary, ) +from icon4py.atm_dyn_iconam.calculate_nabla4 import _calculate_nabla4 from icon4py.common.dimension import ECVDim, EdgeDim, KDim, VertexDim @@ -46,7 +44,7 @@ def _fused_mo_nh_diffusion_stencil_04_05_06( start_2nd_nudge_line_idx_e: int32, ) -> Field[[EdgeDim, KDim], float]: - z_nabla4_e2 = _mo_nh_diffusion_stencil_04( + z_nabla4_e2 = _calculate_nabla4( u_vert, v_vert, primal_normal_vert_v1, @@ -58,7 +56,7 @@ def _fused_mo_nh_diffusion_stencil_04_05_06( vn = where( horz_idx >= start_2nd_nudge_line_idx_e, - _mo_nh_diffusion_stencil_05( + _apply_nabla2_and_nabla4_to_vn( area_edge, kh_smag_e, z_nabla2_e, @@ -68,7 +66,9 @@ def _fused_mo_nh_diffusion_stencil_04_05_06( vn, nudgezone_diff, ), - _mo_nh_diffusion_stencil_06(z_nabla2_e, area_edge, vn, fac_bdydiff_v), + _apply_nabla2_to_vn_in_lateral_boundary( + z_nabla2_e, area_edge, vn, fac_bdydiff_v + ), ) return vn diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py index a2a611a42d..0849dfe56e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py @@ -14,18 +14,14 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field, broadcast, int32, where -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_07 import ( - _mo_nh_diffusion_stencil_07, +from icon4py.atm_dyn_iconam.apply_nabla2_to_w import _apply_nabla2_to_w +from icon4py.atm_dyn_iconam.apply_nabla2_to_w_in_upper_damping_layer import ( + _apply_nabla2_to_w_in_upper_damping_layer, ) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_08 import ( - _mo_nh_diffusion_stencil_08, -) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_09 import ( - _mo_nh_diffusion_stencil_09, -) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_10 import ( - _mo_nh_diffusion_stencil_10, +from icon4py.atm_dyn_iconam.calculate_horizontal_gradients_for_turbulance import ( + _calculate_horizontal_gradients_for_turbulance, ) +from icon4py.atm_dyn_iconam.calculate_nabla2_for_w import _calculate_nabla2_for_w from icon4py.common.dimension import C2E2CODim, CellDim, KDim @@ -56,17 +52,17 @@ def _fused_mo_nh_diffusion_stencil_07_08_09_10( dwdx, dwdy = where( vert_idx > int32(0), - _mo_nh_diffusion_stencil_08(w_old, geofac_grg_x, geofac_grg_y), + _calculate_horizontal_gradients_for_turbulance( + w_old, geofac_grg_x, geofac_grg_y + ), (dwdx, dwdy), ) - z_nabla2_c = _mo_nh_diffusion_stencil_07(w_old, geofac_n2s) + z_nabla2_c = _calculate_nabla2_for_w(w_old, geofac_n2s) w = where( (horz_idx >= interior_idx) & (horz_idx < halo_idx), - _mo_nh_diffusion_stencil_09( - area, z_nabla2_c, geofac_n2s, w_old, diff_multfac_w - ), + _apply_nabla2_to_w(area, z_nabla2_c, geofac_n2s, w_old, diff_multfac_w), w_old, ) @@ -75,7 +71,9 @@ def _fused_mo_nh_diffusion_stencil_07_08_09_10( & (vert_idx < nrdmax) & (horz_idx >= interior_idx) & (horz_idx < halo_idx), - _mo_nh_diffusion_stencil_10(w, diff_multfac_n2w, area, z_nabla2_c), + _apply_nabla2_to_w_in_upper_damping_layer( + w, diff_multfac_n2w, area, z_nabla2_c + ), w, ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py index 375f8031d3..8374f29a94 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py @@ -14,11 +14,11 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_11 import ( - _mo_nh_diffusion_stencil_11, +from icon4py.atm_dyn_iconam.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( + _enhance_diffusion_coefficient_for_grid_point_cold_pools, ) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_12 import ( - _mo_nh_diffusion_stencil_12, +from icon4py.atm_dyn_iconam.temporary_field_for_grid_point_cold_pools_enhancement import ( + _temporary_field_for_grid_point_cold_pools_enhancement, ) from icon4py.common.dimension import CellDim, EdgeDim, KDim @@ -30,8 +30,12 @@ def _fused_mo_nh_diffusion_stencil_11_12( thresh_tdiff: float, kh_smag_e: Field[[EdgeDim, KDim], float], ) -> Field[[EdgeDim, KDim], float]: - enh_diffu_3d = _mo_nh_diffusion_stencil_11(theta_v, theta_ref_mc, thresh_tdiff) - kh_smag_e = _mo_nh_diffusion_stencil_12(kh_smag_e, enh_diffu_3d) + enh_diffu_3d = _temporary_field_for_grid_point_cold_pools_enhancement( + theta_v, theta_ref_mc, thresh_tdiff + ) + kh_smag_e = _enhance_diffusion_coefficient_for_grid_point_cold_pools( + kh_smag_e, enh_diffu_3d + ) return kh_smag_e diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py index 0941bbd774..1b286e0f0c 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py @@ -14,11 +14,9 @@ from functional.ffront.decorator import field_operator, program from functional.ffront.fbuiltins import Field -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_13 import ( - _mo_nh_diffusion_stencil_13, -) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_14 import ( - _mo_nh_diffusion_stencil_14, +from icon4py.atm_dyn_iconam.calculate_nabla2_for_z import _calculate_nabla2_for_z +from icon4py.atm_dyn_iconam.calculate_nabla2_of_theta import ( + _calculate_nabla2_of_theta, ) from icon4py.common.dimension import CEDim, CellDim, EdgeDim, KDim @@ -30,8 +28,8 @@ def _fused_mo_nh_diffusion_stencil_13_14( theta_v: Field[[CellDim, KDim], float], geofac_div: Field[[CEDim], float], ) -> Field[[CellDim, KDim], float]: - z_nabla2_e = _mo_nh_diffusion_stencil_13(kh_smag_e, inv_dual_edge_length, theta_v) - z_temp = _mo_nh_diffusion_stencil_14(z_nabla2_e, geofac_div) + z_nabla2_e = _calculate_nabla2_for_z(kh_smag_e, inv_dual_edge_length, theta_v) + z_temp = _calculate_nabla2_of_theta(z_nabla2_e, geofac_div) return z_temp From 4ab35bcc1b7c21411c0feb837a0b08842311880e Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Wed, 1 Feb 2023 14:47:43 +0100 Subject: [PATCH 08/49] readability refactoring (2): rename fused stencils --- ...ffusion_stencil_04_05_06.py => apply_diffusion_to_vn.py} | 6 +++--- ...py => calculate_diagnostic_quantities_for_turbulence.py} | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) rename atm_dyn_iconam/src/icon4py/atm_dyn_iconam/{fused_mo_nh_diffusion_stencil_04_05_06.py => apply_diffusion_to_vn.py} (96%) rename atm_dyn_iconam/src/icon4py/atm_dyn_iconam/{fused_mo_nh_diffusion_stencil_02_03.py => calculate_diagnostic_quantities_for_turbulence.py} (93%) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py similarity index 96% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py rename to atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py index a2a00a427e..1927ea999a 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_04_05_06.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py @@ -25,7 +25,7 @@ @field_operator -def _fused_mo_nh_diffusion_stencil_04_05_06( +def _apply_diffusion_to_vn( u_vert: Field[[VertexDim, KDim], float], v_vert: Field[[VertexDim, KDim], float], primal_normal_vert_v1: Field[[ECVDim], float], @@ -75,7 +75,7 @@ def _fused_mo_nh_diffusion_stencil_04_05_06( @program -def fused_mo_nh_diffusion_stencil_04_05_06( +def apply_diffusion_to_vn( u_vert: Field[[VertexDim, KDim], float], v_vert: Field[[VertexDim, KDim], float], primal_normal_vert_v1: Field[[ECVDim], float], @@ -93,7 +93,7 @@ def fused_mo_nh_diffusion_stencil_04_05_06( fac_bdydiff_v: float, start_2nd_nudge_line_idx_e: int32, ): - _fused_mo_nh_diffusion_stencil_04_05_06( + _apply_diffusion_to_vn( u_vert, v_vert, primal_normal_vert_v1, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py similarity index 93% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py rename to atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py index 2d74f63e9c..da014532b7 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_02_03.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py @@ -24,7 +24,7 @@ @field_operator -def _fused_mo_nh_diffusion_stencil_02_03( +def _calculate_diagnostic_quantities_for_turbulence( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], e_bln_c_s: Field[[CellDim, C2EDim], float], @@ -40,7 +40,7 @@ def _fused_mo_nh_diffusion_stencil_02_03( @program -def fused_mo_nh_diffusion_stencil_02_03( +def calculate_diagnostic_quantities_for_turbulence( kh_smag_ec: Field[[EdgeDim, KDim], float], vn: Field[[EdgeDim, KDim], float], e_bln_c_s: Field[[CellDim, C2EDim], float], @@ -50,7 +50,7 @@ def fused_mo_nh_diffusion_stencil_02_03( div_ic: Field[[CellDim, KDim], float], hdef_ic: Field[[CellDim, KDim], float], ): - _fused_mo_nh_diffusion_stencil_02_03( + _calculate_diagnostic_quantities_for_turbulence( kh_smag_ec, vn, e_bln_c_s, From 5d0924bed15dcccd4db2266571790ba546b208c0 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 7 Feb 2023 15:15:28 +0100 Subject: [PATCH 09/49] rename fused_mo_nh_diffusion_stencil_07_08_09_10 --- ...to_w_and_compute_horizontal_gradients_for_turbulance.py} | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) rename atm_dyn_iconam/src/icon4py/atm_dyn_iconam/{fused_mo_nh_diffusion_stencil_07_08_09_10.py => apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py} (94%) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py similarity index 94% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py rename to atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py index 0849dfe56e..0dd42164cf 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_07_08_09_10.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py @@ -26,7 +26,7 @@ @field_operator -def _fused_mo_nh_diffusion_stencil_07_08_09_10( +def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( area: Field[[CellDim], float], geofac_n2s: Field[[CellDim, C2E2CODim], float], geofac_grg_x: Field[[CellDim, C2E2CODim], float], @@ -81,7 +81,7 @@ def _fused_mo_nh_diffusion_stencil_07_08_09_10( @program -def fused_mo_nh_diffusion_stencil_07_08_09_10( +def apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( area: Field[[CellDim], float], geofac_n2s: Field[[CellDim, C2E2CODim], float], geofac_grg_x: Field[[CellDim, C2E2CODim], float], @@ -98,7 +98,7 @@ def fused_mo_nh_diffusion_stencil_07_08_09_10( interior_idx: int32, halo_idx: int32, ): - _fused_mo_nh_diffusion_stencil_07_08_09_10( + _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( area, geofac_n2s, geofac_grg_x, From eac329f76f3c30522431ccd66c72b84e030807e8 Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 7 Feb 2023 15:27:43 +0100 Subject: [PATCH 10/49] rename fused_mo_nh_diffusion_stencil_11_12 --- ...ced_diffusion_coefficients_for_grid_point_cold_pools.py} | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) rename atm_dyn_iconam/src/icon4py/atm_dyn_iconam/{fused_mo_nh_diffusion_stencil_11_12.py => calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py} (88%) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py similarity index 88% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py rename to atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py index 8374f29a94..17740e49da 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_11_12.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py @@ -24,7 +24,7 @@ @field_operator -def _fused_mo_nh_diffusion_stencil_11_12( +def _calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools( theta_v: Field[[CellDim, KDim], float], theta_ref_mc: Field[[CellDim, KDim], float], thresh_tdiff: float, @@ -40,12 +40,12 @@ def _fused_mo_nh_diffusion_stencil_11_12( @program -def fused_mo_nh_diffusion_stencil_11_12( +def calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools( theta_v: Field[[CellDim, KDim], float], theta_ref_mc: Field[[CellDim, KDim], float], thresh_tdiff: float, kh_smag_e: Field[[EdgeDim, KDim], float], ): - _fused_mo_nh_diffusion_stencil_11_12( + _calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools( theta_v, theta_ref_mc, thresh_tdiff, kh_smag_e, out=kh_smag_e ) From cd68da96228ef06cf5aa3a672477183d8863a4ad Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 7 Feb 2023 15:38:03 +0100 Subject: [PATCH 11/49] rename fused_mo_nh_diffusion_stencil_13_14 --- ...tencil_13_14.py => calculate_nabla2_for_theta.py} | 12 ++++++------ ...culate_nabla2_of_theta.py => calculate_z_temp.py} | 6 +++--- ...e_nabla2_of_theta.py => test_calculate_z_temp.py} | 12 ++++++------ 3 files changed, 15 insertions(+), 15 deletions(-) rename atm_dyn_iconam/src/icon4py/atm_dyn_iconam/{fused_mo_nh_diffusion_stencil_13_14.py => calculate_nabla2_for_theta.py} (83%) rename atm_dyn_iconam/src/icon4py/atm_dyn_iconam/{calculate_nabla2_of_theta.py => calculate_z_temp.py} (89%) rename atm_dyn_iconam/tests/{test_calculate_nabla2_of_theta.py => test_calculate_z_temp.py} (87%) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py similarity index 83% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py rename to atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py index 1b286e0f0c..3ea2a5a68a 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_13_14.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py @@ -15,32 +15,32 @@ from functional.ffront.fbuiltins import Field from icon4py.atm_dyn_iconam.calculate_nabla2_for_z import _calculate_nabla2_for_z -from icon4py.atm_dyn_iconam.calculate_nabla2_of_theta import ( - _calculate_nabla2_of_theta, +from icon4py.atm_dyn_iconam.calculate_z_temp import ( + _calculate_z_temp, ) from icon4py.common.dimension import CEDim, CellDim, EdgeDim, KDim @field_operator -def _fused_mo_nh_diffusion_stencil_13_14( +def _calculate_nabla2_for_theta( kh_smag_e: Field[[EdgeDim, KDim], float], inv_dual_edge_length: Field[[EdgeDim], float], theta_v: Field[[CellDim, KDim], float], geofac_div: Field[[CEDim], float], ) -> Field[[CellDim, KDim], float]: z_nabla2_e = _calculate_nabla2_for_z(kh_smag_e, inv_dual_edge_length, theta_v) - z_temp = _calculate_nabla2_of_theta(z_nabla2_e, geofac_div) + z_temp = _calculate_z_temp(z_nabla2_e, geofac_div) return z_temp @program -def fused_mo_nh_diffusion_stencil_13_14( +def calculate_nabla2_for_theta( kh_smag_e: Field[[EdgeDim, KDim], float], inv_dual_edge_length: Field[[EdgeDim], float], theta_v: Field[[CellDim, KDim], float], geofac_div: Field[[CEDim], float], z_temp: Field[[CellDim, KDim], float], ): - _fused_mo_nh_diffusion_stencil_13_14( + _calculate_nabla2_for_theta( kh_smag_e, inv_dual_edge_length, theta_v, geofac_div, out=z_temp ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_of_theta.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_z_temp.py similarity index 89% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_of_theta.py rename to atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_z_temp.py index 93047d92fc..22211b555b 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_of_theta.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_z_temp.py @@ -26,7 +26,7 @@ @field_operator -def _calculate_nabla2_of_theta( +def _calculate_z_temp( z_nabla2_e: Field[[EdgeDim, KDim], float], geofac_div: Field[[CEDim], float], ) -> Field[[CellDim, KDim], float]: @@ -35,9 +35,9 @@ def _calculate_nabla2_of_theta( @program -def calculate_nabla2_of_theta( +def calculate_z_temp( z_nabla2_e: Field[[EdgeDim, KDim], float], geofac_div: Field[[CEDim], float], z_temp: Field[[CellDim, KDim], float], ): - _calculate_nabla2_of_theta(z_nabla2_e, geofac_div, out=z_temp) + _calculate_z_temp(z_nabla2_e, geofac_div, out=z_temp) diff --git a/atm_dyn_iconam/tests/test_calculate_nabla2_of_theta.py b/atm_dyn_iconam/tests/test_calculate_z_temp.py similarity index 87% rename from atm_dyn_iconam/tests/test_calculate_nabla2_of_theta.py rename to atm_dyn_iconam/tests/test_calculate_z_temp.py index 344bbc1092..47e3197911 100644 --- a/atm_dyn_iconam/tests/test_calculate_nabla2_of_theta.py +++ b/atm_dyn_iconam/tests/test_calculate_z_temp.py @@ -14,15 +14,15 @@ import numpy as np from functional.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.atm_dyn_iconam.calculate_nabla2_of_theta import ( - calculate_nabla2_of_theta, +from icon4py.atm_dyn_iconam.calculate_z_temp import ( + calculate_z_temp, ) from icon4py.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import as_1D_sparse_field, random_field, zero_field -def calculate_nabla2_of_theta_numpy( +def calculate_z_temp_numpy( c2e: np.array, z_nabla2_e: np.array, geofac_div: np.array ) -> np.array: geofac_div = np.expand_dims(geofac_div, axis=-1) @@ -30,7 +30,7 @@ def calculate_nabla2_of_theta_numpy( return z_temp -def test_calculate_nabla2_of_theta(): +def test_calculate_z_temp(): mesh = SimpleMesh() z_nabla2_e = random_field(mesh, EdgeDim, KDim) @@ -39,10 +39,10 @@ def test_calculate_nabla2_of_theta(): out = zero_field(mesh, CellDim, KDim) - ref = calculate_nabla2_of_theta_numpy( + ref = calculate_z_temp_numpy( mesh.c2e, np.asarray(z_nabla2_e), np.asarray(geofac_div) ) - calculate_nabla2_of_theta( + calculate_z_temp( z_nabla2_e, geofac_div_new, out, From 6f325b701d1408a0bb0f05dfa27b84360a323b3a Mon Sep 17 00:00:00 2001 From: Magdalena Luz Date: Tue, 7 Feb 2023 16:46:14 +0100 Subject: [PATCH 12/49] fix imports from gt4py --- .../src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py | 4 ++-- ...o_w_and_compute_horizontal_gradients_for_turbulance.py | 4 ++-- .../calculate_diagnostic_quantities_for_turbulence.py | 4 ++-- ...ed_diffusion_coefficients_for_grid_point_cold_pools.py | 4 ++-- .../icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py | 8 +++----- atm_dyn_iconam/tests/test_calculate_z_temp.py | 4 +--- 6 files changed, 12 insertions(+), 16 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py index 1927ea999a..999fc80119 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py @@ -11,8 +11,8 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, int32, where +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, int32, where from icon4py.atm_dyn_iconam.apply_nabla2_and_nabla4_to_vn import ( _apply_nabla2_and_nabla4_to_vn, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py index 0dd42164cf..a98ac924a5 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py @@ -11,8 +11,8 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field, broadcast, int32, where +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where from icon4py.atm_dyn_iconam.apply_nabla2_to_w import _apply_nabla2_to_w from icon4py.atm_dyn_iconam.apply_nabla2_to_w_in_upper_damping_layer import ( diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py index da014532b7..5b4b2ac901 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_diagnostic_quantities_for_turbulence.py @@ -11,8 +11,8 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field from icon4py.atm_dyn_iconam.calculate_diagnostics_for_turbulance import ( _calculate_diagnostics_for_turbulance, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py index 17740e49da..98fef2843a 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_enhanced_diffusion_coefficients_for_grid_point_cold_pools.py @@ -11,8 +11,8 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field from icon4py.atm_dyn_iconam.enhance_diffusion_coefficient_for_grid_point_cold_pools import ( _enhance_diffusion_coefficient_for_grid_point_cold_pools, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py index 3ea2a5a68a..a300af542c 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py @@ -11,13 +11,11 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field from icon4py.atm_dyn_iconam.calculate_nabla2_for_z import _calculate_nabla2_for_z -from icon4py.atm_dyn_iconam.calculate_z_temp import ( - _calculate_z_temp, -) +from icon4py.atm_dyn_iconam.calculate_z_temp import _calculate_z_temp from icon4py.common.dimension import CEDim, CellDim, EdgeDim, KDim diff --git a/atm_dyn_iconam/tests/test_calculate_z_temp.py b/atm_dyn_iconam/tests/test_calculate_z_temp.py index 0927a2074f..f9469f951f 100644 --- a/atm_dyn_iconam/tests/test_calculate_z_temp.py +++ b/atm_dyn_iconam/tests/test_calculate_z_temp.py @@ -14,9 +14,7 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.atm_dyn_iconam.calculate_z_temp import ( - calculate_z_temp, -) +from icon4py.atm_dyn_iconam.calculate_z_temp import calculate_z_temp from icon4py.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import as_1D_sparse_field, random_field, zero_field From 1e9c67a42ff2cd497692c3124cdae32630918300 Mon Sep 17 00:00:00 2001 From: Matthias Roethlin Date: Mon, 13 Mar 2023 14:26:10 +0100 Subject: [PATCH 13/49] make spelling internally consistent --- ..._to_w_and_compute_horizontal_gradients_for_turbulance.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py index a98ac924a5..da1d164d01 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py @@ -18,8 +18,8 @@ from icon4py.atm_dyn_iconam.apply_nabla2_to_w_in_upper_damping_layer import ( _apply_nabla2_to_w_in_upper_damping_layer, ) -from icon4py.atm_dyn_iconam.calculate_horizontal_gradients_for_turbulance import ( - _calculate_horizontal_gradients_for_turbulance, +from icon4py.atm_dyn_iconam.calculate_horizontal_gradients_for_turbulence import ( + _calculate_horizontal_gradients_for_turbulence, ) from icon4py.atm_dyn_iconam.calculate_nabla2_for_w import _calculate_nabla2_for_w from icon4py.common.dimension import C2E2CODim, CellDim, KDim @@ -52,7 +52,7 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( dwdx, dwdy = where( vert_idx > int32(0), - _calculate_horizontal_gradients_for_turbulance( + _calculate_horizontal_gradients_for_turbulence( w_old, geofac_grg_x, geofac_grg_y ), (dwdx, dwdy), From ce8877ec33dd75753160ac2efe476c23bb99a39c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Wed, 26 Apr 2023 10:23:50 +0200 Subject: [PATCH 14/49] Renamed calculate_nabla2_of_theta correctly --- .../{calculate_z_temp.py => calculate_nabla2_of_theta.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename atm_dyn_iconam/src/icon4py/atm_dyn_iconam/{calculate_z_temp.py => calculate_nabla2_of_theta.py} (100%) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_z_temp.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_of_theta.py similarity index 100% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_z_temp.py rename to atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_of_theta.py From 176dbf11ced97f11afbb39f7dfb0e861182989de Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Wed, 24 May 2023 16:05:30 +0200 Subject: [PATCH 15/49] Add fused diffusion stencil containing stencil 15 --- .../apply_diffusion_to_theta_and_exner.py | 92 +++++++++++++++++++ .../calculate_nabla2_for_theta.py | 4 +- atm_dyn_iconam/tests/test_calculate_z_temp.py | 4 +- 3 files changed, 98 insertions(+), 2 deletions(-) create mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_theta_and_exner.py diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_theta_and_exner.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_theta_and_exner.py new file mode 100644 index 0000000000..2da666b1b9 --- /dev/null +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_theta_and_exner.py @@ -0,0 +1,92 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later + +from gt4py.next.ffront.decorator import field_operator, program +from gt4py.next.ffront.fbuiltins import Field, int32 + +from icon4py.atm_dyn_iconam.calculate_nabla2_for_z import _calculate_nabla2_for_z +from icon4py.atm_dyn_iconam.calculate_nabla2_of_theta import ( + _calculate_nabla2_of_theta, +) +from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_15 import ( + _mo_nh_diffusion_stencil_15, +) +from icon4py.atm_dyn_iconam.update_theta_and_exner import _update_theta_and_exner +from icon4py.common.dimension import CECDim, CEDim, CellDim, EdgeDim, KDim + + +@field_operator +def _apply_diffusion_to_theta_and_exner( + kh_smag_e: Field[[EdgeDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], + theta_v_in: Field[[CellDim, KDim], float], + geofac_div: Field[[CEDim], float], + mask: Field[[CellDim, KDim], bool], + zd_vertoffset: Field[[CECDim, KDim], int32], + zd_diffcoef: Field[[CellDim, KDim], float], + geofac_n2s_c: Field[[CellDim], float], + geofac_n2s_nbh: Field[[CECDim], float], + vcoef: Field[[CECDim, KDim], float], + area: Field[[CellDim], float], + exner: Field[[CellDim, KDim], float], + rd_o_cvd: float, +) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: + z_nabla2_e = _calculate_nabla2_for_z(kh_smag_e, inv_dual_edge_length, theta_v_in) + z_temp = _calculate_nabla2_of_theta(z_nabla2_e, geofac_div) + z_temp = _mo_nh_diffusion_stencil_15( + mask, + zd_vertoffset, + zd_diffcoef, + geofac_n2s_c, + geofac_n2s_nbh, + vcoef, + theta_v_in, + z_temp, + ) + theta_v, exner = _update_theta_and_exner(z_temp, area, theta_v_in, exner, rd_o_cvd) + return theta_v, exner + + +@program +def apply_diffusion_to_theta_and_exner( + kh_smag_e: Field[[EdgeDim, KDim], float], + inv_dual_edge_length: Field[[EdgeDim], float], + theta_v_in: Field[[CellDim, KDim], float], + geofac_div: Field[[CEDim], float], + mask: Field[[CellDim, KDim], bool], + zd_vertoffset: Field[[CECDim, KDim], int32], + zd_diffcoef: Field[[CellDim, KDim], float], + geofac_n2s_c: Field[[CellDim], float], + geofac_n2s_nbh: Field[[CECDim], float], + vcoef: Field[[CECDim, KDim], float], + area: Field[[CellDim], float], + theta_v: Field[[CellDim, KDim], float], + exner: Field[[CellDim, KDim], float], + rd_o_cvd: float, +): + _apply_diffusion_to_theta_and_exner( + kh_smag_e, + inv_dual_edge_length, + theta_v_in, + geofac_div, + mask, + zd_vertoffset, + zd_diffcoef, + geofac_n2s_c, + geofac_n2s_nbh, + vcoef, + area, + exner, + rd_o_cvd, + out=(theta_v, exner), + ) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py index c26d2f4007..7dba291f33 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/calculate_nabla2_for_theta.py @@ -15,7 +15,9 @@ from gt4py.next.ffront.fbuiltins import Field from icon4py.atm_dyn_iconam.calculate_nabla2_for_z import _calculate_nabla2_for_z -from icon4py.atm_dyn_iconam.calculate_nabla2_of_theta import _calculate_nabla2_of_theta +from icon4py.atm_dyn_iconam.calculate_nabla2_of_theta import ( + _calculate_nabla2_of_theta, +) from icon4py.common.dimension import CEDim, CellDim, EdgeDim, KDim diff --git a/atm_dyn_iconam/tests/test_calculate_z_temp.py b/atm_dyn_iconam/tests/test_calculate_z_temp.py index 85f9a8928d..f4383abc1f 100644 --- a/atm_dyn_iconam/tests/test_calculate_z_temp.py +++ b/atm_dyn_iconam/tests/test_calculate_z_temp.py @@ -14,7 +14,9 @@ import numpy as np from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.atm_dyn_iconam.calculate_nabla2_of_theta import calculate_nabla2_of_theta +from icon4py.atm_dyn_iconam.calculate_nabla2_of_theta import ( + calculate_nabla2_of_theta, +) from icon4py.common.dimension import C2EDim, CEDim, CellDim, EdgeDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh from icon4py.testutils.utils import as_1D_sparse_field, random_field, zero_field From f5ef5e8b919a6b3e977c9ec7e588685b6020b568 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Thu, 25 May 2023 17:38:44 +0200 Subject: [PATCH 16/49] Add global component to the fused diffusion vn update stencil --- .../atm_dyn_iconam/apply_diffusion_to_vn.py | 38 +++++++++++++------ 1 file changed, 27 insertions(+), 11 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py index 999fc80119..5412451132 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py @@ -12,8 +12,11 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, int32, where +from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where +from icon4py.atm_dyn_iconam.apply_nabla2_and_nabla4_global_to_vn import ( + _apply_nabla2_and_nabla4_global_to_vn, +) from icon4py.atm_dyn_iconam.apply_nabla2_and_nabla4_to_vn import ( _apply_nabla2_and_nabla4_to_vn, ) @@ -42,8 +45,8 @@ def _apply_diffusion_to_vn( nudgezone_diff: float, fac_bdydiff_v: float, start_2nd_nudge_line_idx_e: int32, + limited_area: bool, ) -> Field[[EdgeDim, KDim], float]: - z_nabla4_e2 = _calculate_nabla4( u_vert, v_vert, @@ -56,15 +59,26 @@ def _apply_diffusion_to_vn( vn = where( horz_idx >= start_2nd_nudge_line_idx_e, - _apply_nabla2_and_nabla4_to_vn( - area_edge, - kh_smag_e, - z_nabla2_e, - z_nabla4_e2, - diff_multfac_vn, - nudgecoeff_e, - vn, - nudgezone_diff, + where( + broadcast(limited_area, (EdgeDim, KDim)), + _apply_nabla2_and_nabla4_to_vn( + area_edge, + kh_smag_e, + z_nabla2_e, + z_nabla4_e2, + diff_multfac_vn, + nudgecoeff_e, + vn, + nudgezone_diff, + ), + _apply_nabla2_and_nabla4_global_to_vn( + area_edge, + kh_smag_e, + z_nabla2_e, + z_nabla4_e2, + diff_multfac_vn, + vn, + ), ), _apply_nabla2_to_vn_in_lateral_boundary( z_nabla2_e, area_edge, vn, fac_bdydiff_v @@ -92,6 +106,7 @@ def apply_diffusion_to_vn( nudgezone_diff: float, fac_bdydiff_v: float, start_2nd_nudge_line_idx_e: int32, + limited_area: bool, ): _apply_diffusion_to_vn( u_vert, @@ -110,5 +125,6 @@ def apply_diffusion_to_vn( nudgezone_diff, fac_bdydiff_v, start_2nd_nudge_line_idx_e, + limited_area, out=vn, ) From 8976ec6122a4ea8dd33be9beb96dafea0ea6465a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Fri, 26 May 2023 12:05:29 +0200 Subject: [PATCH 17/49] Renamed stencil 15, restructured fused vn stencil for scalar if --- .../apply_diffusion_to_theta_and_exner.py | 6 +- .../atm_dyn_iconam/apply_diffusion_to_vn.py | 67 +++++++++++++------ ...usion_nabla_of_theta_over_steep_points.py} | 6 +- ...usion_nabla_of_theta_over_steep_points.py} | 6 +- 4 files changed, 56 insertions(+), 29 deletions(-) rename atm_dyn_iconam/src/icon4py/atm_dyn_iconam/{mo_nh_diffusion_stencil_15.py => truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py} (93%) rename atm_dyn_iconam/tests/{test_mo_nh_diffusion_stencil_15.py => test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py} (94%) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_theta_and_exner.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_theta_and_exner.py index 2da666b1b9..f607cfb262 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_theta_and_exner.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_theta_and_exner.py @@ -18,8 +18,8 @@ from icon4py.atm_dyn_iconam.calculate_nabla2_of_theta import ( _calculate_nabla2_of_theta, ) -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_15 import ( - _mo_nh_diffusion_stencil_15, +from icon4py.atm_dyn_iconam.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( + _truly_horizontal_diffusion_nabla_of_theta_over_steep_points, ) from icon4py.atm_dyn_iconam.update_theta_and_exner import _update_theta_and_exner from icon4py.common.dimension import CECDim, CEDim, CellDim, EdgeDim, KDim @@ -43,7 +43,7 @@ def _apply_diffusion_to_theta_and_exner( ) -> tuple[Field[[CellDim, KDim], float], Field[[CellDim, KDim], float]]: z_nabla2_e = _calculate_nabla2_for_z(kh_smag_e, inv_dual_edge_length, theta_v_in) z_temp = _calculate_nabla2_of_theta(z_nabla2_e, geofac_div) - z_temp = _mo_nh_diffusion_stencil_15( + z_temp = _truly_horizontal_diffusion_nabla_of_theta_over_steep_points( mask, zd_vertoffset, zd_diffcoef, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py index 5412451132..18f9c7a0e9 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py @@ -27,6 +27,43 @@ from icon4py.common.dimension import ECVDim, EdgeDim, KDim, VertexDim +@field_operator +def _apply_nabla2_and_nabla4_to_vn_switcher( + area_edge: Field[[EdgeDim], float], + kh_smag_e: Field[[EdgeDim, KDim], float], + z_nabla2_e: Field[[EdgeDim, KDim], float], + z_nabla4_e2: Field[[EdgeDim, KDim], float], + diff_multfac_vn: Field[[KDim], float], + nudgecoeff_e: Field[[EdgeDim], float], + vn: Field[[EdgeDim, KDim], float], + nudgezone_diff: float, + limited_area: bool, +) -> Field[[EdgeDim, KDim], float]: + vn = where( + broadcast(limited_area, (EdgeDim, KDim)), + _apply_nabla2_and_nabla4_to_vn( + area_edge, + kh_smag_e, + z_nabla2_e, + z_nabla4_e2, + diff_multfac_vn, + nudgecoeff_e, + vn, + nudgezone_diff, + ), + _apply_nabla2_and_nabla4_global_to_vn( + area_edge, + kh_smag_e, + z_nabla2_e, + z_nabla4_e2, + diff_multfac_vn, + vn, + ), + ) + + return vn + + @field_operator def _apply_diffusion_to_vn( u_vert: Field[[VertexDim, KDim], float], @@ -59,26 +96,16 @@ def _apply_diffusion_to_vn( vn = where( horz_idx >= start_2nd_nudge_line_idx_e, - where( - broadcast(limited_area, (EdgeDim, KDim)), - _apply_nabla2_and_nabla4_to_vn( - area_edge, - kh_smag_e, - z_nabla2_e, - z_nabla4_e2, - diff_multfac_vn, - nudgecoeff_e, - vn, - nudgezone_diff, - ), - _apply_nabla2_and_nabla4_global_to_vn( - area_edge, - kh_smag_e, - z_nabla2_e, - z_nabla4_e2, - diff_multfac_vn, - vn, - ), + _apply_nabla2_and_nabla4_to_vn_switcher( + area_edge, + kh_smag_e, + z_nabla2_e, + z_nabla4_e2, + diff_multfac_vn, + nudgecoeff_e, + vn, + nudgezone_diff, + limited_area, ), _apply_nabla2_to_vn_in_lateral_boundary( z_nabla2_e, area_edge, vn, fac_bdydiff_v diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py similarity index 93% rename from atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py rename to atm_dyn_iconam/src/icon4py/atm_dyn_iconam/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index 9abed41ee1..ab0574e40e 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/mo_nh_diffusion_stencil_15.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -19,7 +19,7 @@ @field_operator -def _mo_nh_diffusion_stencil_15( +def _truly_horizontal_diffusion_nabla_of_theta_over_steep_points( mask: Field[[CellDim, KDim], bool], zd_vertoffset: Field[[CECDim, KDim], int32], zd_diffcoef: Field[[CellDim, KDim], float], @@ -66,7 +66,7 @@ def _mo_nh_diffusion_stencil_15( @program -def mo_nh_diffusion_stencil_15( +def truly_horizontal_diffusion_nabla_of_theta_over_steep_points( mask: Field[[CellDim, KDim], bool], zd_vertoffset: Field[[CECDim, KDim], int32], zd_diffcoef: Field[[CellDim, KDim], float], @@ -76,7 +76,7 @@ def mo_nh_diffusion_stencil_15( theta_v: Field[[CellDim, KDim], float], z_temp: Field[[CellDim, KDim], float], ): - _mo_nh_diffusion_stencil_15( + _truly_horizontal_diffusion_nabla_of_theta_over_steep_points( mask, zd_vertoffset, zd_diffcoef, diff --git a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py b/atm_dyn_iconam/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py similarity index 94% rename from atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py rename to atm_dyn_iconam/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py index 9d4f6661a7..45fc64a3ab 100644 --- a/atm_dyn_iconam/tests/test_mo_nh_diffusion_stencil_15.py +++ b/atm_dyn_iconam/tests/test_truly_horizontal_diffusion_nabla_of_theta_over_steep_points.py @@ -15,8 +15,8 @@ from gt4py.next.ffront.fbuiltins import int32 from gt4py.next.iterator.embedded import StridedNeighborOffsetProvider -from icon4py.atm_dyn_iconam.mo_nh_diffusion_stencil_15 import ( - mo_nh_diffusion_stencil_15, +from icon4py.atm_dyn_iconam.truly_horizontal_diffusion_nabla_of_theta_over_steep_points import ( + truly_horizontal_diffusion_nabla_of_theta_over_steep_points, ) from icon4py.common.dimension import C2E2CDim, CECDim, CellDim, KDim from icon4py.testutils.simple_mesh import SimpleMesh @@ -110,7 +110,7 @@ def test_mo_nh_diffusion_stencil_15(): kstart = 0 kend = mesh.k_level - mo_nh_diffusion_stencil_15( + truly_horizontal_diffusion_nabla_of_theta_over_steep_points( mask, zd_vertoffset_new, zd_diffcoef, From 8e5d414f3f96cc617ff6ff0ab0f6b1cc16b0d76f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Tue, 30 May 2023 16:33:20 +0200 Subject: [PATCH 18/49] Substituted where with if expression --- .../icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py index 18f9c7a0e9..76804aef88 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py @@ -12,7 +12,7 @@ # SPDX-License-Identifier: GPL-3.0-or-later from gt4py.next.ffront.decorator import field_operator, program -from gt4py.next.ffront.fbuiltins import Field, broadcast, int32, where +from gt4py.next.ffront.fbuiltins import Field, int32, where from icon4py.atm_dyn_iconam.apply_nabla2_and_nabla4_global_to_vn import ( _apply_nabla2_and_nabla4_global_to_vn, @@ -39,8 +39,8 @@ def _apply_nabla2_and_nabla4_to_vn_switcher( nudgezone_diff: float, limited_area: bool, ) -> Field[[EdgeDim, KDim], float]: - vn = where( - broadcast(limited_area, (EdgeDim, KDim)), + # TODO: Use normal if-else instead + vn = ( _apply_nabla2_and_nabla4_to_vn( area_edge, kh_smag_e, @@ -50,15 +50,16 @@ def _apply_nabla2_and_nabla4_to_vn_switcher( nudgecoeff_e, vn, nudgezone_diff, - ), - _apply_nabla2_and_nabla4_global_to_vn( + ) + if limited_area + else _apply_nabla2_and_nabla4_global_to_vn( area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, vn, - ), + ) ) return vn From a294505e8a4b5e11937d658e84b58ff53836f20d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Wed, 31 May 2023 00:49:25 +0200 Subject: [PATCH 19/49] Undo renaming of test --- ...test_calculate_z_temp.py => test_calculate_nabla2_of_theta.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename atm_dyn_iconam/tests/{test_calculate_z_temp.py => test_calculate_nabla2_of_theta.py} (100%) diff --git a/atm_dyn_iconam/tests/test_calculate_z_temp.py b/atm_dyn_iconam/tests/test_calculate_nabla2_of_theta.py similarity index 100% rename from atm_dyn_iconam/tests/test_calculate_z_temp.py rename to atm_dyn_iconam/tests/test_calculate_nabla2_of_theta.py From fc60b776b3bfa7645fff0a5cd0aa7e2e3d59ed67 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Wed, 31 May 2023 17:44:34 +0200 Subject: [PATCH 20/49] Magdalena feedback, lateral boundary stencil incorrectly called in global mode --- .../atm_dyn_iconam/apply_diffusion_to_vn.py | 70 ++++++------------- 1 file changed, 22 insertions(+), 48 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py index 76804aef88..fa464cbca2 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py @@ -27,44 +27,6 @@ from icon4py.common.dimension import ECVDim, EdgeDim, KDim, VertexDim -@field_operator -def _apply_nabla2_and_nabla4_to_vn_switcher( - area_edge: Field[[EdgeDim], float], - kh_smag_e: Field[[EdgeDim, KDim], float], - z_nabla2_e: Field[[EdgeDim, KDim], float], - z_nabla4_e2: Field[[EdgeDim, KDim], float], - diff_multfac_vn: Field[[KDim], float], - nudgecoeff_e: Field[[EdgeDim], float], - vn: Field[[EdgeDim, KDim], float], - nudgezone_diff: float, - limited_area: bool, -) -> Field[[EdgeDim, KDim], float]: - # TODO: Use normal if-else instead - vn = ( - _apply_nabla2_and_nabla4_to_vn( - area_edge, - kh_smag_e, - z_nabla2_e, - z_nabla4_e2, - diff_multfac_vn, - nudgecoeff_e, - vn, - nudgezone_diff, - ) - if limited_area - else _apply_nabla2_and_nabla4_global_to_vn( - area_edge, - kh_smag_e, - z_nabla2_e, - z_nabla4_e2, - diff_multfac_vn, - vn, - ) - ) - - return vn - - @field_operator def _apply_diffusion_to_vn( u_vert: Field[[VertexDim, KDim], float], @@ -85,6 +47,7 @@ def _apply_diffusion_to_vn( start_2nd_nudge_line_idx_e: int32, limited_area: bool, ) -> Field[[EdgeDim, KDim], float]: + z_nabla4_e2 = _calculate_nabla4( u_vert, v_vert, @@ -95,22 +58,33 @@ def _apply_diffusion_to_vn( inv_primal_edge_length, ) - vn = where( - horz_idx >= start_2nd_nudge_line_idx_e, - _apply_nabla2_and_nabla4_to_vn_switcher( + # TODO: Use if-else statement instead + vn = ( + where( + horz_idx >= start_2nd_nudge_line_idx_e, + _apply_nabla2_and_nabla4_to_vn( + area_edge, + kh_smag_e, + z_nabla2_e, + z_nabla4_e2, + diff_multfac_vn, + nudgecoeff_e, + vn, + nudgezone_diff, + ), + _apply_nabla2_to_vn_in_lateral_boundary( + z_nabla2_e, area_edge, vn, fac_bdydiff_v + ), + ) + if limited_area + else _apply_nabla2_and_nabla4_global_to_vn( area_edge, kh_smag_e, z_nabla2_e, z_nabla4_e2, diff_multfac_vn, - nudgecoeff_e, vn, - nudgezone_diff, - limited_area, - ), - _apply_nabla2_to_vn_in_lateral_boundary( - z_nabla2_e, area_edge, vn, fac_bdydiff_v - ), + ) ) return vn From 478688a2cffbba628aa3e672ccc8861639763a0f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Wed, 31 May 2023 18:07:20 +0200 Subject: [PATCH 21/49] Drop w from argument list of fused stenil, write conditions always with < and <= --- .../icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py | 2 +- ..._and_compute_horizontal_gradients_for_turbulance.py | 10 ++++------ 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py index fa464cbca2..0c62d06d95 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py @@ -61,7 +61,7 @@ def _apply_diffusion_to_vn( # TODO: Use if-else statement instead vn = ( where( - horz_idx >= start_2nd_nudge_line_idx_e, + start_2nd_nudge_line_idx_e <= horz_idx, _apply_nabla2_and_nabla4_to_vn( area_edge, kh_smag_e, diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py index da1d164d01..4019b4a083 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance.py @@ -32,7 +32,6 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( geofac_grg_x: Field[[CellDim, C2E2CODim], float], geofac_grg_y: Field[[CellDim, C2E2CODim], float], w_old: Field[[CellDim, KDim], float], - w: Field[[CellDim, KDim], float], dwdx: Field[[CellDim, KDim], float], dwdy: Field[[CellDim, KDim], float], diff_multfac_w: float, @@ -51,7 +50,7 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( vert_idx = broadcast(vert_idx, (CellDim, KDim)) dwdx, dwdy = where( - vert_idx > int32(0), + int32(0) < vert_idx, _calculate_horizontal_gradients_for_turbulence( w_old, geofac_grg_x, geofac_grg_y ), @@ -61,15 +60,15 @@ def _apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( z_nabla2_c = _calculate_nabla2_for_w(w_old, geofac_n2s) w = where( - (horz_idx >= interior_idx) & (horz_idx < halo_idx), + (interior_idx <= horz_idx) & (horz_idx < halo_idx), _apply_nabla2_to_w(area, z_nabla2_c, geofac_n2s, w_old, diff_multfac_w), w_old, ) w = where( - (vert_idx > int32(0)) + (int32(0) < vert_idx) & (vert_idx < nrdmax) - & (horz_idx >= interior_idx) + & (interior_idx <= horz_idx) & (horz_idx < halo_idx), _apply_nabla2_to_w_in_upper_damping_layer( w, diff_multfac_n2w, area, z_nabla2_c @@ -104,7 +103,6 @@ def apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance( geofac_grg_x, geofac_grg_y, w_old, - w, dwdx, dwdy, diff_multfac_w, From d2f42f665e9448c67092879e781e017c54c3d2db Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Tue, 6 Jun 2023 13:24:15 +0200 Subject: [PATCH 22/49] Fix domains for global case in fused vn stencil --- .../atm_dyn_iconam/apply_diffusion_to_vn.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py index 0c62d06d95..a635d7c5a5 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py @@ -77,12 +77,17 @@ def _apply_diffusion_to_vn( ), ) if limited_area - else _apply_nabla2_and_nabla4_global_to_vn( - area_edge, - kh_smag_e, - z_nabla2_e, - z_nabla4_e2, - diff_multfac_vn, + else + where( + start_2nd_nudge_line_idx_e <= horz_idx, + _apply_nabla2_and_nabla4_global_to_vn( + area_edge, + kh_smag_e, + z_nabla2_e, + z_nabla4_e2, + diff_multfac_vn, + vn, + ), vn, ) ) From 7df42fdbb386b39661a770c09f1015540423ec8e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Tue, 6 Jun 2023 13:25:22 +0200 Subject: [PATCH 23/49] Deleted example stencil with output fields of different type --- ...fused_mo_nh_diffusion_stencil_01_02_03_rbf | 143 ------------------ 1 file changed, 143 deletions(-) delete mode 100644 atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf deleted file mode 100644 index e33e884f8e..0000000000 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/fused_mo_nh_diffusion_stencil_01_02_03_rbf +++ /dev/null @@ -1,143 +0,0 @@ -# ICON4Py - ICON inspired code in Python and GT4Py -# -# Copyright (c) 2022, ETH Zurich and MeteoSwiss -# All rights reserved. -# -# This file is free software: you can redistribute it and/or modify it under -# the terms of the GNU General Public License as published by the -# Free Software Foundation, either version 3 of the License, or any later -# version. See the LICENSE.txt file at the top-level directory of this -# distribution for a copy of the license or check . -# -# SPDX-License-Identifier: GPL-3.0-or-later - -from functional.ffront.decorator import field_operator, program -from functional.ffront.fbuiltins import Field - -from icon4py.atm_dyn_iconam.mo_intp_rbf_rbf_vec_interpol_vertex import ( - _mo_intp_rbf_rbf_vec_interpol_vertex, -) -from icon4py.atm_dyn_iconam.calculate_nabla2_and_smag_coefficients_for_vn import ( - _calculate_nabla2_and_smag_coefficients_for_vn, -) -from icon4py.atm_dyn_iconam.temporary_fields_for_turbulance_diagnostics import ( - _temporary_fields_for_turbulance_diagnostics, -) -from icon4py.atm_dyn_iconam.calculate_diagnostics_for_turbulance import ( - _calculate_diagnostics_for_turbulance, -) -from icon4py.common.dimension import ( - C2EDim, - CellDim, - ECVDim, - EdgeDim, - KDim, - V2EDim, - VertexDim, -) - - -@field_operator -def _fused_mo_nh_diffusion_stencil_01_02_03_rbf( - diff_multfac_smag: Field[[KDim], float], - tangent_orientation: Field[[EdgeDim], float], - inv_primal_edge_length: Field[[EdgeDim], float], - inv_vert_vert_length: Field[[EdgeDim], float], - u_vert_old: Field[[VertexDim, KDim], float], - v_vert_old: Field[[VertexDim, KDim], float], - primal_normal_vert_x: Field[[ECVDim], float], - primal_normal_vert_y: Field[[ECVDim], float], - dual_normal_vert_x: Field[[ECVDim], float], - dual_normal_vert_y: Field[[ECVDim], float], - vn: Field[[EdgeDim, KDim], float], - smag_limit: Field[[KDim], float], - smag_offset: float, - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], - wgtfac_c: Field[[CellDim, KDim], float], - ptr_coeff_1: Field[[VertexDim, V2EDim], float], - ptr_coeff_2: Field[[VertexDim, V2EDim], float], -) -> tuple[ - Field[[EdgeDim, KDim], float], - Field[[CellDim, KDim], float], - Field[[CellDim, KDim], float], - Field[[VertexDim, KDim], float], - Field[[VertexDim, KDim], float], -]: - - kh_smag_e, kh_smag_ec, z_nabla2_e = _calculate_nabla2_and_smag_coefficients_for_vn( - diff_multfac_smag, - tangent_orientation, - inv_primal_edge_length, - inv_vert_vert_length, - u_vert_old, - v_vert_old, - primal_normal_vert_x, - primal_normal_vert_y, - dual_normal_vert_x, - dual_normal_vert_y, - vn, - smag_limit, - smag_offset, - ) - - kh_c, div = _temporary_fields_for_turbulance_diagnostics( - kh_smag_ec, vn, e_bln_c_s, geofac_div, diff_multfac_smag - ) - - div_ic, hdef_ic = _calculate_diagnostics_for_turbulance(div, kh_c, wgtfac_c) - - u_vert, v_vert = _mo_intp_rbf_rbf_vec_interpol_vertex( - z_nabla2_e, ptr_coeff_1, ptr_coeff_2 - ) - - return kh_smag_e, div_ic, hdef_ic, u_vert, v_vert - - -@program -def fused_mo_nh_diffusion_stencil_01_02_03_rbf( - diff_multfac_smag: Field[[KDim], float], - tangent_orientation: Field[[EdgeDim], float], - inv_primal_edge_length: Field[[EdgeDim], float], - inv_vert_vert_length: Field[[EdgeDim], float], - u_vert_old: Field[[VertexDim, KDim], float], - v_vert_old: Field[[VertexDim, KDim], float], - primal_normal_vert_x: Field[[ECVDim], float], - primal_normal_vert_y: Field[[ECVDim], float], - dual_normal_vert_x: Field[[ECVDim], float], - dual_normal_vert_y: Field[[ECVDim], float], - vn: Field[[EdgeDim, KDim], float], - smag_limit: Field[[KDim], float], - smag_offset: float, - e_bln_c_s: Field[[CellDim, C2EDim], float], - geofac_div: Field[[CellDim, C2EDim], float], - wgtfac_c: Field[[CellDim, KDim], float], - ptr_coeff_1: Field[[VertexDim, V2EDim], float], - ptr_coeff_2: Field[[VertexDim, V2EDim], float], - kh_smag_e: Field[[EdgeDim, KDim], float], - div_ic: Field[[CellDim, KDim], float], - hdef_ic: Field[[CellDim, KDim], float], - u_vert: Field[[VertexDim, KDim], float], - v_vert: Field[[VertexDim, KDim], float], -): - _fused_mo_nh_diffusion_stencil_01_02_03_rbf( - diff_multfac_smag, - tangent_orientation, - inv_primal_edge_length, - inv_vert_vert_length, - u_vert_old, - v_vert_old, - primal_normal_vert_x, - primal_normal_vert_y, - dual_normal_vert_x, - dual_normal_vert_y, - vn, - smag_limit, - smag_offset, - e_bln_c_s, - geofac_div, - wgtfac_c, - ptr_coeff_1, - ptr_coeff_2, - out=(kh_smag_e, div_ic, hdef_ic, u_vert, v_vert), - ) From a8844806740ae7e55a25678157f3491a1dae804d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Christoph=20M=C3=BCller?= Date: Tue, 6 Jun 2023 13:34:28 +0200 Subject: [PATCH 24/49] Apply pre-commit --- .../src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py index a635d7c5a5..64f79b4d83 100644 --- a/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py +++ b/atm_dyn_iconam/src/icon4py/atm_dyn_iconam/apply_diffusion_to_vn.py @@ -77,8 +77,7 @@ def _apply_diffusion_to_vn( ), ) if limited_area - else - where( + else where( start_2nd_nudge_line_idx_e <= horz_idx, _apply_nabla2_and_nabla4_global_to_vn( area_edge, From 084e35cfdcf2eadd2c085107af92757cdc94397c Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Mon, 3 Jul 2023 14:34:27 +0200 Subject: [PATCH 25/49] allow to uncoment with !! --- tools/src/icon4pytools/liskov/parsing/scan.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/src/icon4pytools/liskov/parsing/scan.py b/tools/src/icon4pytools/liskov/parsing/scan.py index 0fb9879749..0c04375538 100644 --- a/tools/src/icon4pytools/liskov/parsing/scan.py +++ b/tools/src/icon4pytools/liskov/parsing/scan.py @@ -62,7 +62,7 @@ def __call__(self, data: Any = None) -> list[ts.RawDirective]: scanned_directives = [] lines = f.readlines() for lnumber, string in enumerate(lines): - if ts.DIRECTIVE_IDENT in string: + if ts.DIRECTIVE_IDENT in string and not '!' + ts.DIRECTIVE_IDENT in string: stripped = string.strip() eol = stripped[-1] scanned = Scanned(string, lnumber) From 115a2300f12565fc82995ecadacc4ca093380d7e Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 4 Jul 2023 09:31:33 +0200 Subject: [PATCH 26/49] allow !$DSL to be commented out and increase verbosity of error msg --- tools/src/icon4pytools/liskov/parsing/scan.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tools/src/icon4pytools/liskov/parsing/scan.py b/tools/src/icon4pytools/liskov/parsing/scan.py index 0c04375538..bc4cdbeb11 100644 --- a/tools/src/icon4pytools/liskov/parsing/scan.py +++ b/tools/src/icon4pytools/liskov/parsing/scan.py @@ -62,7 +62,7 @@ def __call__(self, data: Any = None) -> list[ts.RawDirective]: scanned_directives = [] lines = f.readlines() for lnumber, string in enumerate(lines): - if ts.DIRECTIVE_IDENT in string and not '!' + ts.DIRECTIVE_IDENT in string: + if string.strip().startswith(ts.DIRECTIVE_IDENT): stripped = string.strip() eol = stripped[-1] scanned = Scanned(string, lnumber) @@ -77,12 +77,12 @@ def __call__(self, data: Any = None) -> list[ts.RawDirective]: if ts.DIRECTIVE_IDENT not in next_line: raise DirectiveSyntaxError( f"Error in directive on line number: {lnumber + 1}\n Invalid use of & in single line " - f"directive. " + f"directive in file {self.input_filepath} ." ) continue case _: raise DirectiveSyntaxError( - f"Error in directive on line number: {lnumber + 1}\n Used invalid end of line character." + f"Error in directive on line number: {lnumber + 1}\n Used invalid end of line characterat in file {self.input_filepath} ." ) logger.info(f"Scanning for directives at {self.input_filepath}") return directives From 29f64345cd192149019089ecc52931a8357467f5 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 11 Jul 2023 12:09:44 +0200 Subject: [PATCH 27/49] improve scan --- tools/src/icon4pytools/liskov/parsing/scan.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tools/src/icon4pytools/liskov/parsing/scan.py b/tools/src/icon4pytools/liskov/parsing/scan.py index bc4cdbeb11..95670938e9 100644 --- a/tools/src/icon4pytools/liskov/parsing/scan.py +++ b/tools/src/icon4pytools/liskov/parsing/scan.py @@ -61,11 +61,11 @@ def __call__(self, data: Any = None) -> list[ts.RawDirective]: with self.input_filepath.open() as f: scanned_directives = [] lines = f.readlines() - for lnumber, string in enumerate(lines): - if string.strip().startswith(ts.DIRECTIVE_IDENT): - stripped = string.strip() + for lnumber, line in enumerate(lines): + stripped = line.strip() + if stripped.startswith(ts.DIRECTIVE_IDENT): eol = stripped[-1] - scanned = Scanned(string, lnumber) + scanned = Scanned(line, lnumber) scanned_directives.append(scanned) match eol: From 9083d6950e6a2ffa63de9501df06707718e956ea Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 19 Jul 2023 08:35:36 +0200 Subject: [PATCH 28/49] WIP fused verification working --- tools/src/icon4pytools/liskov/cli.py | 12 +- .../liskov/codegen/integration/deserialise.py | 140 ++++++++++++---- .../liskov/codegen/integration/generate.py | 41 +++++ .../liskov/codegen/integration/interface.py | 37 ++++- .../liskov/codegen/integration/template.py | 153 +++++++++++++++++- .../src/icon4pytools/liskov/external/gt4py.py | 18 ++- .../src/icon4pytools/liskov/parsing/parse.py | 17 ++ .../icon4pytools/liskov/parsing/transform.py | 77 +++++++++ .../liskov/pipeline/collection.py | 15 +- 9 files changed, 462 insertions(+), 48 deletions(-) create mode 100644 tools/src/icon4pytools/liskov/parsing/transform.py diff --git a/tools/src/icon4pytools/liskov/cli.py b/tools/src/icon4pytools/liskov/cli.py index a16e0a428e..4d5d738d84 100644 --- a/tools/src/icon4pytools/liskov/cli.py +++ b/tools/src/icon4pytools/liskov/cli.py @@ -12,13 +12,14 @@ # SPDX-License-Identifier: GPL-3.0-or-later import pathlib +import sys import click from icon4pytools.common.logger import setup_logger from icon4pytools.liskov.external.exceptions import MissingCommandError from icon4pytools.liskov.pipeline.collection import ( - load_gt4py_stencils, + process_stencils, parse_fortran_file, run_code_generation, ) @@ -50,6 +51,11 @@ def main(ctx): is_flag=True, help="Add metadata header with information about program.", ) +@click.option( + "--fused/--unfused", + "-f/-u", default=True, + help="Add metadata header with information about program.", +) @click.argument( "input_path", type=click.Path(exists=True, dir_okay=False, resolve_path=True, path_type=pathlib.Path), @@ -58,10 +64,10 @@ def main(ctx): "output_path", type=click.Path(dir_okay=False, resolve_path=True, path_type=pathlib.Path), ) -def integrate(input_path, output_path, profile, metadatagen): +def integrate(input_path, output_path, fused, profile, metadatagen): mode = "integration" iface = parse_fortran_file(input_path, output_path, mode) - iface_gt4py = load_gt4py_stencils(iface) + iface_gt4py = process_stencils(iface, fused) run_code_generation( input_path, output_path, diff --git a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py index 1884190504..f8e59f6843 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py @@ -11,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +import sys from typing import Any, Optional, Protocol, Type import icon4pytools.liskov.parsing.parse @@ -21,6 +22,7 @@ DeclareData, EndCreateData, EndIfData, + EndDeleteData, EndProfileData, EndStencilData, FieldAssociationData, @@ -28,8 +30,11 @@ InsertData, IntegrationCodeInterface, StartCreateData, + StartDeleteData, StartProfileData, StartStencilData, + StartFusedStencilData, + EndFusedStencilData, UnusedDirective, ) from icon4pytools.liskov.codegen.shared.deserialise import Deserialiser @@ -134,6 +139,14 @@ class EndProfileDataFactory(OptionalMultiUseDataFactory): dtype: Type[EndProfileData] = EndProfileData +class EndDeleteDataFactory(OptionalMultiUseDataFactory): + directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.EndDelete + dtype: Type[EndDeleteData] = EndDeleteData + +class StartDeleteDataFactory(OptionalMultiUseDataFactory): + directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.StartDelete + dtype: Type[StartDeleteData] = StartDeleteData + class StartCreateDataFactory(DataFactoryBase): directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.StartCreate dtype: Type[StartCreateData] = StartCreateData @@ -229,47 +242,28 @@ def __call__(self, parsed: ts.ParsedDict) -> list[EndStencilData]: return deserialised -class StartStencilDataFactory(DataFactoryBase): - directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.StartStencil - dtype: Type[StartStencilData] = StartStencilData - - def __call__(self, parsed: ts.ParsedDict) -> list[StartStencilData]: - """Create and return a list of StartStencilData objects from the parsed directives. - - Args: - parsed (ParsedDict): Dictionary of parsed directives and their associated content. +class EndFusedStencilDataFactory(DataFactoryBase): + directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.EndFusedStencil + dtype: Type[EndFusedStencilData] = EndFusedStencilData - Returns: - List[StartStencilData]: List of StartStencilData objects created from the parsed directives. - """ + def __call__(self, parsed: ts.ParsedDict) -> list[EndFusedStencilData]: deserialised = [] - field_dimensions = flatten_list_of_dicts( - [DeclareDataFactory.get_field_dimensions(dim) for dim in parsed["content"]["Declare"]] - ) - directives = extract_directive(parsed["directives"], self.directive_cls) - for i, directive in enumerate(directives): - named_args = parsed["content"]["StartStencil"][i] - acc_present = string_to_bool(pop_item_from_dict(named_args, "accpresent", "true")) - mergecopy = string_to_bool(pop_item_from_dict(named_args, "mergecopy", "false")) - copies = string_to_bool(pop_item_from_dict(named_args, "copies", "true")) + extracted = extract_directive(parsed["directives"], self.directive_cls) + for i, directive in enumerate(extracted): + named_args = parsed["content"]["EndFusedStencil"][i] stencil_name = _extract_stencil_name(named_args, directive) - bounds = self._make_bounds(named_args) - fields = self._make_fields(named_args, field_dimensions) - fields_w_tolerance = self._update_tolerances(named_args, fields) - deserialised.append( self.dtype( name=stencil_name, - fields=fields_w_tolerance, - bounds=bounds, startln=directive.startln, - acc_present=acc_present, - mergecopy=mergecopy, - copies=copies, ) ) return deserialised + + +class StartStencilDataFactoryBase(DataFactoryBase): + @staticmethod def _make_bounds(named_args: dict) -> BoundsData: """Extract stencil bounds from directive arguments.""" @@ -356,6 +350,88 @@ def _update_tolerances( setattr(f, tol, association) return fields +class StartStencilDataFactory(StartStencilDataFactoryBase): + directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.StartStencil + dtype: Type[StartStencilData] = StartStencilData + + def __call__(self, parsed: ts.ParsedDict) -> list[StartStencilData]: + """Create and return a list of StartStencilData objects from the parsed directives. + + Args: + parsed (ParsedDict): Dictionary of parsed directives and their associated content. + + Returns: + List[StartStencilData]: List of StartStencilData objects created from the parsed directives. + """ + deserialised = [] + field_dimensions = flatten_list_of_dicts( + [DeclareDataFactory.get_field_dimensions(dim) for dim in parsed["content"]["Declare"]] + ) + directives = extract_directive(parsed["directives"], self.directive_cls) + for i, directive in enumerate(directives): + named_args = parsed["content"]["StartStencil"][i] + acc_present = string_to_bool(pop_item_from_dict(named_args, "accpresent", "true")) + mergecopy = string_to_bool(pop_item_from_dict(named_args, "mergecopy", "false")) + copies = string_to_bool(pop_item_from_dict(named_args, "copies", "true")) + stencil_name = _extract_stencil_name(named_args, directive) + bounds = self._make_bounds(named_args) + fields = self._make_fields(named_args, field_dimensions) + fields_w_tolerance = self._update_tolerances(named_args, fields) + + deserialised.append( + self.dtype( + name=stencil_name, + fields=fields_w_tolerance, + bounds=bounds, + startln=directive.startln, + acc_present=acc_present, + mergecopy=mergecopy, + copies=copies, + ) + ) + return deserialised + + + +class StartFusedStencilDataFactory(StartStencilDataFactoryBase): + directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.StartFusedStencil + dtype: Type[StartFusedStencilData] = StartFusedStencilData + + def __call__(self, parsed: ts.ParsedDict) -> list[StartFusedStencilData]: + """Create and return a list of StartStencilData objects from the parsed directives. + + Args: + parsed (ParsedDict): Dictionary of parsed directives and their associated content. + + Returns: + List[StartStencilData]: List of StartStencilData objects created from the parsed directives. + """ + deserialised = [] + field_dimensions = flatten_list_of_dicts( + [DeclareDataFactory.get_field_dimensions(dim) for dim in parsed["content"]["Declare"]] + ) + directives = extract_directive(parsed["directives"], self.directive_cls) + for i, directive in enumerate(directives): + named_args = parsed["content"]["StartFusedStencil"][i] + acc_present = string_to_bool(pop_item_from_dict(named_args, "accpresent", "true")) + stencil_name = _extract_stencil_name(named_args, directive) + bounds = self._make_bounds(named_args) + fields = self._make_fields(named_args, field_dimensions) + fields_w_tolerance = self._update_tolerances(named_args, fields) + + deserialised.append( + self.dtype( + name=stencil_name, + fields=fields_w_tolerance, + bounds=bounds, + startln=directive.startln, + acc_present=acc_present, + ) + ) + return deserialised + + + class InsertDataFactory(DataFactoryBase): directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.Insert @@ -380,6 +456,10 @@ class IntegrationCodeDeserialiser(Deserialiser): "Declare": DeclareDataFactory(), "StartStencil": StartStencilDataFactory(), "EndStencil": EndStencilDataFactory(), + "StartFusedStencil": StartFusedStencilDataFactory(), + "EndFusedStencil": EndFusedStencilDataFactory(), + "StartDelete": StartDeleteDataFactory(), + "EndDelete": EndDeleteDataFactory(), "EndIf": EndIfDataFactory(), "StartProfile": StartProfileDataFactory(), "EndProfile": EndProfileDataFactory(), diff --git a/tools/src/icon4pytools/liskov/codegen/integration/generate.py b/tools/src/icon4pytools/liskov/codegen/integration/generate.py index 9705435be3..859709b7db 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/generate.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/generate.py @@ -11,6 +11,7 @@ # # SPDX-License-Identifier: GPL-3.0-or-later +import sys from typing import Any from icon4pytools.common.logger import setup_logger @@ -30,6 +31,8 @@ EndProfileStatementGenerator, EndStencilStatement, EndStencilStatementGenerator, + EndFusedStencilStatement, + EndFusedStencilStatementGenerator, ImportsStatement, ImportsStatementGenerator, InsertStatement, @@ -42,6 +45,8 @@ StartProfileStatementGenerator, StartStencilStatement, StartStencilStatementGenerator, + StartFusedStencilStatement, + StartFusedStencilStatementGenerator, ) from icon4pytools.liskov.codegen.shared.generate import CodeGenerator from icon4pytools.liskov.codegen.shared.types import GeneratedCode @@ -75,6 +80,8 @@ def __call__(self, data: Any = None) -> list[GeneratedCode]: self._generate_declare() self._generate_start_stencil() self._generate_end_stencil() + self._generate_start_fused_stencil() + self._generate_end_fused_stencil() self._generate_endif() self._generate_profile() self._generate_insert() @@ -148,6 +155,24 @@ def _generate_start_stencil(self) -> None: ) i += 1 + def _generate_start_fused_stencil(self) -> None: + """Generate f90 integration code surrounding a stencil. + + Args: + profile: A boolean indicating whether to include profiling calls in the generated code. + """ + + for stencil in self.interface.StartFusedStencil: + # while i < len(self.interface.StartStencil): + # stencil = self.interface.StartStencil[i] + logger.info(f"Generating START FUSED statement for {stencil.name}") + self._generate( + StartFusedStencilStatement, + StartFusedStencilStatementGenerator, + stencil.startln, + stencil_data=stencil, + ) + def _generate_end_stencil(self) -> None: """Generate f90 integration code surrounding a stencil. @@ -167,6 +192,22 @@ def _generate_end_stencil(self) -> None: noaccenddata=self.interface.EndStencil[i].noaccenddata, ) + def _generate_end_fused_stencil(self) -> None: + """Generate f90 integration code surrounding a stencil. + + Args: + profile: A boolean indicating whether to include profiling calls in the generated code. + """ + for i, stencil in enumerate(self.interface.StartFusedStencil): + logger.info(f"Generating END Fused statement for {stencil.name}") + self._generate( + EndFusedStencilStatement, + EndFusedStencilStatementGenerator, + self.interface.EndFusedStencil[i].startln, + stencil_data=stencil, + ) + # sys.exit(1) + def _generate_imports(self) -> None: """Generate f90 code for import statements.""" logger.info("Generating IMPORT statement.") diff --git a/tools/src/icon4pytools/liskov/codegen/integration/interface.py b/tools/src/icon4pytools/liskov/codegen/integration/interface.py index b48e0d6a6b..1b4391b32e 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/interface.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/interface.py @@ -76,25 +76,50 @@ class StartProfileData(CodeGenInput): class EndProfileData(CodeGenInput): ... - @dataclass -class StartStencilData(CodeGenInput): +class StartBasicStencilData(CodeGenInput): name: str fields: list[FieldAssociationData] - bounds: BoundsData acc_present: Optional[bool] + bounds: BoundsData + + +@dataclass +class StartStencilData(StartBasicStencilData): mergecopy: Optional[bool] copies: Optional[bool] +@dataclass +class StartFusedStencilData(StartBasicStencilData): + ... @dataclass -class EndStencilData(CodeGenInput): +class EndStencilBaseData(CodeGenInput): name: str + +@dataclass +class EndStencilData(EndStencilBaseData): noendif: Optional[bool] noprofile: Optional[bool] noaccenddata: Optional[bool] +@dataclass +class EndFusedStencilData(EndStencilBaseData): + ... + + +@dataclass +class StartDeleteData(CodeGenInput): + def __init__(self, startStencil: StartStencilData): + self.startln = startStencil.startln + + +@dataclass +class EndDeleteData(CodeGenInput): + def __init__(self, endStencil: EndStencilData): + self.startln = endStencil.startln + @dataclass class InsertData(CodeGenInput): content: str @@ -104,6 +129,10 @@ class InsertData(CodeGenInput): class IntegrationCodeInterface: StartStencil: Sequence[StartStencilData] EndStencil: Sequence[EndStencilData] + StartFusedStencil: Sequence[StartFusedStencilData] + EndFusedStencil: Sequence[EndFusedStencilData] + StartDelete: Sequence[StartDeleteData] | UnusedDirective + EndDelete: Sequence[EndDeleteData] | UnusedDirective Declare: Sequence[DeclareData] Imports: ImportsData StartCreate: Sequence[StartCreateData] | UnusedDirective diff --git a/tools/src/icon4pytools/liskov/codegen/integration/template.py b/tools/src/icon4pytools/liskov/codegen/integration/template.py index dba9543864..54126f4b02 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/template.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/template.py @@ -20,7 +20,7 @@ from gt4py.eve.codegen import TemplatedGenerator from icon4pytools.liskov.codegen.integration.exceptions import UndeclaredFieldError -from icon4pytools.liskov.codegen.integration.interface import DeclareData, StartStencilData +from icon4pytools.liskov.codegen.integration.interface import DeclareData, StartStencilData, StartFusedStencilData from icon4pytools.liskov.external.metadata import CodeMetadata @@ -120,7 +120,97 @@ def __post_init__(self) -> None: # type: ignore fields=[f for f in all_fields if f.rel_tol or f.abs_tol] ) +class EndFusedStencilStatement(eve.Node): + stencil_data: StartFusedStencilData + name: str = eve.datamodels.field(init=False) + input_fields: InputFields = eve.datamodels.field(init=False) + output_fields: OutputFields = eve.datamodels.field(init=False) + tolerance_fields: ToleranceFields = eve.datamodels.field(init=False) + bounds_fields: BoundsFields = eve.datamodels.field(init=False) + + def __post_init__(self) -> None: # type: ignore + all_fields = [Field(**asdict(f)) for f in self.stencil_data.fields] + self.bounds_fields = BoundsFields(**asdict(self.stencil_data.bounds)) + self.name = self.stencil_data.name + self.input_fields = InputFields(fields=[f for f in all_fields if f.inp]) + self.output_fields = OutputFields(fields=[f for f in all_fields if f.out]) + self.tolerance_fields = ToleranceFields( + fields=[f for f in all_fields if f.rel_tol or f.abs_tol] + ) + + +class EndFusedStencilStatementGenerator(TemplatedGenerator): + EndFusedStencilStatement = as_jinja( + """ + call wrap_run_{{ name }}( & + {{ input_fields }} + {{ output_fields }} + {{ tolerance_fields }} + {{ bounds_fields }} + + !$ACC END DATA + """ + ) + + InputFields = as_jinja( + """ + {%- for field in _this_node.fields %} + {%- if field.out %} + + {%- else %} + {{ field.variable }}={{ field.association }},& + {%- endif -%} + {%- endfor %} + """ + ) + + OutputFields = as_jinja( + """ + {%- for field in _this_node.fields %} + {{ field.variable }}={{ field.association }},& + {{ field.variable }}_before={{ field.variable }}_before{{ field.rh_index }},& + {%- endfor %} + """ + ) + + def visit_OutputFields(self, out: OutputFields) -> OutputFields: # type: ignore + for f in out.fields: # type: ignore + idx = render_index(f.dims) + split_idx = idx.split(",") + + if len(split_idx) >= 3: + split_idx[-1] = "1" + + f.rh_index = enclose_in_parentheses(",".join(split_idx)) + return self.generic_visit(out) + + ToleranceFields = as_jinja( + """ + {%- if _this_node.fields|length < 1 -%} + + {%- else -%} + + {%- for f in _this_node.fields -%} + {% if f.rel_tol %} + {{ f.variable }}_rel_tol={{ f.rel_tol }}, & + {%- endif -%} + {% if f.abs_tol %} + {{ f.variable }}_abs_tol={{ f.abs_tol }}, & + {% endif %} + {%- endfor -%} + + {%- endif -%} + """ + ) + + BoundsFields = as_jinja( + """vertical_lower={{ vlower }}, & + vertical_upper={{ vupper }}, & + horizontal_lower={{ hlower }}, & + horizontal_upper={{ hupper }}) + """ + ) class EndStencilStatementGenerator(TemplatedGenerator): EndStencilStatement = as_jinja( """ @@ -270,6 +360,44 @@ def make_copy_declaration(f: Field) -> CopyDeclaration: rh_index=rh_idx, ) +class StartFusedStencilStatement(eve.Node): + stencil_data: StartFusedStencilData + copy_declarations: list[CopyDeclaration] = eve.datamodels.field(init=False) + + def __post_init__(self) -> None: # type: ignore + all_fields = [Field(**asdict(f)) for f in self.stencil_data.fields] + self.copy_declarations = [self.make_copy_declaration(f) for f in all_fields if f.out] + self.acc_present = "PRESENT" if self.stencil_data.acc_present else "NONE" + + @staticmethod + def make_copy_declaration(f: Field) -> CopyDeclaration: + if f.dims is None: + raise UndeclaredFieldError(f"{f.variable} was not declared!") + + lh_idx = render_index(f.dims) + + # get length of association index + association_dims = get_array_dims(f.association).split(",") + n_association_dims = len(association_dims) + + offset = len(",".join(association_dims)) + 2 + truncated_association = f.association[:-offset] + + if n_association_dims > f.dims: + rh_idx = f"{lh_idx},{association_dims[-1]}" + else: + rh_idx = f"{lh_idx}" + + lh_idx = enclose_in_parentheses(lh_idx) + rh_idx = enclose_in_parentheses(rh_idx) + + return CopyDeclaration( + variable=f.variable, + association=truncated_association, + lh_index=lh_idx, + rh_index=rh_idx, + ) + def render_index(n: int) -> str: """ @@ -314,6 +442,29 @@ class StartStencilStatementGenerator(TemplatedGenerator): ) +class StartFusedStencilStatementGenerator(TemplatedGenerator): + StartFusedStencilStatement = as_jinja( + """ + + !$ACC DATA CREATE( & + {%- for d in _this_node.copy_declarations %} + !$ACC {{ d.variable }}_before {%- if not loop.last -%}, & {% else %} ) & {%- endif -%} + {%- endfor %} + !$ACC IF ( i_am_accel_node ) + + #ifdef __DSL_VERIFY + !$ACC KERNELS IF( i_am_accel_node ) DEFAULT({{ _this_node.acc_present }}) ASYNC(1) + {%- for d in _this_node.copy_declarations %} + {{ d.variable }}_before{{ d.lh_index }} = {{ d.association }}{{ d.rh_index }} + {%- endfor %} + !$ACC END KERNELS + #endif + + """ + ) + + + class ImportsStatement(eve.Node): stencils: list[StartStencilData] stencil_names: list[str] = eve.datamodels.field(init=False) diff --git a/tools/src/icon4pytools/liskov/external/gt4py.py b/tools/src/icon4pytools/liskov/external/gt4py.py index 4ad14ef9d3..f26cf02ebb 100644 --- a/tools/src/icon4pytools/liskov/external/gt4py.py +++ b/tools/src/icon4pytools/liskov/external/gt4py.py @@ -13,7 +13,7 @@ import importlib from inspect import getmembers -from typing import Any +from typing import Any, Sequence from gt4py.next.ffront.decorator import Program @@ -22,6 +22,7 @@ from icon4pytools.liskov.codegen.integration.interface import IntegrationCodeInterface from icon4pytools.liskov.external.exceptions import IncompatibleFieldError, UnknownStencilError from icon4pytools.liskov.pipeline.definition import Step +from icon4pytools.liskov.codegen.integration.interface import StartBasicStencilData logger = setup_logger(__name__) @@ -36,7 +37,13 @@ def __init__(self, parsed: IntegrationCodeInterface): def __call__(self, data: Any = None) -> IntegrationCodeInterface: logger.info("Updating parsed fields with data from icon4py stencils...") - for s in self.parsed.StartStencil: + self._set_in_out_field(self.parsed.StartStencil) + self._set_in_out_field(self.parsed.StartFusedStencil) + + return self.parsed + + def _set_in_out_field(self, startStencil: Sequence[StartBasicStencilData]) -> None: + for s in startStencil: gt4py_program = self._collect_icon4py_stencil(s.name) gt4py_stencil_info = get_stencil_info(gt4py_program) gt4py_fields = gt4py_stencil_info.fields @@ -44,12 +51,11 @@ def __call__(self, data: Any = None) -> IntegrationCodeInterface: try: field_info = gt4py_fields[f.variable] except KeyError: - raise IncompatibleFieldError( - f"Used field variable name that is incompatible with the expected field names defined in {s.name} in icon4pytools." - ) + error_msg = f"Used field variable name ({f.variable}) that is incompatible with the expected field names defined in {s.name} in icon4pytools." + raise IncompatibleFieldError(error_msg) f.out = field_info.out f.inp = field_info.inp - return self.parsed + def _collect_icon4py_stencil(self, stencil_name: str) -> Program: """Collect and return the ICON4PY stencil program with the given name.""" diff --git a/tools/src/icon4pytools/liskov/parsing/parse.py b/tools/src/icon4pytools/liskov/parsing/parse.py index 02759621a5..0ba79092d3 100644 --- a/tools/src/icon4pytools/liskov/parsing/parse.py +++ b/tools/src/icon4pytools/liskov/parsing/parse.py @@ -186,6 +186,11 @@ class StartStencil(WithArguments): class EndStencil(WithArguments): pattern = "END STENCIL" +class StartFusedStencil(WithArguments): + pattern = "START FUSED STENCIL" + +class EndFusedStencil(WithArguments): + pattern = "END FUSED STENCIL" class Declare(WithArguments): pattern = "DECLARE" @@ -215,6 +220,14 @@ class EndProfile(WithoutArguments): pattern = "END PROFILE" +class StartDelete(WithoutArguments): + pattern = 'START DELETE' + + +class EndDelete(WithoutArguments): + pattern = 'END DELETE' + + class Insert(FreeForm): pattern = "INSERT" @@ -222,6 +235,10 @@ class Insert(FreeForm): SUPPORTED_DIRECTIVES: Sequence[Type[ParsedDirective]] = [ StartStencil, EndStencil, + StartFusedStencil, + EndFusedStencil, + StartDelete, + EndDelete, Imports, Declare, StartCreate, diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py new file mode 100644 index 0000000000..405a536b14 --- /dev/null +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -0,0 +1,77 @@ +# ICON4Py - ICON inspired code in Python and GT4Py +# +# Copyright (c) 2022, ETH Zurich and MeteoSwiss +# All rights reserved. +# +# This file is free software: you can redistribute it and/or modify it under +# the terms of the GNU General Public License as published by the +# Free Software Foundation, either version 3 of the License, or any later +# version. See the LICENSE.txt file at the top-level directory of this +# distribution for a copy of the license or check . +# +# SPDX-License-Identifier: GPL-3.0-or-later +import collections +import shutil +import sys +from typing import Any + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Optional, Sequence, Type + +import icon4pytools.liskov.parsing.types as ts +from icon4pytools.common.logger import setup_logger +from icon4pytools.liskov.parsing.exceptions import UnsupportedDirectiveError +from icon4pytools.liskov.parsing.types import ParsedDirective, RawDirective +from icon4pytools.liskov.parsing.validation import VALIDATORS +from icon4pytools.liskov.pipeline.definition import Step +from icon4pytools.liskov.codegen.integration.interface import IntegrationCodeInterface, StartDeleteData, EndDeleteData, UnusedDirective + + + +logger = setup_logger(__name__) + + +class TransformFuseStencils(Step): + def __init__(self, parsed: IntegrationCodeInterface, fused: bool) -> None: + self.parsed = parsed + self.fused = fused + + def __call__(self, data: Any=None) -> ts.ParsedDict: + """Parse the directives and return a dictionary of parsed directives and their associated content. + + Returns: + ParsedType: Dictionary of parsed directives and their associated content. + """ + if self.fused: + logger.info(f"Transforming single stencils.") + + removeStartStencil = [] + removeEndStencil = [] + + for startFused, endFused in zip(self.parsed.StartFusedStencil, self.parsed.EndFusedStencil, strict=True): + for startSingle, endSingle in zip(self.parsed.StartStencil, self.parsed.EndStencil, strict=True): + # print('try: ', startSingle.name) + if( startFused.startln < startSingle.startln and startSingle.startln < endFused.startln and + startFused.startln < endSingle.startln and endSingle.startln < endFused.startln): + try: + self.parsed.StartDelete.append(StartDeleteData(startSingle)) + except AttributeError: + self.parsed.StartDelete = [StartDeleteData(startSingle)] + try: + self.parsed.EndDelete.append(EndDeleteData(endSingle)) + except AttributeError: + self.parsed.EndDelete = [EndDeleteData(endSingle)] + removeStartStencil.append(startSingle) + removeEndStencil.append(endSingle) + + self.parsed.StartStencil = [x for x in self.parsed.StartStencil if x not in removeStartStencil] + self.parsed.EndStencil = [x for x in self.parsed.EndStencil if x not in removeEndStencil] + + else: + logger.info(f"Removing fused stencils.") + + self.parsed.StartFusedStencil = [] + self.parsed.EndFusedStencil = [] + + return self.parsed \ No newline at end of file diff --git a/tools/src/icon4pytools/liskov/pipeline/collection.py b/tools/src/icon4pytools/liskov/pipeline/collection.py index 006fa8a0e9..9dbceb5da2 100644 --- a/tools/src/icon4pytools/liskov/pipeline/collection.py +++ b/tools/src/icon4pytools/liskov/pipeline/collection.py @@ -20,6 +20,8 @@ from icon4pytools.liskov.codegen.shared.write import CodegenWriter from icon4pytools.liskov.external.gt4py import UpdateFieldsWithGt4PyStencils from icon4pytools.liskov.parsing.parse import DirectivesParser +# from icon4pytools.liskov.parsing.process import RemoveFusedStencils, RemoveSingleStencils +from icon4pytools.liskov.parsing.transform import TransformFuseStencils from icon4pytools.liskov.parsing.scan import DirectivesScanner from icon4pytools.liskov.pipeline.definition import Step, linear_pipeline @@ -68,16 +70,21 @@ def parse_fortran_file( @linear_pipeline -def load_gt4py_stencils(parsed: IntegrationCodeInterface) -> list[Step]: - """Execute a pipeline to update fields of a IntegrationCodeInterface object with GT4Py stencils. +def process_stencils( + parsed: IntegrationCodeInterface, fused: bool +) -> list[Step]: + """Execute a pipeline to transform stencils to fused or unfused execution, and a pipeline Args: parsed: The input IntegrationCodeInterface object. + fused: The input mode if fused or unfused Returns: - The updated object with fields containing information from GT4Py stencils. + The updated and transformed object with fields containing information from GT4Py stencils. """ - return [UpdateFieldsWithGt4PyStencils(parsed)] + + return [TransformFuseStencils(parsed, fused), UpdateFieldsWithGt4PyStencils(parsed)] + @linear_pipeline From c0321e20986286f26de75c195821792430f8f10e Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 19 Jul 2023 08:46:30 +0200 Subject: [PATCH 29/49] styel up --- tools/src/icon4pytools/liskov/cli.py | 6 +- .../liskov/codegen/integration/deserialise.py | 15 ++-- .../liskov/codegen/integration/generate.py | 13 +-- .../liskov/codegen/integration/interface.py | 9 +- .../liskov/codegen/integration/template.py | 11 ++- .../src/icon4pytools/liskov/external/gt4py.py | 7 +- .../src/icon4pytools/liskov/parsing/parse.py | 7 +- .../icon4pytools/liskov/parsing/transform.py | 82 ++++++++++--------- .../liskov/pipeline/collection.py | 11 +-- 9 files changed, 84 insertions(+), 77 deletions(-) diff --git a/tools/src/icon4pytools/liskov/cli.py b/tools/src/icon4pytools/liskov/cli.py index 4d5d738d84..6d80439381 100644 --- a/tools/src/icon4pytools/liskov/cli.py +++ b/tools/src/icon4pytools/liskov/cli.py @@ -12,15 +12,14 @@ # SPDX-License-Identifier: GPL-3.0-or-later import pathlib -import sys import click from icon4pytools.common.logger import setup_logger from icon4pytools.liskov.external.exceptions import MissingCommandError from icon4pytools.liskov.pipeline.collection import ( - process_stencils, parse_fortran_file, + process_stencils, run_code_generation, ) @@ -53,7 +52,8 @@ def main(ctx): ) @click.option( "--fused/--unfused", - "-f/-u", default=True, + "-f/-u", + default=True, help="Add metadata header with information about program.", ) @click.argument( diff --git a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py index f8e59f6843..faa6857b87 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py @@ -11,7 +11,6 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -import sys from typing import Any, Optional, Protocol, Type import icon4pytools.liskov.parsing.parse @@ -21,8 +20,9 @@ BoundsData, DeclareData, EndCreateData, - EndIfData, EndDeleteData, + EndFusedStencilData, + EndIfData, EndProfileData, EndStencilData, FieldAssociationData, @@ -31,10 +31,9 @@ IntegrationCodeInterface, StartCreateData, StartDeleteData, + StartFusedStencilData, StartProfileData, StartStencilData, - StartFusedStencilData, - EndFusedStencilData, UnusedDirective, ) from icon4pytools.liskov.codegen.shared.deserialise import Deserialiser @@ -143,10 +142,12 @@ class EndDeleteDataFactory(OptionalMultiUseDataFactory): directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.EndDelete dtype: Type[EndDeleteData] = EndDeleteData + class StartDeleteDataFactory(OptionalMultiUseDataFactory): directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.StartDelete dtype: Type[StartDeleteData] = StartDeleteData + class StartCreateDataFactory(DataFactoryBase): directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.StartCreate dtype: Type[StartCreateData] = StartCreateData @@ -261,9 +262,7 @@ def __call__(self, parsed: ts.ParsedDict) -> list[EndFusedStencilData]: return deserialised - class StartStencilDataFactoryBase(DataFactoryBase): - @staticmethod def _make_bounds(named_args: dict) -> BoundsData: """Extract stencil bounds from directive arguments.""" @@ -350,6 +349,7 @@ def _update_tolerances( setattr(f, tol, association) return fields + class StartStencilDataFactory(StartStencilDataFactoryBase): directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.StartStencil dtype: Type[StartStencilData] = StartStencilData @@ -392,7 +392,6 @@ def __call__(self, parsed: ts.ParsedDict) -> list[StartStencilData]: return deserialised - class StartFusedStencilDataFactory(StartStencilDataFactoryBase): directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.StartFusedStencil dtype: Type[StartFusedStencilData] = StartFusedStencilData @@ -431,8 +430,6 @@ def __call__(self, parsed: ts.ParsedDict) -> list[StartFusedStencilData]: return deserialised - - class InsertDataFactory(DataFactoryBase): directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.Insert dtype: Type[InsertData] = InsertData diff --git a/tools/src/icon4pytools/liskov/codegen/integration/generate.py b/tools/src/icon4pytools/liskov/codegen/integration/generate.py index 859709b7db..216c3e4bc7 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/generate.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/generate.py @@ -11,7 +11,6 @@ # # SPDX-License-Identifier: GPL-3.0-or-later -import sys from typing import Any from icon4pytools.common.logger import setup_logger @@ -25,14 +24,14 @@ DeclareStatementGenerator, EndCreateStatement, EndCreateStatementGenerator, + EndFusedStencilStatement, + EndFusedStencilStatementGenerator, EndIfStatement, EndIfStatementGenerator, EndProfileStatement, EndProfileStatementGenerator, EndStencilStatement, EndStencilStatementGenerator, - EndFusedStencilStatement, - EndFusedStencilStatementGenerator, ImportsStatement, ImportsStatementGenerator, InsertStatement, @@ -41,12 +40,12 @@ MetadataStatementGenerator, StartCreateStatement, StartCreateStatementGenerator, + StartFusedStencilStatement, + StartFusedStencilStatementGenerator, StartProfileStatement, StartProfileStatementGenerator, StartStencilStatement, StartStencilStatementGenerator, - StartFusedStencilStatement, - StartFusedStencilStatementGenerator, ) from icon4pytools.liskov.codegen.shared.generate import CodeGenerator from icon4pytools.liskov.codegen.shared.types import GeneratedCode @@ -161,10 +160,7 @@ def _generate_start_fused_stencil(self) -> None: Args: profile: A boolean indicating whether to include profiling calls in the generated code. """ - for stencil in self.interface.StartFusedStencil: - # while i < len(self.interface.StartStencil): - # stencil = self.interface.StartStencil[i] logger.info(f"Generating START FUSED statement for {stencil.name}") self._generate( StartFusedStencilStatement, @@ -206,7 +202,6 @@ def _generate_end_fused_stencil(self) -> None: self.interface.EndFusedStencil[i].startln, stencil_data=stencil, ) - # sys.exit(1) def _generate_imports(self) -> None: """Generate f90 code for import statements.""" diff --git a/tools/src/icon4pytools/liskov/codegen/integration/interface.py b/tools/src/icon4pytools/liskov/codegen/integration/interface.py index 1b4391b32e..4dd1925316 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/interface.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/interface.py @@ -76,6 +76,7 @@ class StartProfileData(CodeGenInput): class EndProfileData(CodeGenInput): ... + @dataclass class StartBasicStencilData(CodeGenInput): name: str @@ -89,14 +90,17 @@ class StartStencilData(StartBasicStencilData): mergecopy: Optional[bool] copies: Optional[bool] + @dataclass class StartFusedStencilData(StartBasicStencilData): - ... + ... + @dataclass class EndStencilBaseData(CodeGenInput): name: str + @dataclass class EndStencilData(EndStencilBaseData): noendif: Optional[bool] @@ -106,7 +110,7 @@ class EndStencilData(EndStencilBaseData): @dataclass class EndFusedStencilData(EndStencilBaseData): - ... + ... @dataclass @@ -120,6 +124,7 @@ class EndDeleteData(CodeGenInput): def __init__(self, endStencil: EndStencilData): self.startln = endStencil.startln + @dataclass class InsertData(CodeGenInput): content: str diff --git a/tools/src/icon4pytools/liskov/codegen/integration/template.py b/tools/src/icon4pytools/liskov/codegen/integration/template.py index 54126f4b02..f733dfa777 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/template.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/template.py @@ -20,7 +20,11 @@ from gt4py.eve.codegen import TemplatedGenerator from icon4pytools.liskov.codegen.integration.exceptions import UndeclaredFieldError -from icon4pytools.liskov.codegen.integration.interface import DeclareData, StartStencilData, StartFusedStencilData +from icon4pytools.liskov.codegen.integration.interface import ( + DeclareData, + StartFusedStencilData, + StartStencilData, +) from icon4pytools.liskov.external.metadata import CodeMetadata @@ -120,6 +124,7 @@ def __post_init__(self) -> None: # type: ignore fields=[f for f in all_fields if f.rel_tol or f.abs_tol] ) + class EndFusedStencilStatement(eve.Node): stencil_data: StartFusedStencilData @@ -211,6 +216,8 @@ def visit_OutputFields(self, out: OutputFields) -> OutputFields: # type: ignore horizontal_upper={{ hupper }}) """ ) + + class EndStencilStatementGenerator(TemplatedGenerator): EndStencilStatement = as_jinja( """ @@ -360,6 +367,7 @@ def make_copy_declaration(f: Field) -> CopyDeclaration: rh_index=rh_idx, ) + class StartFusedStencilStatement(eve.Node): stencil_data: StartFusedStencilData copy_declarations: list[CopyDeclaration] = eve.datamodels.field(init=False) @@ -464,7 +472,6 @@ class StartFusedStencilStatementGenerator(TemplatedGenerator): ) - class ImportsStatement(eve.Node): stencils: list[StartStencilData] stencil_names: list[str] = eve.datamodels.field(init=False) diff --git a/tools/src/icon4pytools/liskov/external/gt4py.py b/tools/src/icon4pytools/liskov/external/gt4py.py index f26cf02ebb..f627708d44 100644 --- a/tools/src/icon4pytools/liskov/external/gt4py.py +++ b/tools/src/icon4pytools/liskov/external/gt4py.py @@ -19,10 +19,12 @@ from icon4pytools.common.logger import setup_logger from icon4pytools.icon4pygen.metadata import get_stencil_info -from icon4pytools.liskov.codegen.integration.interface import IntegrationCodeInterface +from icon4pytools.liskov.codegen.integration.interface import ( + IntegrationCodeInterface, + StartBasicStencilData, +) from icon4pytools.liskov.external.exceptions import IncompatibleFieldError, UnknownStencilError from icon4pytools.liskov.pipeline.definition import Step -from icon4pytools.liskov.codegen.integration.interface import StartBasicStencilData logger = setup_logger(__name__) @@ -56,7 +58,6 @@ def _set_in_out_field(self, startStencil: Sequence[StartBasicStencilData]) -> No f.out = field_info.out f.inp = field_info.inp - def _collect_icon4py_stencil(self, stencil_name: str) -> Program: """Collect and return the ICON4PY stencil program with the given name.""" err_counter = 0 diff --git a/tools/src/icon4pytools/liskov/parsing/parse.py b/tools/src/icon4pytools/liskov/parsing/parse.py index 0ba79092d3..7d60812b08 100644 --- a/tools/src/icon4pytools/liskov/parsing/parse.py +++ b/tools/src/icon4pytools/liskov/parsing/parse.py @@ -186,12 +186,15 @@ class StartStencil(WithArguments): class EndStencil(WithArguments): pattern = "END STENCIL" + class StartFusedStencil(WithArguments): pattern = "START FUSED STENCIL" + class EndFusedStencil(WithArguments): pattern = "END FUSED STENCIL" + class Declare(WithArguments): pattern = "DECLARE" @@ -221,11 +224,11 @@ class EndProfile(WithoutArguments): class StartDelete(WithoutArguments): - pattern = 'START DELETE' + pattern = "START DELETE" class EndDelete(WithoutArguments): - pattern = 'END DELETE' + pattern = "END DELETE" class Insert(FreeForm): diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index 405a536b14..6b96840f7b 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -10,23 +10,16 @@ # distribution for a copy of the license or check . # # SPDX-License-Identifier: GPL-3.0-or-later -import collections -import shutil -import sys from typing import Any -from dataclasses import dataclass, field -from pathlib import Path -from typing import Optional, Sequence, Type - import icon4pytools.liskov.parsing.types as ts from icon4pytools.common.logger import setup_logger -from icon4pytools.liskov.parsing.exceptions import UnsupportedDirectiveError -from icon4pytools.liskov.parsing.types import ParsedDirective, RawDirective -from icon4pytools.liskov.parsing.validation import VALIDATORS +from icon4pytools.liskov.codegen.integration.interface import ( + EndDeleteData, + IntegrationCodeInterface, + StartDeleteData, +) from icon4pytools.liskov.pipeline.definition import Step -from icon4pytools.liskov.codegen.integration.interface import IntegrationCodeInterface, StartDeleteData, EndDeleteData, UnusedDirective - logger = setup_logger(__name__) @@ -34,44 +27,55 @@ class TransformFuseStencils(Step): def __init__(self, parsed: IntegrationCodeInterface, fused: bool) -> None: - self.parsed = parsed - self.fused = fused + self.parsed = parsed + self.fused = fused - def __call__(self, data: Any=None) -> ts.ParsedDict: + def __call__(self, data: Any = None) -> ts.ParsedDict: """Parse the directives and return a dictionary of parsed directives and their associated content. Returns: ParsedType: Dictionary of parsed directives and their associated content. """ if self.fused: - logger.info(f"Transforming single stencils.") + logger.info("Transforming single stencils.") - removeStartStencil = [] - removeEndStencil = [] + removeStartStencil = [] + removeEndStencil = [] - for startFused, endFused in zip(self.parsed.StartFusedStencil, self.parsed.EndFusedStencil, strict=True): - for startSingle, endSingle in zip(self.parsed.StartStencil, self.parsed.EndStencil, strict=True): - # print('try: ', startSingle.name) - if( startFused.startln < startSingle.startln and startSingle.startln < endFused.startln and - startFused.startln < endSingle.startln and endSingle.startln < endFused.startln): - try: - self.parsed.StartDelete.append(StartDeleteData(startSingle)) - except AttributeError: - self.parsed.StartDelete = [StartDeleteData(startSingle)] - try: - self.parsed.EndDelete.append(EndDeleteData(endSingle)) - except AttributeError: - self.parsed.EndDelete = [EndDeleteData(endSingle)] - removeStartStencil.append(startSingle) - removeEndStencil.append(endSingle) + for startFused, endFused in zip( + self.parsed.StartFusedStencil, self.parsed.EndFusedStencil, strict=True + ): + for startSingle, endSingle in zip( + self.parsed.StartStencil, self.parsed.EndStencil, strict=True + ): + if ( + startFused.startln < startSingle.startln + and startSingle.startln < endFused.startln + and startFused.startln < endSingle.startln + and endSingle.startln < endFused.startln + ): + try: + self.parsed.StartDelete.append(StartDeleteData(startSingle)) + except AttributeError: + self.parsed.StartDelete = [StartDeleteData(startSingle)] + try: + self.parsed.EndDelete.append(EndDeleteData(endSingle)) + except AttributeError: + self.parsed.EndDelete = [EndDeleteData(endSingle)] + removeStartStencil.append(startSingle) + removeEndStencil.append(endSingle) - self.parsed.StartStencil = [x for x in self.parsed.StartStencil if x not in removeStartStencil] - self.parsed.EndStencil = [x for x in self.parsed.EndStencil if x not in removeEndStencil] + self.parsed.StartStencil = [ + x for x in self.parsed.StartStencil if x not in removeStartStencil + ] + self.parsed.EndStencil = [ + x for x in self.parsed.EndStencil if x not in removeEndStencil + ] else: - logger.info(f"Removing fused stencils.") + logger.info("Removing fused stencils.") - self.parsed.StartFusedStencil = [] - self.parsed.EndFusedStencil = [] + self.parsed.StartFusedStencil = [] + self.parsed.EndFusedStencil = [] - return self.parsed \ No newline at end of file + return self.parsed diff --git a/tools/src/icon4pytools/liskov/pipeline/collection.py b/tools/src/icon4pytools/liskov/pipeline/collection.py index 9dbceb5da2..ba9b776b3f 100644 --- a/tools/src/icon4pytools/liskov/pipeline/collection.py +++ b/tools/src/icon4pytools/liskov/pipeline/collection.py @@ -20,9 +20,8 @@ from icon4pytools.liskov.codegen.shared.write import CodegenWriter from icon4pytools.liskov.external.gt4py import UpdateFieldsWithGt4PyStencils from icon4pytools.liskov.parsing.parse import DirectivesParser -# from icon4pytools.liskov.parsing.process import RemoveFusedStencils, RemoveSingleStencils -from icon4pytools.liskov.parsing.transform import TransformFuseStencils from icon4pytools.liskov.parsing.scan import DirectivesScanner +from icon4pytools.liskov.parsing.transform import TransformFuseStencils from icon4pytools.liskov.pipeline.definition import Step, linear_pipeline @@ -70,10 +69,8 @@ def parse_fortran_file( @linear_pipeline -def process_stencils( - parsed: IntegrationCodeInterface, fused: bool -) -> list[Step]: - """Execute a pipeline to transform stencils to fused or unfused execution, and a pipeline +def process_stencils(parsed: IntegrationCodeInterface, fused: bool) -> list[Step]: + """Execute a pipeline to transform stencils to fused or unfused execution, and a pipeline. Args: parsed: The input IntegrationCodeInterface object. @@ -82,11 +79,9 @@ def process_stencils( Returns: The updated and transformed object with fields containing information from GT4Py stencils. """ - return [TransformFuseStencils(parsed, fused), UpdateFieldsWithGt4PyStencils(parsed)] - @linear_pipeline def run_code_generation( input_filepath: Path, From 98e6a805bb74e51e0527b5056f6bf366dbb487ba Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 19 Jul 2023 09:59:01 +0200 Subject: [PATCH 30/49] WIP substition mode --- .../liskov/codegen/integration/generate.py | 69 ++++++++++++------- .../liskov/codegen/integration/template.py | 13 ++++ 2 files changed, 57 insertions(+), 25 deletions(-) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/generate.py b/tools/src/icon4pytools/liskov/codegen/integration/generate.py index 216c3e4bc7..60b9bd0153 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/generate.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/generate.py @@ -42,6 +42,10 @@ StartCreateStatementGenerator, StartFusedStencilStatement, StartFusedStencilStatementGenerator, + StartDeleteStatement, + StartDeleteStatementGenerator, + EndDeleteStatement, + EndDeleteStatementGenerator, StartProfileStatement, StartProfileStatementGenerator, StartStencilStatement, @@ -81,6 +85,7 @@ def __call__(self, data: Any = None) -> list[GeneratedCode]: self._generate_end_stencil() self._generate_start_fused_stencil() self._generate_end_fused_stencil() + self._generate_delete() self._generate_endif() self._generate_profile() self._generate_insert() @@ -154,20 +159,6 @@ def _generate_start_stencil(self) -> None: ) i += 1 - def _generate_start_fused_stencil(self) -> None: - """Generate f90 integration code surrounding a stencil. - - Args: - profile: A boolean indicating whether to include profiling calls in the generated code. - """ - for stencil in self.interface.StartFusedStencil: - logger.info(f"Generating START FUSED statement for {stencil.name}") - self._generate( - StartFusedStencilStatement, - StartFusedStencilStatementGenerator, - stencil.startln, - stencil_data=stencil, - ) def _generate_end_stencil(self) -> None: """Generate f90 integration code surrounding a stencil. @@ -188,20 +179,48 @@ def _generate_end_stencil(self) -> None: noaccenddata=self.interface.EndStencil[i].noaccenddata, ) + def _generate_start_fused_stencil(self) -> None: + """Generate f90 integration code surrounding a fused stencil. + """ + if self.interface.StartFusedStencil != UnusedDirective: + for stencil in self.interface.StartFusedStencil: + logger.info(f"Generating START FUSED statement for {stencil.name}") + self._generate( + StartFusedStencilStatement, + StartFusedStencilStatementGenerator, + stencil.startln, + stencil_data=stencil, + ) + def _generate_end_fused_stencil(self) -> None: - """Generate f90 integration code surrounding a stencil. + """Generate f90 integration code surrounding a fused stencil. + """ + if self.interface.EndFusedStencil != UnusedDirective: + for i, stencil in enumerate(self.interface.StartFusedStencil): + logger.info(f"Generating END Fused statement for {stencil.name}") + self._generate( + EndFusedStencilStatement, + EndFusedStencilStatementGenerator, + self.interface.EndFusedStencil[i].startln, + stencil_data=stencil, + ) - Args: - profile: A boolean indicating whether to include profiling calls in the generated code. + def _generate_delete(self) -> None: + """Generate f90 integration code for delete section. """ - for i, stencil in enumerate(self.interface.StartFusedStencil): - logger.info(f"Generating END Fused statement for {stencil.name}") - self._generate( - EndFusedStencilStatement, - EndFusedStencilStatementGenerator, - self.interface.EndFusedStencil[i].startln, - stencil_data=stencil, - ) + if self.interface.StartDelete != UnusedDirective: + logger.info("Generating DELETE statement.") + for start, end in zip(self.interface.StartDelete, self.interface.EndDelete, strict=True): + self._generate( + StartDeleteStatement, + StartDeleteStatementGenerator, + start.startln, + ) + self._generate( + EndDeleteStatement, + EndDeleteStatementGenerator, + end.startln, + ) def _generate_imports(self) -> None: """Generate f90 code for import statements.""" diff --git a/tools/src/icon4pytools/liskov/codegen/integration/template.py b/tools/src/icon4pytools/liskov/codegen/integration/template.py index f733dfa777..3c0717e448 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/template.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/template.py @@ -512,6 +512,19 @@ class EndCreateStatementGenerator(TemplatedGenerator): EndCreateStatement = as_jinja("!$ACC END DATA") +class StartDeleteStatement(eve.Node): + ... + +class StartDeleteStatementGenerator(TemplatedGenerator): + StartDeleteStatement = as_jinja("#ifdef __DSL_VERIFY") + + +class EndDeleteStatement(eve.Node): + ... + +class EndDeleteStatementGenerator(TemplatedGenerator): + EndDeleteStatement = as_jinja("#endif") + class EndIfStatement(eve.Node): ... From 37f4d6936179002c6b98286bd0bb60bf77cf9767 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 19 Jul 2023 10:01:41 +0200 Subject: [PATCH 31/49] cleanup style --- .../liskov/codegen/integration/generate.py | 76 +++++++++---------- .../liskov/codegen/integration/template.py | 3 + 2 files changed, 40 insertions(+), 39 deletions(-) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/generate.py b/tools/src/icon4pytools/liskov/codegen/integration/generate.py index 60b9bd0153..ddb5270e9b 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/generate.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/generate.py @@ -24,6 +24,8 @@ DeclareStatementGenerator, EndCreateStatement, EndCreateStatementGenerator, + EndDeleteStatement, + EndDeleteStatementGenerator, EndFusedStencilStatement, EndFusedStencilStatementGenerator, EndIfStatement, @@ -40,12 +42,10 @@ MetadataStatementGenerator, StartCreateStatement, StartCreateStatementGenerator, - StartFusedStencilStatement, - StartFusedStencilStatementGenerator, StartDeleteStatement, StartDeleteStatementGenerator, - EndDeleteStatement, - EndDeleteStatementGenerator, + StartFusedStencilStatement, + StartFusedStencilStatementGenerator, StartProfileStatement, StartProfileStatementGenerator, StartStencilStatement, @@ -159,7 +159,6 @@ def _generate_start_stencil(self) -> None: ) i += 1 - def _generate_end_stencil(self) -> None: """Generate f90 integration code surrounding a stencil. @@ -180,47 +179,46 @@ def _generate_end_stencil(self) -> None: ) def _generate_start_fused_stencil(self) -> None: - """Generate f90 integration code surrounding a fused stencil. - """ + """Generate f90 integration code surrounding a fused stencil.""" if self.interface.StartFusedStencil != UnusedDirective: - for stencil in self.interface.StartFusedStencil: - logger.info(f"Generating START FUSED statement for {stencil.name}") - self._generate( - StartFusedStencilStatement, - StartFusedStencilStatementGenerator, - stencil.startln, - stencil_data=stencil, - ) + for stencil in self.interface.StartFusedStencil: + logger.info(f"Generating START FUSED statement for {stencil.name}") + self._generate( + StartFusedStencilStatement, + StartFusedStencilStatementGenerator, + stencil.startln, + stencil_data=stencil, + ) def _generate_end_fused_stencil(self) -> None: - """Generate f90 integration code surrounding a fused stencil. - """ + """Generate f90 integration code surrounding a fused stencil.""" if self.interface.EndFusedStencil != UnusedDirective: - for i, stencil in enumerate(self.interface.StartFusedStencil): - logger.info(f"Generating END Fused statement for {stencil.name}") - self._generate( - EndFusedStencilStatement, - EndFusedStencilStatementGenerator, - self.interface.EndFusedStencil[i].startln, - stencil_data=stencil, - ) + for i, stencil in enumerate(self.interface.StartFusedStencil): + logger.info(f"Generating END Fused statement for {stencil.name}") + self._generate( + EndFusedStencilStatement, + EndFusedStencilStatementGenerator, + self.interface.EndFusedStencil[i].startln, + stencil_data=stencil, + ) def _generate_delete(self) -> None: - """Generate f90 integration code for delete section. - """ + """Generate f90 integration code for delete section.""" if self.interface.StartDelete != UnusedDirective: - logger.info("Generating DELETE statement.") - for start, end in zip(self.interface.StartDelete, self.interface.EndDelete, strict=True): - self._generate( - StartDeleteStatement, - StartDeleteStatementGenerator, - start.startln, - ) - self._generate( - EndDeleteStatement, - EndDeleteStatementGenerator, - end.startln, - ) + logger.info("Generating DELETE statement.") + for start, end in zip( + self.interface.StartDelete, self.interface.EndDelete, strict=True + ): + self._generate( + StartDeleteStatement, + StartDeleteStatementGenerator, + start.startln, + ) + self._generate( + EndDeleteStatement, + EndDeleteStatementGenerator, + end.startln, + ) def _generate_imports(self) -> None: """Generate f90 code for import statements.""" diff --git a/tools/src/icon4pytools/liskov/codegen/integration/template.py b/tools/src/icon4pytools/liskov/codegen/integration/template.py index 3c0717e448..dfe602e882 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/template.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/template.py @@ -515,6 +515,7 @@ class EndCreateStatementGenerator(TemplatedGenerator): class StartDeleteStatement(eve.Node): ... + class StartDeleteStatementGenerator(TemplatedGenerator): StartDeleteStatement = as_jinja("#ifdef __DSL_VERIFY") @@ -522,9 +523,11 @@ class StartDeleteStatementGenerator(TemplatedGenerator): class EndDeleteStatement(eve.Node): ... + class EndDeleteStatementGenerator(TemplatedGenerator): EndDeleteStatement = as_jinja("#endif") + class EndIfStatement(eve.Node): ... From 1bfabd5f9da8a7932803e9a8ac7b52deb916dda3 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 19 Jul 2023 10:41:04 +0200 Subject: [PATCH 32/49] fix import for fused --- tools/src/icon4pytools/liskov/codegen/integration/generate.py | 2 +- tools/src/icon4pytools/liskov/codegen/integration/template.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/generate.py b/tools/src/icon4pytools/liskov/codegen/integration/generate.py index ddb5270e9b..6c23289bc6 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/generate.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/generate.py @@ -227,7 +227,7 @@ def _generate_imports(self) -> None: ImportsStatement, ImportsStatementGenerator, self.interface.Imports.startln, - stencils=self.interface.StartStencil, + stencils=self.interface.StartStencil+self.interface.StartFusedStencil, ) def _generate_create(self) -> None: diff --git a/tools/src/icon4pytools/liskov/codegen/integration/template.py b/tools/src/icon4pytools/liskov/codegen/integration/template.py index dfe602e882..519ac7c941 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/template.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/template.py @@ -24,6 +24,7 @@ DeclareData, StartFusedStencilData, StartStencilData, + StartBasicStencilData, ) from icon4pytools.liskov.external.metadata import CodeMetadata @@ -473,7 +474,7 @@ class StartFusedStencilStatementGenerator(TemplatedGenerator): class ImportsStatement(eve.Node): - stencils: list[StartStencilData] + stencils: list[StartBasicStencilData] stencil_names: list[str] = eve.datamodels.field(init=False) def __post_init__(self) -> None: # type: ignore From eec5f6d1ac673b1b50f790eda1e9db5a121c267f Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 19 Jul 2023 17:07:18 +0200 Subject: [PATCH 33/49] workaround for data regions --- .../liskov/codegen/integration/template.py | 91 +++++++++++++------ 1 file changed, 63 insertions(+), 28 deletions(-) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/template.py b/tools/src/icon4pytools/liskov/codegen/integration/template.py index 519ac7c941..598a3814c0 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/template.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/template.py @@ -33,6 +33,8 @@ def enclose_in_parentheses(string: str) -> str: return f"({string})" + + class BoundsFields(eve.Node): vlower: str vupper: str @@ -126,26 +128,6 @@ def __post_init__(self) -> None: # type: ignore ) -class EndFusedStencilStatement(eve.Node): - stencil_data: StartFusedStencilData - - name: str = eve.datamodels.field(init=False) - input_fields: InputFields = eve.datamodels.field(init=False) - output_fields: OutputFields = eve.datamodels.field(init=False) - tolerance_fields: ToleranceFields = eve.datamodels.field(init=False) - bounds_fields: BoundsFields = eve.datamodels.field(init=False) - - def __post_init__(self) -> None: # type: ignore - all_fields = [Field(**asdict(f)) for f in self.stencil_data.fields] - self.bounds_fields = BoundsFields(**asdict(self.stencil_data.bounds)) - self.name = self.stencil_data.name - self.input_fields = InputFields(fields=[f for f in all_fields if f.inp]) - self.output_fields = OutputFields(fields=[f for f in all_fields if f.out]) - self.tolerance_fields = ToleranceFields( - fields=[f for f in all_fields if f.rel_tol or f.abs_tol] - ) - - class EndFusedStencilStatementGenerator(TemplatedGenerator): EndFusedStencilStatement = as_jinja( """ @@ -155,7 +137,11 @@ class EndFusedStencilStatementGenerator(TemplatedGenerator): {{ tolerance_fields }} {{ bounds_fields }} - !$ACC END DATA + !$ACC EXIT DATA DELETE( & + {%- for d in _this_node.copy_declarations %} + !$ACC {{ d.variable }}_before {%- if not loop.last -%}, & {% else %} ) & {%- endif -%} + {%- endfor %} + !$ACC IF ( i_am_accel_node ) """ ) @@ -301,6 +287,9 @@ def visit_OutputFields(self, out: OutputFields) -> OutputFields: # type: ignore class Declaration(Assign): ... +class CopyDeclaration(Declaration): + lh_index: str + rh_index: str class DeclareStatement(eve.Node): declare_data: DeclareData @@ -324,11 +313,6 @@ class DeclareStatementGenerator(TemplatedGenerator): ) -class CopyDeclaration(Declaration): - lh_index: str - rh_index: str - - class StartStencilStatement(eve.Node): stencil_data: StartStencilData profile: bool @@ -407,6 +391,57 @@ def make_copy_declaration(f: Field) -> CopyDeclaration: rh_index=rh_idx, ) +class EndFusedStencilStatement(eve.Node): + stencil_data: StartFusedStencilData + + name: str = eve.datamodels.field(init=False) + input_fields: InputFields = eve.datamodels.field(init=False) + output_fields: OutputFields = eve.datamodels.field(init=False) + tolerance_fields: ToleranceFields = eve.datamodels.field(init=False) + bounds_fields: BoundsFields = eve.datamodels.field(init=False) + copy_declarations: list[CopyDeclaration] = eve.datamodels.field(init=False) + + def __post_init__(self) -> None: # type: ignore + all_fields = [Field(**asdict(f)) for f in self.stencil_data.fields] + self.copy_declarations = [self.make_copy_declaration(f) for f in all_fields if f.out] + self.bounds_fields = BoundsFields(**asdict(self.stencil_data.bounds)) + self.name = self.stencil_data.name + self.input_fields = InputFields(fields=[f for f in all_fields if f.inp]) + self.output_fields = OutputFields(fields=[f for f in all_fields if f.out]) + self.tolerance_fields = ToleranceFields( + fields=[f for f in all_fields if f.rel_tol or f.abs_tol] + ) + + @staticmethod + def make_copy_declaration(f: Field) -> CopyDeclaration: + if f.dims is None: + raise UndeclaredFieldError(f"{f.variable} was not declared!") + + lh_idx = render_index(f.dims) + + # get length of association index + association_dims = get_array_dims(f.association).split(",") + n_association_dims = len(association_dims) + + offset = len(",".join(association_dims)) + 2 + truncated_association = f.association[:-offset] + + if n_association_dims > f.dims: + rh_idx = f"{lh_idx},{association_dims[-1]}" + else: + rh_idx = f"{lh_idx}" + + lh_idx = enclose_in_parentheses(lh_idx) + rh_idx = enclose_in_parentheses(rh_idx) + + return CopyDeclaration( + variable=f.variable, + association=truncated_association, + lh_index=lh_idx, + rh_index=rh_idx, + ) + + def render_index(n: int) -> str: """ @@ -455,14 +490,14 @@ class StartFusedStencilStatementGenerator(TemplatedGenerator): StartFusedStencilStatement = as_jinja( """ - !$ACC DATA CREATE( & + !$ACC ENTER DATA CREATE( & {%- for d in _this_node.copy_declarations %} !$ACC {{ d.variable }}_before {%- if not loop.last -%}, & {% else %} ) & {%- endif -%} {%- endfor %} !$ACC IF ( i_am_accel_node ) #ifdef __DSL_VERIFY - !$ACC KERNELS IF( i_am_accel_node ) DEFAULT({{ _this_node.acc_present }}) ASYNC(1) + !$ACC KERNELS IF( i_am_accel_node ) DEFAULT(PRESENT) ASYNC(1) {%- for d in _this_node.copy_declarations %} {{ d.variable }}_before{{ d.lh_index }} = {{ d.association }}{{ d.rh_index }} {%- endfor %} From c2685a87faef99a16fe712db798d567fb5f58602 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 19 Jul 2023 17:08:19 +0200 Subject: [PATCH 34/49] fix style --- .../icon4pytools/liskov/codegen/integration/generate.py | 2 +- .../icon4pytools/liskov/codegen/integration/template.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/generate.py b/tools/src/icon4pytools/liskov/codegen/integration/generate.py index 6c23289bc6..40324d138f 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/generate.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/generate.py @@ -227,7 +227,7 @@ def _generate_imports(self) -> None: ImportsStatement, ImportsStatementGenerator, self.interface.Imports.startln, - stencils=self.interface.StartStencil+self.interface.StartFusedStencil, + stencils=self.interface.StartStencil + self.interface.StartFusedStencil, ) def _generate_create(self) -> None: diff --git a/tools/src/icon4pytools/liskov/codegen/integration/template.py b/tools/src/icon4pytools/liskov/codegen/integration/template.py index 598a3814c0..3706a5956b 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/template.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/template.py @@ -22,9 +22,9 @@ from icon4pytools.liskov.codegen.integration.exceptions import UndeclaredFieldError from icon4pytools.liskov.codegen.integration.interface import ( DeclareData, + StartBasicStencilData, StartFusedStencilData, StartStencilData, - StartBasicStencilData, ) from icon4pytools.liskov.external.metadata import CodeMetadata @@ -33,8 +33,6 @@ def enclose_in_parentheses(string: str) -> str: return f"({string})" - - class BoundsFields(eve.Node): vlower: str vupper: str @@ -287,10 +285,12 @@ def visit_OutputFields(self, out: OutputFields) -> OutputFields: # type: ignore class Declaration(Assign): ... + class CopyDeclaration(Declaration): lh_index: str rh_index: str + class DeclareStatement(eve.Node): declare_data: DeclareData declarations: list[Declaration] = eve.datamodels.field(init=False) @@ -391,6 +391,7 @@ def make_copy_declaration(f: Field) -> CopyDeclaration: rh_index=rh_idx, ) + class EndFusedStencilStatement(eve.Node): stencil_data: StartFusedStencilData @@ -442,7 +443,6 @@ def make_copy_declaration(f: Field) -> CopyDeclaration: ) - def render_index(n: int) -> str: """ Render a string of comma-separated colon characters, used to define the shape of an array in Fortran. From e1139a9425a2d8b172e963ca77e8e8299f379305 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 25 Jul 2023 12:04:02 +0200 Subject: [PATCH 35/49] refactor 1 --- .../liskov/codegen/integration/template.py | 51 ++++++++++++++----- 1 file changed, 37 insertions(+), 14 deletions(-) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/template.py b/tools/src/icon4pytools/liskov/codegen/integration/template.py index 3706a5956b..20ff37dd2e 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/template.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/template.py @@ -83,6 +83,7 @@ def get_array_dims(association: str) -> str: return "".join(list(dims)) + class MetadataStatement(eve.Node): metadata: CodeMetadata @@ -101,20 +102,21 @@ class MetadataStatementGenerator(TemplatedGenerator): """ ) +class EndBasicStencilStatement(eve.Node): + name: str = eve.datamodels.field(init=False) + input_fields: InputFields = eve.datamodels.field(init=False) + output_fields: OutputFields = eve.datamodels.field(init=False) + tolerance_fields: ToleranceFields = eve.datamodels.field(init=False) + bounds_fields: BoundsFields = eve.datamodels.field(init=False) + -class EndStencilStatement(eve.Node): +class EndStencilStatement(EndBasicStencilStatement): stencil_data: StartStencilData profile: bool noendif: Optional[bool] noprofile: Optional[bool] noaccenddata: Optional[bool] - name: str = eve.datamodels.field(init=False) - input_fields: InputFields = eve.datamodels.field(init=False) - output_fields: OutputFields = eve.datamodels.field(init=False) - tolerance_fields: ToleranceFields = eve.datamodels.field(init=False) - bounds_fields: BoundsFields = eve.datamodels.field(init=False) - def __post_init__(self) -> None: # type: ignore all_fields = [Field(**asdict(f)) for f in self.stencil_data.fields] self.bounds_fields = BoundsFields(**asdict(self.stencil_data.bounds)) @@ -290,6 +292,33 @@ class CopyDeclaration(Declaration): lh_index: str rh_index: str +def _make_copy_declaration(f: Field) -> CopyDeclaration: + if f.dims is None: + raise UndeclaredFieldError(f"{f.variable} was not declared!") + + lh_idx = render_index(f.dims) + + # get length of association index + association_dims = get_array_dims(f.association).split(",") + n_association_dims = len(association_dims) + + offset = len(",".join(association_dims)) + 2 + truncated_association = f.association[:-offset] + + if n_association_dims > f.dims: + rh_idx = f"{lh_idx},{association_dims[-1]}" + else: + rh_idx = f"{lh_idx}" + + lh_idx = enclose_in_parentheses(lh_idx) + rh_idx = enclose_in_parentheses(rh_idx) + + return CopyDeclaration( + variable=f.variable, + association=truncated_association, + lh_index=lh_idx, + rh_index=rh_idx, + ) class DeclareStatement(eve.Node): declare_data: DeclareData @@ -392,14 +421,8 @@ def make_copy_declaration(f: Field) -> CopyDeclaration: ) -class EndFusedStencilStatement(eve.Node): +class EndFusedStencilStatement(EndBasicStencilStatement): stencil_data: StartFusedStencilData - - name: str = eve.datamodels.field(init=False) - input_fields: InputFields = eve.datamodels.field(init=False) - output_fields: OutputFields = eve.datamodels.field(init=False) - tolerance_fields: ToleranceFields = eve.datamodels.field(init=False) - bounds_fields: BoundsFields = eve.datamodels.field(init=False) copy_declarations: list[CopyDeclaration] = eve.datamodels.field(init=False) def __post_init__(self) -> None: # type: ignore From 2f4864e3b99c8332a81d5ba1e68297f100924fc2 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 25 Jul 2023 15:43:10 +0200 Subject: [PATCH 36/49] more refactor --- .../liskov/codegen/integration/template.py | 94 +------------------ 1 file changed, 3 insertions(+), 91 deletions(-) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/template.py b/tools/src/icon4pytools/liskov/codegen/integration/template.py index 20ff37dd2e..3fce6ff263 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/template.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/template.py @@ -349,77 +349,18 @@ class StartStencilStatement(eve.Node): def __post_init__(self) -> None: # type: ignore all_fields = [Field(**asdict(f)) for f in self.stencil_data.fields] - self.copy_declarations = [self.make_copy_declaration(f) for f in all_fields if f.out] + self.copy_declarations = [_make_copy_declaration(f) for f in all_fields if f.out] self.acc_present = "PRESENT" if self.stencil_data.acc_present else "NONE" - @staticmethod - def make_copy_declaration(f: Field) -> CopyDeclaration: - if f.dims is None: - raise UndeclaredFieldError(f"{f.variable} was not declared!") - - lh_idx = render_index(f.dims) - - # get length of association index - association_dims = get_array_dims(f.association).split(",") - n_association_dims = len(association_dims) - - offset = len(",".join(association_dims)) + 2 - truncated_association = f.association[:-offset] - - if n_association_dims > f.dims: - rh_idx = f"{lh_idx},{association_dims[-1]}" - else: - rh_idx = f"{lh_idx}" - - lh_idx = enclose_in_parentheses(lh_idx) - rh_idx = enclose_in_parentheses(rh_idx) - - return CopyDeclaration( - variable=f.variable, - association=truncated_association, - lh_index=lh_idx, - rh_index=rh_idx, - ) - - class StartFusedStencilStatement(eve.Node): stencil_data: StartFusedStencilData copy_declarations: list[CopyDeclaration] = eve.datamodels.field(init=False) def __post_init__(self) -> None: # type: ignore all_fields = [Field(**asdict(f)) for f in self.stencil_data.fields] - self.copy_declarations = [self.make_copy_declaration(f) for f in all_fields if f.out] + self.copy_declarations = [_make_copy_declaration(f) for f in all_fields if f.out] self.acc_present = "PRESENT" if self.stencil_data.acc_present else "NONE" - @staticmethod - def make_copy_declaration(f: Field) -> CopyDeclaration: - if f.dims is None: - raise UndeclaredFieldError(f"{f.variable} was not declared!") - - lh_idx = render_index(f.dims) - - # get length of association index - association_dims = get_array_dims(f.association).split(",") - n_association_dims = len(association_dims) - - offset = len(",".join(association_dims)) + 2 - truncated_association = f.association[:-offset] - - if n_association_dims > f.dims: - rh_idx = f"{lh_idx},{association_dims[-1]}" - else: - rh_idx = f"{lh_idx}" - - lh_idx = enclose_in_parentheses(lh_idx) - rh_idx = enclose_in_parentheses(rh_idx) - - return CopyDeclaration( - variable=f.variable, - association=truncated_association, - lh_index=lh_idx, - rh_index=rh_idx, - ) - class EndFusedStencilStatement(EndBasicStencilStatement): stencil_data: StartFusedStencilData @@ -427,7 +368,7 @@ class EndFusedStencilStatement(EndBasicStencilStatement): def __post_init__(self) -> None: # type: ignore all_fields = [Field(**asdict(f)) for f in self.stencil_data.fields] - self.copy_declarations = [self.make_copy_declaration(f) for f in all_fields if f.out] + self.copy_declarations = [_make_copy_declaration(f) for f in all_fields if f.out] self.bounds_fields = BoundsFields(**asdict(self.stencil_data.bounds)) self.name = self.stencil_data.name self.input_fields = InputFields(fields=[f for f in all_fields if f.inp]) @@ -436,35 +377,6 @@ def __post_init__(self) -> None: # type: ignore fields=[f for f in all_fields if f.rel_tol or f.abs_tol] ) - @staticmethod - def make_copy_declaration(f: Field) -> CopyDeclaration: - if f.dims is None: - raise UndeclaredFieldError(f"{f.variable} was not declared!") - - lh_idx = render_index(f.dims) - - # get length of association index - association_dims = get_array_dims(f.association).split(",") - n_association_dims = len(association_dims) - - offset = len(",".join(association_dims)) + 2 - truncated_association = f.association[:-offset] - - if n_association_dims > f.dims: - rh_idx = f"{lh_idx},{association_dims[-1]}" - else: - rh_idx = f"{lh_idx}" - - lh_idx = enclose_in_parentheses(lh_idx) - rh_idx = enclose_in_parentheses(rh_idx) - - return CopyDeclaration( - variable=f.variable, - association=truncated_association, - lh_index=lh_idx, - rh_index=rh_idx, - ) - def render_index(n: int) -> str: """ From 12b6248da3e7b7f68e17bc133eccc158de1f8ead Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Tue, 25 Jul 2023 15:47:07 +0200 Subject: [PATCH 37/49] improve help test --- tools/src/icon4pytools/liskov/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/src/icon4pytools/liskov/cli.py b/tools/src/icon4pytools/liskov/cli.py index 6d80439381..e6685c3ad5 100644 --- a/tools/src/icon4pytools/liskov/cli.py +++ b/tools/src/icon4pytools/liskov/cli.py @@ -54,7 +54,7 @@ def main(ctx): "--fused/--unfused", "-f/-u", default=True, - help="Add metadata header with information about program.", + help="Adds fused or unfused stencils.", ) @click.argument( "input_path", From 3d37a5ad3ec792c19e32043a3bf3f3936ddd8e84 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 26 Jul 2023 08:55:32 +0200 Subject: [PATCH 38/49] test comenting out DSL --- tools/tests/liskov/fortran_samples.py | 49 +++++++++++++++++++++++++++ tools/tests/liskov/test_cli.py | 2 ++ tools/tests/liskov/test_external.py | 4 +++ tools/tests/liskov/test_generation.py | 4 +++ tools/tests/liskov/test_parser.py | 5 +-- tools/tests/liskov/test_validation.py | 8 ++--- 6 files changed, 66 insertions(+), 6 deletions(-) diff --git a/tools/tests/liskov/fortran_samples.py b/tools/tests/liskov/fortran_samples.py index 858b3a2089..785ab2ab59 100644 --- a/tools/tests/liskov/fortran_samples.py +++ b/tools/tests/liskov/fortran_samples.py @@ -69,6 +69,55 @@ !$DSL END CREATE() """ +SINGLE_STENCIL_WITH_COMMENTS = """\ + ! Use !$DSL statements, they are great. They can be easily commented out by: + + !!$DSL IMPORTS() + + ! $DSL START CREATE() + + !$DSL IMPORTS() + + !$DSL START CREATE() + + !$DSL DECLARE(vn=nproma,p_patch%nlev,p_patch%nblks_e; suffix=dsl) + + !$DSL DECLARE(vn=nproma,p_patch%nlev,p_patch%nblks_e; a=nproma,p_patch%nlev,p_patch%nblks_e; & + !$DSL b=nproma,p_patch%nlev,p_patch%nblks_e; type=REAL(vp)) + + !$DSL START STENCIL(name=apply_nabla2_to_vn_in_lateral_boundary; & + !$DSL z_nabla2_e=z_nabla2_e(:,:,1); area_edge=p_patch%edges%area_edge(:,1); & + !$DSL fac_bdydiff_v=fac_bdydiff_v; vn=p_nh_prog%vn(:,:,1); & + !$DSL vertical_lower=1; vertical_upper=nlev; & + !$DSL horizontal_lower=i_startidx; horizontal_upper=i_endidx; & + !$DSL accpresent=True) + !$OMP DO PRIVATE(je,jk,jb,i_startidx,i_endidx) ICON_OMP_DEFAULT_SCHEDULE + DO jb = i_startblk,i_endblk + + CALL get_indices_e(p_patch, jb, i_startblk, i_endblk, & + i_startidx, i_endidx, start_bdydiff_e, grf_bdywidth_e) + + !$ACC PARALLEL IF( i_am_accel_node .AND. acc_on ) DEFAULT(NONE) ASYNC(1) + vn_before(:,:,:) = p_nh_prog%vn(:,:,:) + !$ACC END PARALLEL + + !$ACC PARALLEL LOOP DEFAULT(NONE) GANG VECTOR COLLAPSE(2) ASYNC(1) IF( i_am_accel_node .AND. acc_on ) + DO jk = 1, nlev + !DIR$ IVDEP + DO je = i_startidx, i_endidx + p_nh_prog%vn(je,jk,jb) = & + p_nh_prog%vn(je,jk,jb) + & + z_nabla2_e(je,jk,jb) * & + p_patch%edges%area_edge(je,jb)*fac_bdydiff_v + ENDDO + ENDDO + !$DSL START PROFILE(name=apply_nabla2_to_vn_in_lateral_boundary) + !$ACC END PARALLEL LOOP + !$DSL END PROFILE() + !$DSL END STENCIL(name=apply_nabla2_to_vn_in_lateral_boundary; noprofile=True) + !$DSL END CREATE() + """ + MULTIPLE_STENCILS = """\ !$DSL IMPORTS() diff --git a/tools/tests/liskov/test_cli.py b/tools/tests/liskov/test_cli.py index edd178f62e..028170e023 100644 --- a/tools/tests/liskov/test_cli.py +++ b/tools/tests/liskov/test_cli.py @@ -24,6 +24,7 @@ MULTIPLE_STENCILS, NO_DIRECTIVES_STENCIL, SINGLE_STENCIL, + SINGLE_STENCIL_WITH_COMMENTS, ) @@ -37,6 +38,7 @@ def outfile(tmp_path): files = [ ("NO_DIRECTIVES", NO_DIRECTIVES_STENCIL), ("SINGLE", SINGLE_STENCIL), + ("COMMENTS", SINGLE_STENCIL_WITH_COMMENTS), ("CONSECUTIVE", CONSECUTIVE_STENCIL), ("FREE_FORM", FREE_FORM_STENCIL), ("MULTIPLE", MULTIPLE_STENCILS), diff --git a/tools/tests/liskov/test_external.py b/tools/tests/liskov/test_external.py index 1e0de1b5ca..297bd6ab48 100644 --- a/tools/tests/liskov/test_external.py +++ b/tools/tests/liskov/test_external.py @@ -82,6 +82,10 @@ def test_stencil_collector_invalid_member(): Imports=None, Declare=None, EndStencil=None, + StartFusedStencil=None, + EndFusedStencil=None, + StartDelete=None, + EndDelete=None, StartCreate=None, EndCreate=None, EndIf=None, diff --git a/tools/tests/liskov/test_generation.py b/tools/tests/liskov/test_generation.py index dc16b09c6a..bfe9397b34 100644 --- a/tools/tests/liskov/test_generation.py +++ b/tools/tests/liskov/test_generation.py @@ -93,6 +93,10 @@ def integration_code_interface(): return IntegrationCodeInterface( StartStencil=[start_stencil_data], EndStencil=[end_stencil_data], + StartFusedStencil=[], + EndFusedStencil=[], + StartDelete=[], + EndDelete=[], Declare=[declare_data], Imports=imports_data, StartCreate=[start_create_data], diff --git a/tools/tests/liskov/test_parser.py b/tools/tests/liskov/test_parser.py index 5bb2518e13..c9373c36ec 100644 --- a/tools/tests/liskov/test_parser.py +++ b/tools/tests/liskov/test_parser.py @@ -22,7 +22,7 @@ from icon4pytools.liskov.parsing.parse import DirectivesParser from .conftest import insert_new_lines, scan_for_directives -from .fortran_samples import MULTIPLE_STENCILS, NO_DIRECTIVES_STENCIL, SINGLE_STENCIL +from .fortran_samples import MULTIPLE_STENCILS, NO_DIRECTIVES_STENCIL, SINGLE_STENCIL, SINGLE_STENCIL_WITH_COMMENTS def test_parse_no_input(): @@ -76,7 +76,7 @@ def test_parse_single_directive(directive, string, startln, endln, expected_cont @mark.parametrize( "stencil, num_directives, num_content", - [(SINGLE_STENCIL, 9, 8), (MULTIPLE_STENCILS, 11, 7)], + [(SINGLE_STENCIL, 9, 8), (SINGLE_STENCIL_WITH_COMMENTS, 9, 8), (MULTIPLE_STENCILS, 11, 7)], ) def test_file_parsing(make_f90_tmpfile, stencil, num_directives, num_content): fpath = make_f90_tmpfile(content=stencil) @@ -108,6 +108,7 @@ def test_directive_parser_no_directives_found(make_f90_tmpfile): "stencil, directive", [ (SINGLE_STENCIL, "!$DSL FOO()"), + (SINGLE_STENCIL_WITH_COMMENTS, "!$DSL FOO()"), (MULTIPLE_STENCILS, "!$DSL BAR()"), ], ) diff --git a/tools/tests/liskov/test_validation.py b/tools/tests/liskov/test_validation.py index d6fbc4f927..0e5c8eb531 100644 --- a/tools/tests/liskov/test_validation.py +++ b/tools/tests/liskov/test_validation.py @@ -24,7 +24,7 @@ from icon4pytools.liskov.parsing.validation import DirectiveSyntaxValidator from .conftest import insert_new_lines, scan_for_directives -from .fortran_samples import MULTIPLE_STENCILS, SINGLE_STENCIL +from .fortran_samples import MULTIPLE_STENCILS, SINGLE_STENCIL, SINGLE_STENCIL_WITH_COMMENTS @mark.parametrize( @@ -73,7 +73,7 @@ def test_directive_syntax_validator(directive): ], ) def test_directive_semantics_validation_repeated_directives(make_f90_tmpfile, directive): - fpath = make_f90_tmpfile(content=SINGLE_STENCIL) + fpath = make_f90_tmpfile(content=SINGLE_STENCIL_WITH_COMMENTS) opath = fpath.with_suffix(".gen") insert_new_lines(fpath, [directive]) directives = scan_for_directives(fpath) @@ -93,7 +93,7 @@ def test_directive_semantics_validation_repeated_directives(make_f90_tmpfile, di ], ) def test_directive_semantics_validation_repeated_stencil(make_f90_tmpfile, directive): - fpath = make_f90_tmpfile(content=SINGLE_STENCIL) + fpath = make_f90_tmpfile(content=SINGLE_STENCIL_WITH_COMMENTS) opath = fpath.with_suffix(".gen") insert_new_lines(fpath, [directive]) directives = scan_for_directives(fpath) @@ -109,7 +109,7 @@ def test_directive_semantics_validation_repeated_stencil(make_f90_tmpfile, direc ], ) def test_directive_semantics_validation_required_directives(make_f90_tmpfile, directive): - new = SINGLE_STENCIL.replace(directive, "") + new = SINGLE_STENCIL_WITH_COMMENTS.replace(directive, "") fpath = make_f90_tmpfile(content=new) opath = fpath.with_suffix(".gen") directives = scan_for_directives(fpath) From bbf9505ecb43cb6297b753f1e904daa00f6ebc69 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 26 Jul 2023 10:36:12 +0200 Subject: [PATCH 39/49] more tests --- tools/tests/liskov/fortran_samples.py | 48 +++++++++++++++++++++++++++ tools/tests/liskov/test_cli.py | 4 ++- 2 files changed, 51 insertions(+), 1 deletion(-) diff --git a/tools/tests/liskov/fortran_samples.py b/tools/tests/liskov/fortran_samples.py index 785ab2ab59..696768d830 100644 --- a/tools/tests/liskov/fortran_samples.py +++ b/tools/tests/liskov/fortran_samples.py @@ -261,6 +261,54 @@ """ +FUSED_STENCIL = """\ + !$DSL IMPORTS() + + !$DSL INSERT(INTEGER :: start_interior_idx_c, end_interior_idx_c, start_nudging_idx_c, end_halo_1_idx_c) + + !$DSL DECLARE(kh_smag_e=nproma,p_patch%nlev,p_patch%nblks_e; & + !$DSL kh_smag_ec=nproma,p_patch%nlev,p_patch%nblks_e; & + !$DSL z_nabla2_e=nproma,p_patch%nlev,p_patch%nblks_e; & + !$DSL kh_c=nproma,p_patch%nlev; & + !$DSL div=nproma,p_patch%nlev; & + !$DSL div_ic=nproma,p_patch%nlev,p_patch%nblks_c; & + !$DSL hdef_ic=nproma,p_patch%nlev,p_patch%nblks_c; & + !$DSL z_nabla4_e2=nproma,p_patch%nlev; & + !$DSL vn=nproma,p_patch%nlev,p_patch%nblks_e; & + !$DSL z_nabla2_c=nproma,p_patch%nlev,p_patch%nblks_e; & + !$DSL dwdx=nproma,p_patch%nlev,p_patch%nblks_c; & + !$DSL dwdy=nproma,p_patch%nlev,p_patch%nblks_c; & + !$DSL w=nproma,p_patch%nlev,p_patch%nblks_c; & + !$DSL enh_diffu_3d=nproma,p_patch%nlev,p_patch%nblks_c; & + !$DSL z_temp=nproma,p_patch%nlev,p_patch%nblks_c; & + !$DSL theta_v=nproma,p_patch%nlev,p_patch%nblks_c; & + !$DSL exner=nproma,p_patch%nlev,p_patch%nblks_c) + + !$DSL START FUSED STENCIL(name=calculate_diagnostic_quantities_for_turbulence; & + !$DSL kh_smag_ec=kh_smag_ec(:,:,1); vn=p_nh_prog%vn(:,:,1); e_bln_c_s=p_int%e_bln_c_s(:,:,1); & + !$DSL geofac_div=p_int%geofac_div(:,:,1); diff_multfac_smag=diff_multfac_smag(:); & + !$DSL wgtfac_c=p_nh_metrics%wgtfac_c(:,:,1); div_ic=p_nh_diag%div_ic(:,:,1); & + !$DSL hdef_ic=p_nh_diag%hdef_ic(:,:,1); & + !$DSL div_ic_abs_tol=1e-18_wp; vertical_lower=2; & + !$DSL vertical_upper=nlev; horizontal_lower=i_startidx; horizontal_upper=i_endidx) + + !$DSL START STENCIL(name=temporary_fields_for_turbulence_diagnostics; kh_smag_ec=kh_smag_ec(:,:,1); vn=p_nh_prog%vn(:,:,1); & + !$DSL e_bln_c_s=p_int%e_bln_c_s(:,:,1); geofac_div=p_int%geofac_div(:,:,1); & + !$DSL diff_multfac_smag=diff_multfac_smag(:); kh_c=kh_c(:,:); div=div(:,:); & + !$DSL vertical_lower=1; vertical_upper=nlev; horizontal_lower=i_startidx; & + !$DSL horizontal_upper=i_endidx) + + !$DSL END STENCIL(name=temporary_fields_for_turbulence_diagnostics) + + !$DSL START STENCIL(name=calculate_diagnostics_for_turbulence; div=div; kh_c=kh_c; wgtfac_c=p_nh_metrics%wgtfac_c(:,:,1); & + !$DSL div_ic=p_nh_diag%div_ic(:,:,1); hdef_ic=p_nh_diag%hdef_ic(:,:,1); div_ic_abs_tol=1e-18_wp; & + !$DSL vertical_lower=2; vertical_upper=nlev; horizontal_lower=i_startidx; horizontal_upper=i_endidx) + + !$DSL END STENCIL(name=calculate_diagnostics_for_turbulence) + + !$DSL END FUSED STENCIL(name=calculate_diagnostic_quantities_for_turbulence) + """ + FREE_FORM_STENCIL = """\ !$DSL IMPORTS() diff --git a/tools/tests/liskov/test_cli.py b/tools/tests/liskov/test_cli.py index 028170e023..87c591b3f5 100644 --- a/tools/tests/liskov/test_cli.py +++ b/tools/tests/liskov/test_cli.py @@ -25,6 +25,7 @@ NO_DIRECTIVES_STENCIL, SINGLE_STENCIL, SINGLE_STENCIL_WITH_COMMENTS, + FUSED_STENCIL, ) @@ -42,9 +43,10 @@ def outfile(tmp_path): ("CONSECUTIVE", CONSECUTIVE_STENCIL), ("FREE_FORM", FREE_FORM_STENCIL), ("MULTIPLE", MULTIPLE_STENCILS), + ("FUSED", FUSED_STENCIL), ] -flags = {"serialise": ["--multinode"], "integrate": ["-p", "-m"]} +flags = {"serialise": ["--multinode"], "integrate": ["-p", "-m", "-f", "-u"]} for file_name, file_content in files: for cmd in flags.keys(): From 49045462380e5e8e6aafdc68b11a567713a828c8 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 26 Jul 2023 10:45:29 +0200 Subject: [PATCH 40/49] more tests --- tools/tests/liskov/test_parser.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tools/tests/liskov/test_parser.py b/tools/tests/liskov/test_parser.py index c9373c36ec..a8a58ee3bf 100644 --- a/tools/tests/liskov/test_parser.py +++ b/tools/tests/liskov/test_parser.py @@ -22,7 +22,7 @@ from icon4pytools.liskov.parsing.parse import DirectivesParser from .conftest import insert_new_lines, scan_for_directives -from .fortran_samples import MULTIPLE_STENCILS, NO_DIRECTIVES_STENCIL, SINGLE_STENCIL, SINGLE_STENCIL_WITH_COMMENTS +from .fortran_samples import MULTIPLE_STENCILS, NO_DIRECTIVES_STENCIL, SINGLE_STENCIL, SINGLE_STENCIL_WITH_COMMENTS, FUSED_STENCIL def test_parse_no_input(): @@ -76,7 +76,7 @@ def test_parse_single_directive(directive, string, startln, endln, expected_cont @mark.parametrize( "stencil, num_directives, num_content", - [(SINGLE_STENCIL, 9, 8), (SINGLE_STENCIL_WITH_COMMENTS, 9, 8), (MULTIPLE_STENCILS, 11, 7)], + [(SINGLE_STENCIL, 9, 8), (SINGLE_STENCIL_WITH_COMMENTS, 9, 8), (MULTIPLE_STENCILS, 11, 7), (FUSED_STENCIL, 9, 7)], ) def test_file_parsing(make_f90_tmpfile, stencil, num_directives, num_content): fpath = make_f90_tmpfile(content=stencil) From cb8068332d2689dcfa9f5468c886ca14f8571b57 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 26 Jul 2023 13:39:22 +0200 Subject: [PATCH 41/49] WIP --- tools/tests/liskov/fortran_samples.py | 135 ++++++++++++++++++++++++++ tools/tests/liskov/test_cli.py | 2 + tools/tests/liskov/test_validation.py | 2 +- 3 files changed, 138 insertions(+), 1 deletion(-) diff --git a/tools/tests/liskov/fortran_samples.py b/tools/tests/liskov/fortran_samples.py index 696768d830..ce9bdd2975 100644 --- a/tools/tests/liskov/fortran_samples.py +++ b/tools/tests/liskov/fortran_samples.py @@ -309,6 +309,141 @@ !$DSL END FUSED STENCIL(name=calculate_diagnostic_quantities_for_turbulence) """ + +FUSED_STENCILS = """\ + !$DSL IMPORTS() + + !$DSL INSERT(INTEGER :: start_interior_idx_c, end_interior_idx_c, start_nudging_idx_c, end_halo_1_idx_c) + + !$DSL DECLARE(kh_smag_e=nproma,p_patch%nlev,p_patch%nblks_e; & + !$DSL kh_smag_ec=nproma,p_patch%nlev,p_patch%nblks_e; & + !$DSL z_nabla2_e=nproma,p_patch%nlev,p_patch%nblks_e; & + !$DSL kh_c=nproma,p_patch%nlev; & + !$DSL div=nproma,p_patch%nlev; & + !$DSL div_ic=nproma,p_patch%nlev,p_patch%nblks_c; & + !$DSL hdef_ic=nproma,p_patch%nlev,p_patch%nblks_c; & + !$DSL z_nabla4_e2=nproma,p_patch%nlev; & + !$DSL vn=nproma,p_patch%nlev,p_patch%nblks_e; & + !$DSL z_nabla2_c=nproma,p_patch%nlev,p_patch%nblks_e; & + !$DSL dwdx=nproma,p_patch%nlev,p_patch%nblks_c; & + !$DSL dwdy=nproma,p_patch%nlev,p_patch%nblks_c; & + !$DSL w=nproma,p_patch%nlev,p_patch%nblks_c; & + !$DSL enh_diffu_3d=nproma,p_patch%nlev,p_patch%nblks_c; & + !$DSL z_temp=nproma,p_patch%nlev,p_patch%nblks_c; & + !$DSL theta_v=nproma,p_patch%nlev,p_patch%nblks_c; & + !$DSL exner=nproma,p_patch%nlev,p_patch%nblks_c) + + !$DSL INSERT(start_2nd_nudge_line_idx_e = i_startidx) + !$DSL INSERT(end_interior_idx_e = i_endidx) + ! Compute nabla4(v) + + !$DSL START FUSED STENCIL(name=apply_diffusion_to_vn; & + !$DSL u_vert=u_vert(:,:,1); & + !$DSL v_vert=v_vert(:,:,1); & + !$DSL primal_normal_vert_v1=p_patch%edges%primal_normal_vert_x(:,:,1); & + !$DSL primal_normal_vert_v2=p_patch%edges%primal_normal_vert_y(:,:,1); & + !$DSL z_nabla2_e=z_nabla2_e(:,:,1); & + !$DSL inv_vert_vert_length=p_patch%edges%inv_vert_vert_length(:,1); & + !$DSL inv_primal_edge_length=p_patch%edges%inv_primal_edge_length(:,1); & + !$DSL area_edge=p_patch%edges%area_edge(:,1); & + !$DSL kh_smag_e=kh_smag_e(:,:,1); & + !$DSL diff_multfac_vn=diff_multfac_vn(:); & + !$DSL nudgecoeff_e=p_int%nudgecoeff_e(:,1); & + !$DSL vn=p_nh_prog%vn(:,:,1); & + !$DSL horz_idx=horz_idx(:); & + !$DSL nudgezone_diff=nudgezone_diff; & + !$DSL fac_bdydiff_v=fac_bdydiff_v; & + !$DSL start_2nd_nudge_line_idx_e=start_2nd_nudge_line_idx_e-1; & + !$DSL limited_area=l_limited_area; & + !$DSL vn_rel_tol=1e-11_wp; & + !$DSL vertical_lower=1; & + !$DSL vertical_upper=nlev; & + !$DSL horizontal_lower=start_bdydiff_idx_e; & + !$DSL horizontal_upper=end_interior_idx_e) + + !$DSL START STENCIL(name=calculate_nabla4; u_vert=u_vert(:,:,1); v_vert=v_vert(:,:,1); & + !$DSL primal_normal_vert_v1=p_patch%edges%primal_normal_vert_x(:,:,1); & + !$DSL primal_normal_vert_v2=p_patch%edges%primal_normal_vert_y(:,:,1); & + !$DSL z_nabla2_e=z_nabla2_e(:,:,1); inv_vert_vert_length=p_patch%edges%inv_vert_vert_length(:,1); & + !$DSL inv_primal_edge_length=p_patch%edges%inv_primal_edge_length(:,1); z_nabla4_e2_abs_tol=1e-27_wp; & + !$DSL z_nabla4_e2=z_nabla4_e2(:, :); vertical_lower=1; vertical_upper=nlev; horizontal_lower=i_startidx; & + !$DSL horizontal_upper=i_endidx) + + !$DSL END STENCIL(name=calculate_nabla4) + + !$DSL START STENCIL(name=apply_nabla2_and_nabla4_to_vn; nudgezone_diff=nudgezone_diff; area_edge=p_patch%edges%area_edge(:,1); & + !$DSL kh_smag_e=kh_smag_e(:,:,1); z_nabla2_e=z_nabla2_e(:,:,1); z_nabla4_e2=z_nabla4_e2(:,:); & + !$DSL diff_multfac_vn=diff_multfac_vn(:); nudgecoeff_e=p_int%nudgecoeff_e(:,1); vn=p_nh_prog%vn(:,:,1); vn_rel_tol=1e-11_wp; & + !$DSL vertical_lower=1; vertical_upper=nlev; horizontal_lower=i_startidx; horizontal_upper=i_endidx) + + !$DSL END STENCIL(name=apply_nabla2_and_nabla4_to_vn) + + !$DSL START STENCIL(name=apply_nabla2_and_nabla4_global_to_vn; area_edge=p_patch%edges%area_edge(:,1); kh_smag_e=kh_smag_e(:,:,1); & + !$DSL z_nabla2_e=z_nabla2_e(:,:,1); z_nabla4_e2=z_nabla4_e2(:,:); diff_multfac_vn=diff_multfac_vn(:); vn=p_nh_prog%vn(:,:,1); & + !$DSL vn_rel_tol=1e-10_wp; vertical_lower=1; vertical_upper=nlev; horizontal_lower=i_startidx; horizontal_upper=i_endidx) + + !$DSL END STENCIL(name=apply_nabla2_and_nabla4_global_to_vn) + + !$DSL INSERT(start_bdydiff_idx_e = i_startidx) + + !$DSL START STENCIL(name=apply_nabla2_to_vn_in_lateral_boundary; z_nabla2_e=z_nabla2_e(:,:,1); area_edge=p_patch%edges%area_edge(:,1); & + !$DSL fac_bdydiff_v=fac_bdydiff_v; vn=p_nh_prog%vn(:,:,1); vn_abs_tol=1e-14_wp; vertical_lower=1; & + !$DSL vertical_upper=nlev; horizontal_lower=i_startidx; horizontal_upper=i_endidx) + + !$DSL END FUSED STENCIL(name=apply_diffusion_to_vn) + + + !$DSL INSERT(start_nudging_idx_c = i_startidx) + !$DSL INSERT(end_halo_1_idx_c = i_endidx) + + !$DSL INSERT(!$ACC PARALLEL IF( i_am_accel_node ) DEFAULT(PRESENT) ASYNC(1)) + !$DSL INSERT(w_old(:,:,:) = p_nh_prog%w(:,:,:)) + !$DSL INSERT(!$ACC END PARALLEL) + + !$DSL START FUSED STENCIL(name=apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance; & + !$DSL area=p_patch%cells%area(:,1); geofac_grg_x=p_int%geofac_grg(:,:,1,1); & + !$DSL geofac_grg_y=p_int%geofac_grg(:,:,1,2); geofac_n2s=p_int%geofac_n2s(:,:,1); & + !$DSL w_old=w_old(:,:,1); w=p_nh_prog%w(:,:,1); diff_multfac_w=diff_multfac_w; & + !$DSL diff_multfac_n2w=diff_multfac_n2w(:); vert_idx=vert_idx(:); & + !$DSL horz_idx=horz_idx(:); nrdmax=nrdmax(jg); interior_idx=start_interior_idx_c-1; & + !$DSL halo_idx=end_interior_idx_c; dwdx=p_nh_diag%dwdx(:,:,1); & + !$DSL dwdy=p_nh_diag%dwdy(:,:,1); & + !$DSL w_rel_tol=1e-09_wp; dwdx_rel_tol=1e-09_wp; dwdy_abs_tol=1e-09_wp; & + !$DSL vertical_lower=1; vertical_upper=nlev; horizontal_lower=start_nudging_idx_c; & + !$DSL horizontal_upper=end_halo_1_idx_c) + + !$DSL START STENCIL(name=calculate_nabla2_for_w; w=p_nh_prog%w(:,:,1); geofac_n2s=p_int%geofac_n2s(:,:,1); & + !$DSL z_nabla2_c=z_nabla2_c(:,:,1); z_nabla2_c_abs_tol=1e-21_wp; vertical_lower=1; vertical_upper=nlev; & + !$DSL horizontal_lower=i_startidx; horizontal_upper=i_endidx) + + !$DSL END STENCIL(name=calculate_nabla2_for_w) + + !$DSL START STENCIL(name=calculate_horizontal_gradients_for_turbulence; w=p_nh_prog%w(:,:,1); geofac_grg_x=p_int%geofac_grg(:,:,1,1); geofac_grg_y=p_int%geofac_grg(:,:,1,2); & + !$DSL dwdx=p_nh_diag%dwdx(:,:,1); dwdy=p_nh_diag%dwdy(:,:,1); dwdx_rel_tol=1e-09_wp; dwdy_rel_tol=1e-09_wp; vertical_lower=2; & + !$DSL vertical_upper=nlev; horizontal_lower=i_startidx; horizontal_upper=i_endidx) + + !$DSL END STENCIL(name=calculate_horizontal_gradients_for_turbulence) + + !$DSL INSERT(start_interior_idx_c = i_startidx) + !$DSL INSERT(end_interior_idx_c = i_endidx) + + !$DSL START STENCIL(name=apply_nabla2_to_w; diff_multfac_w=diff_multfac_w; area=p_patch%cells%area(:,1); & + !$DSL z_nabla2_c=z_nabla2_c(:,:,1); geofac_n2s=p_int%geofac_n2s(:,:,1); w=p_nh_prog%w(:,:,1); & + !$DSL w_abs_tol=1e-15_wp; vertical_lower=1; vertical_upper=nlev; horizontal_lower=i_startidx; & + !$DSL horizontal_upper=i_endidx) + + !$DSL END STENCIL(name=apply_nabla2_to_w) + + !$DSL START STENCIL(name=apply_nabla2_to_w_in_upper_damping_layer; w=p_nh_prog%w(:,:,1); diff_multfac_n2w=diff_multfac_n2w(:); & + !$DSL cell_area=p_patch%cells%area(:,1); z_nabla2_c=z_nabla2_c(:,:,1); vertical_lower=2; w_abs_tol=1e-16_wp; w_rel_tol=1e-10_wp; & + !$DSL vertical_upper=nrdmax(jg); horizontal_lower=i_startidx; horizontal_upper=i_endidx) + + !$DSL END STENCIL(name=apply_nabla2_to_w_in_upper_damping_layer) + + !$DSL END FUSED STENCIL(name=apply_diffusion_to_w_and_compute_horizontal_gradients_for_turbulance) + """ + + FREE_FORM_STENCIL = """\ !$DSL IMPORTS() diff --git a/tools/tests/liskov/test_cli.py b/tools/tests/liskov/test_cli.py index 87c591b3f5..5944c4f006 100644 --- a/tools/tests/liskov/test_cli.py +++ b/tools/tests/liskov/test_cli.py @@ -26,6 +26,7 @@ SINGLE_STENCIL, SINGLE_STENCIL_WITH_COMMENTS, FUSED_STENCIL, + FUSED_STENCILS, ) @@ -44,6 +45,7 @@ def outfile(tmp_path): ("FREE_FORM", FREE_FORM_STENCIL), ("MULTIPLE", MULTIPLE_STENCILS), ("FUSED", FUSED_STENCIL), + ("FUSEDS", FUSED_STENCILS), ] flags = {"serialise": ["--multinode"], "integrate": ["-p", "-m", "-f", "-u"]} diff --git a/tools/tests/liskov/test_validation.py b/tools/tests/liskov/test_validation.py index 0e5c8eb531..d8d4f633bc 100644 --- a/tools/tests/liskov/test_validation.py +++ b/tools/tests/liskov/test_validation.py @@ -24,7 +24,7 @@ from icon4pytools.liskov.parsing.validation import DirectiveSyntaxValidator from .conftest import insert_new_lines, scan_for_directives -from .fortran_samples import MULTIPLE_STENCILS, SINGLE_STENCIL, SINGLE_STENCIL_WITH_COMMENTS +from .fortran_samples import MULTIPLE_STENCILS, SINGLE_STENCIL, SINGLE_STENCIL_WITH_COMMENTS,FUSED_STENCIL @mark.parametrize( From 3e9db8c18915764f9aab8e173116b807a182d475 Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Thu, 27 Jul 2023 12:23:19 +0200 Subject: [PATCH 42/49] fix test --- tools/tests/liskov/fortran_samples.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tools/tests/liskov/fortran_samples.py b/tools/tests/liskov/fortran_samples.py index ce9bdd2975..14faba852c 100644 --- a/tools/tests/liskov/fortran_samples.py +++ b/tools/tests/liskov/fortran_samples.py @@ -390,6 +390,8 @@ !$DSL fac_bdydiff_v=fac_bdydiff_v; vn=p_nh_prog%vn(:,:,1); vn_abs_tol=1e-14_wp; vertical_lower=1; & !$DSL vertical_upper=nlev; horizontal_lower=i_startidx; horizontal_upper=i_endidx) + !$DSL END STENCIL(name=apply_nabla2_to_vn_in_lateral_boundary) + !$DSL END FUSED STENCIL(name=apply_diffusion_to_vn) From 0fb7dabe8b50dc55260db35b261058f7db16467b Mon Sep 17 00:00:00 2001 From: Daniel Hupp Date: Wed, 2 Aug 2023 13:09:49 +0200 Subject: [PATCH 43/49] add one test --- .../liskov/codegen/integration/interface.py | 9 ++++++++- tools/tests/liskov/test_directives_deserialiser.py | 11 +++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/interface.py b/tools/src/icon4pytools/liskov/codegen/integration/interface.py index 4dd1925316..c09f600bac 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/interface.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/interface.py @@ -121,8 +121,15 @@ def __init__(self, startStencil: StartStencilData): @dataclass class EndDeleteData(CodeGenInput): - def __init__(self, endStencil: EndStencilData): + def __init__(self, endStencil: EndStencilData=None, startln:int=None): + if endStencil is None and startln is None: + self.startln = 0 + elif endStencil is not None and startln is not None: + raise Exception() + elif endStencil is not None: self.startln = endStencil.startln + elif startln is not None: + self.startln = startln @dataclass diff --git a/tools/tests/liskov/test_directives_deserialiser.py b/tools/tests/liskov/test_directives_deserialiser.py index 13a235f6f5..9e8f1ee992 100644 --- a/tools/tests/liskov/test_directives_deserialiser.py +++ b/tools/tests/liskov/test_directives_deserialiser.py @@ -22,11 +22,13 @@ EndIfDataFactory, EndProfileDataFactory, EndStencilDataFactory, + EndDeleteDataFactory, ImportsDataFactory, InsertDataFactory, StartCreateDataFactory, StartProfileDataFactory, StartStencilDataFactory, + StartFusedStencilDataFactory, ) from icon4pytools.liskov.codegen.integration.interface import ( BoundsData, @@ -35,6 +37,7 @@ EndIfData, EndProfileData, EndStencilData, + EndDeleteData, FieldAssociationData, ImportsData, InsertData, @@ -83,6 +86,14 @@ 5, EndProfileData, ), + ( + EndDeleteDataFactory, + ts.EndDelete, + "END DELETE", + 6, + 6, + EndDeleteData, + ), ], ) def test_data_factories_no_args(factory_class, directive_type, string, startln, endln, expected): From 2211331fe994097d1d94f9061325f7c62b0adcc3 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Tue, 8 Aug 2023 17:20:17 +0200 Subject: [PATCH 44/49] Refactor StencilTransformer --- .../liskov/codegen/integration/interface.py | 18 +-- .../liskov/codegen/integration/template.py | 5 +- .../icon4pytools/liskov/parsing/transform.py | 109 +++++++++++------- .../liskov/pipeline/collection.py | 14 ++- tools/tests/liskov/fortran_samples.py | 4 +- tools/tests/liskov/test_cli.py | 8 +- tools/tests/liskov/test_parser.py | 20 +++- tools/tests/liskov/test_validation.py | 8 +- 8 files changed, 115 insertions(+), 71 deletions(-) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/interface.py b/tools/src/icon4pytools/liskov/codegen/integration/interface.py index c09f600bac..b2dbb5b4df 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/interface.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/interface.py @@ -121,15 +121,15 @@ def __init__(self, startStencil: StartStencilData): @dataclass class EndDeleteData(CodeGenInput): - def __init__(self, endStencil: EndStencilData=None, startln:int=None): - if endStencil is None and startln is None: - self.startln = 0 - elif endStencil is not None and startln is not None: - raise Exception() - elif endStencil is not None: - self.startln = endStencil.startln - elif startln is not None: - self.startln = startln + def __init__(self, endStencil: EndStencilData = None, startln: int = None): + if endStencil is None and startln is None: + self.startln = 0 + elif endStencil is not None and startln is not None: + raise Exception() + elif endStencil is not None: + self.startln = endStencil.startln + elif startln is not None: + self.startln = startln @dataclass diff --git a/tools/src/icon4pytools/liskov/codegen/integration/template.py b/tools/src/icon4pytools/liskov/codegen/integration/template.py index 3fce6ff263..0a4c0895ec 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/template.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/template.py @@ -83,7 +83,6 @@ def get_array_dims(association: str) -> str: return "".join(list(dims)) - class MetadataStatement(eve.Node): metadata: CodeMetadata @@ -102,6 +101,7 @@ class MetadataStatementGenerator(TemplatedGenerator): """ ) + class EndBasicStencilStatement(eve.Node): name: str = eve.datamodels.field(init=False) input_fields: InputFields = eve.datamodels.field(init=False) @@ -292,6 +292,7 @@ class CopyDeclaration(Declaration): lh_index: str rh_index: str + def _make_copy_declaration(f: Field) -> CopyDeclaration: if f.dims is None: raise UndeclaredFieldError(f"{f.variable} was not declared!") @@ -320,6 +321,7 @@ def _make_copy_declaration(f: Field) -> CopyDeclaration: rh_index=rh_idx, ) + class DeclareStatement(eve.Node): declare_data: DeclareData declarations: list[Declaration] = eve.datamodels.field(init=False) @@ -352,6 +354,7 @@ def __post_init__(self) -> None: # type: ignore self.copy_declarations = [_make_copy_declaration(f) for f in all_fields if f.out] self.acc_present = "PRESENT" if self.stencil_data.acc_present else "NONE" + class StartFusedStencilStatement(eve.Node): stencil_data: StartFusedStencilData copy_declarations: list[CopyDeclaration] = eve.datamodels.field(init=False) diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index 6b96840f7b..4ff17e4cbf 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -15,67 +15,92 @@ import icon4pytools.liskov.parsing.types as ts from icon4pytools.common.logger import setup_logger from icon4pytools.liskov.codegen.integration.interface import ( - EndDeleteData, IntegrationCodeInterface, - StartDeleteData, + StartStencilData, + UnusedDirective, + EndStencilData, + StartFusedStencilData, + EndFusedStencilData, ) from icon4pytools.liskov.pipeline.definition import Step - logger = setup_logger(__name__) -class TransformFuseStencils(Step): +class StencilTransformer(Step): def __init__(self, parsed: IntegrationCodeInterface, fused: bool) -> None: self.parsed = parsed self.fused = fused def __call__(self, data: Any = None) -> ts.ParsedDict: - """Parse the directives and return a dictionary of parsed directives and their associated content. + """Transform stencils in the parse tree based on the 'fused' flag, transforming or removing as necessary. + + This method processes stencils present in the 'parsed' object according to the 'fused' + flag. If 'fused' is True, it identifies and processes stencils that are eligible for + deletion. If 'fused' is False, it removes fused stencils. + + Args: + data (Any): Optional data to be passed. Default is None. Returns: - ParsedType: Dictionary of parsed directives and their associated content. + ts.ParsedDict: The parsed directives along with any modifications applied. """ if self.fused: - logger.info("Transforming single stencils.") + logger.info("Transforming stencils for deletion.") + self._process_stencils_for_deletion() + else: + logger.info("Removing fused stencils.") + self._remove_fused_stencils() + + return self.parsed - removeStartStencil = [] - removeEndStencil = [] + def _process_stencils_for_deletion(self) -> None: + stencils_to_remove = [] - for startFused, endFused in zip( - self.parsed.StartFusedStencil, self.parsed.EndFusedStencil, strict=True + for start_fused, end_fused in zip( + self.parsed.StartFusedStencil, self.parsed.EndFusedStencil, strict=True + ): + for start_single, end_single in zip( + self.parsed.StartStencil, self.parsed.EndStencil, strict=True ): - for startSingle, endSingle in zip( - self.parsed.StartStencil, self.parsed.EndStencil, strict=True - ): - if ( - startFused.startln < startSingle.startln - and startSingle.startln < endFused.startln - and startFused.startln < endSingle.startln - and endSingle.startln < endFused.startln - ): - try: - self.parsed.StartDelete.append(StartDeleteData(startSingle)) - except AttributeError: - self.parsed.StartDelete = [StartDeleteData(startSingle)] - try: - self.parsed.EndDelete.append(EndDeleteData(endSingle)) - except AttributeError: - self.parsed.EndDelete = [EndDeleteData(endSingle)] - removeStartStencil.append(startSingle) - removeEndStencil.append(endSingle) - - self.parsed.StartStencil = [ - x for x in self.parsed.StartStencil if x not in removeStartStencil - ] - self.parsed.EndStencil = [ - x for x in self.parsed.EndStencil if x not in removeEndStencil - ] + if self._stencil_is_removable(start_fused, end_fused, start_single, end_single): + self._add_to_delete_directives(start_single, end_single) + stencils_to_remove += [start_single, end_single] - else: - logger.info("Removing fused stencils.") + self._remove_stencils(stencils_to_remove) - self.parsed.StartFusedStencil = [] - self.parsed.EndFusedStencil = [] + def _stencil_is_removable( + self, + start_fused: StartFusedStencilData, + end_fused: EndFusedStencilData, + start_single: StartStencilData, + end_single: EndStencilData, + ) -> bool: + return ( + start_fused.startln < start_single.startln + and start_single.startln < end_fused.startln + and start_fused.startln < end_single.startln + and end_single.startln < end_fused.startln + ) - return self.parsed + def _add_to_delete_directives( + self, start_single: StartStencilData, end_single: EndStencilData + ) -> None: + for attr, param in zip(["StartDelete", "EndDelete"], [start_single, end_single]): + directive = getattr(self.parsed, attr) + if directive == UnusedDirective: + directive = [] + directive.append(param) + setattr(self.parsed, attr, directive) + + def _remove_stencils(self, stencils_to_remove: list[StartStencilData | EndStencilData]) -> None: + attributes_to_modify = ["StartStencil", "EndStencil"] + + for attr_name in attributes_to_modify: + current_stencil_list = getattr(self.parsed, attr_name) + modified_stencil_list = [_ for _ in current_stencil_list if _ not in stencils_to_remove] + setattr(self.parsed, attr_name, modified_stencil_list) + + def _remove_fused_stencils(self) -> None: + self.parsed.StartFusedStencil = [] + self.parsed.EndFusedStencil = [] diff --git a/tools/src/icon4pytools/liskov/pipeline/collection.py b/tools/src/icon4pytools/liskov/pipeline/collection.py index ba9b776b3f..8af70e54ac 100644 --- a/tools/src/icon4pytools/liskov/pipeline/collection.py +++ b/tools/src/icon4pytools/liskov/pipeline/collection.py @@ -21,7 +21,7 @@ from icon4pytools.liskov.external.gt4py import UpdateFieldsWithGt4PyStencils from icon4pytools.liskov.parsing.parse import DirectivesParser from icon4pytools.liskov.parsing.scan import DirectivesScanner -from icon4pytools.liskov.parsing.transform import TransformFuseStencils +from icon4pytools.liskov.parsing.transform import StencilTransformer from icon4pytools.liskov.pipeline.definition import Step, linear_pipeline @@ -70,16 +70,20 @@ def parse_fortran_file( @linear_pipeline def process_stencils(parsed: IntegrationCodeInterface, fused: bool) -> list[Step]: - """Execute a pipeline to transform stencils to fused or unfused execution, and a pipeline. + """Execute a linear pipeline to transform stencils and produce either fused or unfused execution. + + This function takes an input `parsed` object of type `IntegrationCodeInterface` and a `fused` boolean flag. + It then executes a linear pipeline, consisting of two steps: transformation of stencils for fusion or unfusion, + and updating fields with information from GT4Py stencils. Args: - parsed: The input IntegrationCodeInterface object. - fused: The input mode if fused or unfused + parsed (IntegrationCodeInterface): The input object containing parsed integration code. + fused (bool): A boolean flag indicating whether to produce fused (True) or unfused (False) execution. Returns: The updated and transformed object with fields containing information from GT4Py stencils. """ - return [TransformFuseStencils(parsed, fused), UpdateFieldsWithGt4PyStencils(parsed)] + return [StencilTransformer(parsed, fused), UpdateFieldsWithGt4PyStencils(parsed)] @linear_pipeline diff --git a/tools/tests/liskov/fortran_samples.py b/tools/tests/liskov/fortran_samples.py index 14faba852c..50e5b717a8 100644 --- a/tools/tests/liskov/fortran_samples.py +++ b/tools/tests/liskov/fortran_samples.py @@ -261,7 +261,7 @@ """ -FUSED_STENCIL = """\ +SINGLE_FUSED = """\ !$DSL IMPORTS() !$DSL INSERT(INTEGER :: start_interior_idx_c, end_interior_idx_c, start_nudging_idx_c, end_halo_1_idx_c) @@ -310,7 +310,7 @@ """ -FUSED_STENCILS = """\ +MULTIPLE_FUSED = """\ !$DSL IMPORTS() !$DSL INSERT(INTEGER :: start_interior_idx_c, end_interior_idx_c, start_nudging_idx_c, end_halo_1_idx_c) diff --git a/tools/tests/liskov/test_cli.py b/tools/tests/liskov/test_cli.py index 5944c4f006..7bab42fb1c 100644 --- a/tools/tests/liskov/test_cli.py +++ b/tools/tests/liskov/test_cli.py @@ -25,8 +25,8 @@ NO_DIRECTIVES_STENCIL, SINGLE_STENCIL, SINGLE_STENCIL_WITH_COMMENTS, - FUSED_STENCIL, - FUSED_STENCILS, + SINGLE_FUSED, + MULTIPLE_FUSED, ) @@ -44,8 +44,8 @@ def outfile(tmp_path): ("CONSECUTIVE", CONSECUTIVE_STENCIL), ("FREE_FORM", FREE_FORM_STENCIL), ("MULTIPLE", MULTIPLE_STENCILS), - ("FUSED", FUSED_STENCIL), - ("FUSEDS", FUSED_STENCILS), + ("SINGLE_FUSED", SINGLE_FUSED), + ("MULTIPLE_FUSED", MULTIPLE_FUSED), ] flags = {"serialise": ["--multinode"], "integrate": ["-p", "-m", "-f", "-u"]} diff --git a/tools/tests/liskov/test_parser.py b/tools/tests/liskov/test_parser.py index a8a58ee3bf..deb2c4a468 100644 --- a/tools/tests/liskov/test_parser.py +++ b/tools/tests/liskov/test_parser.py @@ -14,15 +14,20 @@ from collections import defaultdict -import pytest -from pytest import mark - import icon4pytools.liskov.parsing.parse as ts +import pytest from icon4pytools.liskov.parsing.exceptions import UnsupportedDirectiveError from icon4pytools.liskov.parsing.parse import DirectivesParser +from pytest import mark from .conftest import insert_new_lines, scan_for_directives -from .fortran_samples import MULTIPLE_STENCILS, NO_DIRECTIVES_STENCIL, SINGLE_STENCIL, SINGLE_STENCIL_WITH_COMMENTS, FUSED_STENCIL +from .fortran_samples import ( + MULTIPLE_STENCILS, + NO_DIRECTIVES_STENCIL, + SINGLE_STENCIL, + SINGLE_STENCIL_WITH_COMMENTS, + SINGLE_FUSED, +) def test_parse_no_input(): @@ -76,7 +81,12 @@ def test_parse_single_directive(directive, string, startln, endln, expected_cont @mark.parametrize( "stencil, num_directives, num_content", - [(SINGLE_STENCIL, 9, 8), (SINGLE_STENCIL_WITH_COMMENTS, 9, 8), (MULTIPLE_STENCILS, 11, 7), (FUSED_STENCIL, 9, 7)], + [ + (SINGLE_STENCIL, 9, 8), + (SINGLE_STENCIL_WITH_COMMENTS, 9, 8), + (MULTIPLE_STENCILS, 11, 7), + (SINGLE_FUSED, 9, 7), + ], ) def test_file_parsing(make_f90_tmpfile, stencil, num_directives, num_content): fpath = make_f90_tmpfile(content=stencil) diff --git a/tools/tests/liskov/test_validation.py b/tools/tests/liskov/test_validation.py index d8d4f633bc..e287f6ac9f 100644 --- a/tools/tests/liskov/test_validation.py +++ b/tools/tests/liskov/test_validation.py @@ -12,8 +12,6 @@ # SPDX-License-Identifier: GPL-3.0-or-later import pytest -from pytest import mark - from icon4pytools.liskov.parsing.exceptions import ( DirectiveSyntaxError, RepeatedDirectiveError, @@ -22,9 +20,13 @@ ) from icon4pytools.liskov.parsing.parse import Declare, DirectivesParser, Imports, StartStencil from icon4pytools.liskov.parsing.validation import DirectiveSyntaxValidator +from pytest import mark from .conftest import insert_new_lines, scan_for_directives -from .fortran_samples import MULTIPLE_STENCILS, SINGLE_STENCIL, SINGLE_STENCIL_WITH_COMMENTS,FUSED_STENCIL +from .fortran_samples import ( + MULTIPLE_STENCILS, + SINGLE_STENCIL_WITH_COMMENTS, +) @mark.parametrize( From aa9f5fb52897de77139e00d61f821e6d4656e9d8 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 9 Aug 2023 10:33:48 +0200 Subject: [PATCH 45/49] simplify interface --- .../liskov/codegen/integration/interface.py | 27 +++++++------------ .../liskov/codegen/integration/template.py | 4 +-- .../src/icon4pytools/liskov/external/gt4py.py | 6 ++--- tools/src/icon4pytools/liskov/parsing/scan.py | 4 ++- .../icon4pytools/liskov/parsing/transform.py | 14 +++++++--- tools/tests/liskov/fortran_samples.py | 3 ++- tools/tests/liskov/test_cli.py | 14 +++++----- 7 files changed, 36 insertions(+), 36 deletions(-) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/interface.py b/tools/src/icon4pytools/liskov/codegen/integration/interface.py index b2dbb5b4df..93903cab78 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/interface.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/interface.py @@ -78,7 +78,7 @@ class EndProfileData(CodeGenInput): @dataclass -class StartBasicStencilData(CodeGenInput): +class BaseStartStencilData(CodeGenInput): name: str fields: list[FieldAssociationData] acc_present: Optional[bool] @@ -86,50 +86,41 @@ class StartBasicStencilData(CodeGenInput): @dataclass -class StartStencilData(StartBasicStencilData): +class StartStencilData(BaseStartStencilData): mergecopy: Optional[bool] copies: Optional[bool] @dataclass -class StartFusedStencilData(StartBasicStencilData): +class StartFusedStencilData(BaseStartStencilData): ... @dataclass -class EndStencilBaseData(CodeGenInput): +class BaseEndStencilData(CodeGenInput): name: str @dataclass -class EndStencilData(EndStencilBaseData): +class EndStencilData(BaseEndStencilData): noendif: Optional[bool] noprofile: Optional[bool] noaccenddata: Optional[bool] @dataclass -class EndFusedStencilData(EndStencilBaseData): +class EndFusedStencilData(BaseEndStencilData): ... @dataclass class StartDeleteData(CodeGenInput): - def __init__(self, startStencil: StartStencilData): - self.startln = startStencil.startln + startln: int @dataclass -class EndDeleteData(CodeGenInput): - def __init__(self, endStencil: EndStencilData = None, startln: int = None): - if endStencil is None and startln is None: - self.startln = 0 - elif endStencil is not None and startln is not None: - raise Exception() - elif endStencil is not None: - self.startln = endStencil.startln - elif startln is not None: - self.startln = startln +class EndDeleteData(StartDeleteData): + ... @dataclass diff --git a/tools/src/icon4pytools/liskov/codegen/integration/template.py b/tools/src/icon4pytools/liskov/codegen/integration/template.py index 0a4c0895ec..7bf30d55a0 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/template.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/template.py @@ -22,7 +22,7 @@ from icon4pytools.liskov.codegen.integration.exceptions import UndeclaredFieldError from icon4pytools.liskov.codegen.integration.interface import ( DeclareData, - StartBasicStencilData, + BaseStartStencilData, StartFusedStencilData, StartStencilData, ) @@ -447,7 +447,7 @@ class StartFusedStencilStatementGenerator(TemplatedGenerator): class ImportsStatement(eve.Node): - stencils: list[StartBasicStencilData] + stencils: list[BaseStartStencilData] stencil_names: list[str] = eve.datamodels.field(init=False) def __post_init__(self) -> None: # type: ignore diff --git a/tools/src/icon4pytools/liskov/external/gt4py.py b/tools/src/icon4pytools/liskov/external/gt4py.py index 62d1d345b9..d24d46ebe5 100644 --- a/tools/src/icon4pytools/liskov/external/gt4py.py +++ b/tools/src/icon4pytools/liskov/external/gt4py.py @@ -22,7 +22,7 @@ from icon4pytools.icon4pygen.metadata import get_stencil_info from icon4pytools.liskov.codegen.integration.interface import ( IntegrationCodeInterface, - StartBasicStencilData, + BaseStartStencilData, ) from icon4pytools.liskov.external.exceptions import IncompatibleFieldError, UnknownStencilError from icon4pytools.liskov.pipeline.definition import Step @@ -45,7 +45,7 @@ def __call__(self, data: Any = None) -> IntegrationCodeInterface: return self.parsed - def _set_in_out_field(self, startStencil: Sequence[StartBasicStencilData]) -> None: + def _set_in_out_field(self, startStencil: Sequence[BaseStartStencilData]) -> None: for s in startStencil: gt4py_program = self._collect_icon4py_stencil(s.name) gt4py_stencil_info = get_stencil_info(gt4py_program) @@ -54,7 +54,7 @@ def _set_in_out_field(self, startStencil: Sequence[StartBasicStencilData]) -> No try: field_info = gt4py_fields[f.variable] except KeyError: - error_msg = f"Used field variable name ({f.variable}) that is incompatible with the expected field names defined in {s.name} in icon4pytools." + error_msg = f"Used field variable name ({f.variable}) that is incompatible with the expected field names defined in {s.name} in icon4py." raise IncompatibleFieldError(error_msg) f.out = field_info.out f.inp = field_info.inp diff --git a/tools/src/icon4pytools/liskov/parsing/scan.py b/tools/src/icon4pytools/liskov/parsing/scan.py index 95670938e9..a4da42309f 100644 --- a/tools/src/icon4pytools/liskov/parsing/scan.py +++ b/tools/src/icon4pytools/liskov/parsing/scan.py @@ -36,7 +36,9 @@ def __init__(self, input_filepath: Path) -> None: A directive must start with !$DSL ( with the directive arguments delimited by a ;. The directive if on multiple lines must include a & at the end of the line. The directive - must always be closed by a closing bracket ). + must always be closed by a closing bracket ). A directive can be + commented out by using a ! before the directive, + for example, !!$DSL means the directive is disabled. Example: !$DSL IMPORTS() diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index 4ff17e4cbf..9578a0bf06 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -20,7 +20,7 @@ UnusedDirective, EndStencilData, StartFusedStencilData, - EndFusedStencilData, + EndFusedStencilData, StartDeleteData, EndDeleteData, ) from icon4pytools.liskov.pipeline.definition import Step @@ -64,7 +64,7 @@ def _process_stencils_for_deletion(self) -> None: self.parsed.StartStencil, self.parsed.EndStencil, strict=True ): if self._stencil_is_removable(start_fused, end_fused, start_single, end_single): - self._add_to_delete_directives(start_single, end_single) + self._create_delete_directives(start_single, end_single) stencils_to_remove += [start_single, end_single] self._remove_stencils(stencils_to_remove) @@ -83,14 +83,20 @@ def _stencil_is_removable( and end_single.startln < end_fused.startln ) - def _add_to_delete_directives( + def _create_delete_directives( self, start_single: StartStencilData, end_single: EndStencilData ) -> None: for attr, param in zip(["StartDelete", "EndDelete"], [start_single, end_single]): directive = getattr(self.parsed, attr) if directive == UnusedDirective: directive = [] - directive.append(param) + + if attr == "StartDelete": + cls = StartDeleteData + elif attr == "EndDelete": + cls = EndDeleteData + + directive.append(cls(startln=param.startln)) setattr(self.parsed, attr, directive) def _remove_stencils(self, stencils_to_remove: list[StartStencilData | EndStencilData]) -> None: diff --git a/tools/tests/liskov/fortran_samples.py b/tools/tests/liskov/fortran_samples.py index 50e5b717a8..7ed24e5954 100644 --- a/tools/tests/liskov/fortran_samples.py +++ b/tools/tests/liskov/fortran_samples.py @@ -312,7 +312,8 @@ MULTIPLE_FUSED = """\ !$DSL IMPORTS() - + !$DSL START DELETE() + !$DSL END DELETE() !$DSL INSERT(INTEGER :: start_interior_idx_c, end_interior_idx_c, start_nudging_idx_c, end_halo_1_idx_c) !$DSL DECLARE(kh_smag_e=nproma,p_patch%nlev,p_patch%nblks_e; & diff --git a/tools/tests/liskov/test_cli.py b/tools/tests/liskov/test_cli.py index 7bab42fb1c..9046555fbb 100644 --- a/tools/tests/liskov/test_cli.py +++ b/tools/tests/liskov/test_cli.py @@ -38,13 +38,13 @@ def outfile(tmp_path): test_cases = [] files = [ - ("NO_DIRECTIVES", NO_DIRECTIVES_STENCIL), - ("SINGLE", SINGLE_STENCIL), - ("COMMENTS", SINGLE_STENCIL_WITH_COMMENTS), - ("CONSECUTIVE", CONSECUTIVE_STENCIL), - ("FREE_FORM", FREE_FORM_STENCIL), - ("MULTIPLE", MULTIPLE_STENCILS), - ("SINGLE_FUSED", SINGLE_FUSED), + # ("NO_DIRECTIVES", NO_DIRECTIVES_STENCIL), + # ("SINGLE", SINGLE_STENCIL), + # ("COMMENTS", SINGLE_STENCIL_WITH_COMMENTS), + # ("CONSECUTIVE", CONSECUTIVE_STENCIL), + # ("FREE_FORM", FREE_FORM_STENCIL), + # ("MULTIPLE", MULTIPLE_STENCILS), + # ("SINGLE_FUSED", SINGLE_FUSED), ("MULTIPLE_FUSED", MULTIPLE_FUSED), ] From 53127b5f8e0e3ddf1170240d1f24d4e646ca04ea Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 9 Aug 2023 11:18:51 +0200 Subject: [PATCH 46/49] Simplify codegen template --- .../liskov/codegen/integration/template.py | 83 +++---------------- .../icon4pytools/liskov/parsing/transform.py | 4 +- tools/tests/liskov/test_cli.py | 14 ++-- 3 files changed, 23 insertions(+), 78 deletions(-) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/template.py b/tools/src/icon4pytools/liskov/codegen/integration/template.py index 7bf30d55a0..b8ec74b7c9 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/template.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/template.py @@ -128,23 +128,7 @@ def __post_init__(self) -> None: # type: ignore ) -class EndFusedStencilStatementGenerator(TemplatedGenerator): - EndFusedStencilStatement = as_jinja( - """ - call wrap_run_{{ name }}( & - {{ input_fields }} - {{ output_fields }} - {{ tolerance_fields }} - {{ bounds_fields }} - - !$ACC EXIT DATA DELETE( & - {%- for d in _this_node.copy_declarations %} - !$ACC {{ d.variable }}_before {%- if not loop.last -%}, & {% else %} ) & {%- endif -%} - {%- endfor %} - !$ACC IF ( i_am_accel_node ) - """ - ) - +class BaseEndStencilStatementGenerator(TemplatedGenerator): InputFields = as_jinja( """ {%- for field in _this_node.fields %} @@ -205,7 +189,7 @@ def visit_OutputFields(self, out: OutputFields) -> OutputFields: # type: ignore ) -class EndStencilStatementGenerator(TemplatedGenerator): +class EndStencilStatementGenerator(BaseEndStencilStatementGenerator): EndStencilStatement = as_jinja( """ {%- if _this_node.profile %} @@ -224,62 +208,21 @@ class EndStencilStatementGenerator(TemplatedGenerator): """ ) - InputFields = as_jinja( - """ - {%- for field in _this_node.fields %} - {%- if field.out %} - {%- else %} - {{ field.variable }}={{ field.association }},& - {%- endif -%} - {%- endfor %} +class EndFusedStencilStatementGenerator(BaseEndStencilStatementGenerator): + EndFusedStencilStatement = as_jinja( """ - ) + call wrap_run_{{ name }}( & + {{ input_fields }} + {{ output_fields }} + {{ tolerance_fields }} + {{ bounds_fields }} - OutputFields = as_jinja( - """ - {%- for field in _this_node.fields %} - {{ field.variable }}={{ field.association }},& - {{ field.variable }}_before={{ field.variable }}_before{{ field.rh_index }},& + !$ACC EXIT DATA DELETE( & + {%- for d in _this_node.copy_declarations %} + !$ACC {{ d.variable }}_before {%- if not loop.last -%}, & {% else %} ) & {%- endif -%} {%- endfor %} - """ - ) - - def visit_OutputFields(self, out: OutputFields) -> OutputFields: # type: ignore - for f in out.fields: # type: ignore - idx = render_index(f.dims) - split_idx = idx.split(",") - - if len(split_idx) >= 3: - split_idx[-1] = "1" - - f.rh_index = enclose_in_parentheses(",".join(split_idx)) - return self.generic_visit(out) - - ToleranceFields = as_jinja( - """ - {%- if _this_node.fields|length < 1 -%} - - {%- else -%} - - {%- for f in _this_node.fields -%} - {% if f.rel_tol %} - {{ f.variable }}_rel_tol={{ f.rel_tol }}, & - {%- endif -%} - {% if f.abs_tol %} - {{ f.variable }}_abs_tol={{ f.abs_tol }}, & - {% endif %} - {%- endfor -%} - - {%- endif -%} - """ - ) - - BoundsFields = as_jinja( - """vertical_lower={{ vlower }}, & - vertical_upper={{ vupper }}, & - horizontal_lower={{ hlower }}, & - horizontal_upper={{ hupper }}) + !$ACC IF ( i_am_accel_node ) """ ) diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index 9578a0bf06..8f8c5300c4 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -20,7 +20,9 @@ UnusedDirective, EndStencilData, StartFusedStencilData, - EndFusedStencilData, StartDeleteData, EndDeleteData, + EndFusedStencilData, + StartDeleteData, + EndDeleteData, ) from icon4pytools.liskov.pipeline.definition import Step diff --git a/tools/tests/liskov/test_cli.py b/tools/tests/liskov/test_cli.py index 9046555fbb..7bab42fb1c 100644 --- a/tools/tests/liskov/test_cli.py +++ b/tools/tests/liskov/test_cli.py @@ -38,13 +38,13 @@ def outfile(tmp_path): test_cases = [] files = [ - # ("NO_DIRECTIVES", NO_DIRECTIVES_STENCIL), - # ("SINGLE", SINGLE_STENCIL), - # ("COMMENTS", SINGLE_STENCIL_WITH_COMMENTS), - # ("CONSECUTIVE", CONSECUTIVE_STENCIL), - # ("FREE_FORM", FREE_FORM_STENCIL), - # ("MULTIPLE", MULTIPLE_STENCILS), - # ("SINGLE_FUSED", SINGLE_FUSED), + ("NO_DIRECTIVES", NO_DIRECTIVES_STENCIL), + ("SINGLE", SINGLE_STENCIL), + ("COMMENTS", SINGLE_STENCIL_WITH_COMMENTS), + ("CONSECUTIVE", CONSECUTIVE_STENCIL), + ("FREE_FORM", FREE_FORM_STENCIL), + ("MULTIPLE", MULTIPLE_STENCILS), + ("SINGLE_FUSED", SINGLE_FUSED), ("MULTIPLE_FUSED", MULTIPLE_FUSED), ] From 0c645f24a6f552f32a3c468f38e9e519ba415b1f Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 9 Aug 2023 11:51:48 +0200 Subject: [PATCH 47/49] Simplify deserialise step --- .../liskov/codegen/integration/deserialise.py | 122 ++++++++---------- 1 file changed, 56 insertions(+), 66 deletions(-) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py index 9548f7548e..f8a56a30e8 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py @@ -261,6 +261,48 @@ def __call__(self, parsed: ts.ParsedDict) -> list[EndFusedStencilData]: class StartStencilDataFactoryBase(DataFactoryBase): + def create_stencil_data( + self, + parsed: ts.ParsedDict, + field_dimensions: list[dict[str, Any]], + directives: list[ts.ParsedDirective], + directive_cls: Type[ts.ParsedDirective], + dtype: Type[StartStencilData | StartFusedStencilData], + ) -> list[StartStencilData | StartFusedStencilData]: + """Create and return a list of StartStencilData or StartFusedStencilData objects from parsed directives.""" + deserialised = [] + for i, directive in enumerate(directives): + named_args = parsed["content"][directive_cls.__name__][i] + additional_attrs = self._pop_additional_attributes(dtype, named_args) + acc_present = string_to_bool(pop_item_from_dict(named_args, "accpresent", "true")) + stencil_name = _extract_stencil_name(named_args, directive) + bounds = self._make_bounds(named_args) + fields = self._make_fields(named_args, field_dimensions) + fields_w_tolerance = self._update_tolerances(named_args, fields) + + deserialised.append( + dtype( + name=stencil_name, + fields=fields_w_tolerance, + bounds=bounds, + startln=directive.startln, + acc_present=acc_present, + **additional_attrs, + ) + ) + return deserialised + + def _pop_additional_attributes( + self, dtype: Type[StartStencilData | StartFusedStencilData], named_args: dict[str, Any] + ) -> dict: + """Pop and return additional attributes specific to StartStencilData.""" + additional_attrs = {} + if dtype == StartStencilData: + mergecopy = string_to_bool(pop_item_from_dict(named_args, "mergecopy", "false")) + copies = string_to_bool(pop_item_from_dict(named_args, "copies", "true")) + additional_attrs = {"mergecopy": mergecopy, "copies": copies} + return additional_attrs + @staticmethod def _make_bounds(named_args: dict) -> BoundsData: """Extract stencil bounds from directive arguments.""" @@ -348,84 +390,32 @@ def _update_tolerances( return fields -class StartStencilDataFactory(StartStencilDataFactoryBase): - directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.StartStencil - dtype: Type[StartStencilData] = StartStencilData - - def __call__(self, parsed: ts.ParsedDict) -> list[StartStencilData]: - """Create and return a list of StartStencilData objects from the parsed directives. - - Args: - parsed (ParsedDict): Dictionary of parsed directives and their associated content. +class StartFusedStencilDataFactory(StartStencilDataFactoryBase): + directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.StartFusedStencil + dtype: Type[StartFusedStencilData] = StartFusedStencilData - Returns: - List[StartStencilData]: List of StartStencilData objects created from the parsed directives. - """ - deserialised = [] + def __call__(self, parsed: ts.ParsedDict) -> list[StartFusedStencilData]: field_dimensions = flatten_list_of_dicts( [DeclareDataFactory.get_field_dimensions(dim) for dim in parsed["content"]["Declare"]] ) directives = extract_directive(parsed["directives"], self.directive_cls) - for i, directive in enumerate(directives): - named_args = parsed["content"]["StartStencil"][i] - acc_present = string_to_bool(pop_item_from_dict(named_args, "accpresent", "true")) - mergecopy = string_to_bool(pop_item_from_dict(named_args, "mergecopy", "false")) - copies = string_to_bool(pop_item_from_dict(named_args, "copies", "true")) - stencil_name = _extract_stencil_name(named_args, directive) - bounds = self._make_bounds(named_args) - fields = self._make_fields(named_args, field_dimensions) - fields_w_tolerance = self._update_tolerances(named_args, fields) - - deserialised.append( - self.dtype( - name=stencil_name, - fields=fields_w_tolerance, - bounds=bounds, - startln=directive.startln, - acc_present=acc_present, - mergecopy=mergecopy, - copies=copies, - ) - ) - return deserialised - - -class StartFusedStencilDataFactory(StartStencilDataFactoryBase): - directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.StartFusedStencil - dtype: Type[StartFusedStencilData] = StartFusedStencilData + return self.create_stencil_data( + parsed, field_dimensions, directives, self.directive_cls, self.dtype + ) - def __call__(self, parsed: ts.ParsedDict) -> list[StartFusedStencilData]: - """Create and return a list of StartStencilData objects from the parsed directives. - Args: - parsed (ParsedDict): Dictionary of parsed directives and their associated content. +class StartStencilDataFactory(StartStencilDataFactoryBase): + directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.StartStencil + dtype: Type[StartStencilData] = StartStencilData - Returns: - List[StartStencilData]: List of StartStencilData objects created from the parsed directives. - """ - deserialised = [] + def __call__(self, parsed: ts.ParsedDict) -> list[StartStencilData]: field_dimensions = flatten_list_of_dicts( [DeclareDataFactory.get_field_dimensions(dim) for dim in parsed["content"]["Declare"]] ) directives = extract_directive(parsed["directives"], self.directive_cls) - for i, directive in enumerate(directives): - named_args = parsed["content"]["StartFusedStencil"][i] - acc_present = string_to_bool(pop_item_from_dict(named_args, "accpresent", "true")) - stencil_name = _extract_stencil_name(named_args, directive) - bounds = self._make_bounds(named_args) - fields = self._make_fields(named_args, field_dimensions) - fields_w_tolerance = self._update_tolerances(named_args, fields) - - deserialised.append( - self.dtype( - name=stencil_name, - fields=fields_w_tolerance, - bounds=bounds, - startln=directive.startln, - acc_present=acc_present, - ) - ) - return deserialised + return self.create_stencil_data( + parsed, field_dimensions, directives, self.directive_cls, self.dtype + ) class InsertDataFactory(DataFactoryBase): From 16133fb21dec797ab4118cf7f698f311703eaef5 Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 9 Aug 2023 12:01:07 +0200 Subject: [PATCH 48/49] Further simplify deserialiser --- .../liskov/codegen/integration/deserialise.py | 30 ++++++++----------- 1 file changed, 12 insertions(+), 18 deletions(-) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py index f8a56a30e8..f4118c9e95 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/deserialise.py @@ -261,6 +261,18 @@ def __call__(self, parsed: ts.ParsedDict) -> list[EndFusedStencilData]: class StartStencilDataFactoryBase(DataFactoryBase): + directive_cls: Type[ts.ParsedDirective] = None + dtype: Type[StartFusedStencilData] = None + + def __call__(self, parsed: ts.ParsedDict) -> list[StartStencilData]: + field_dimensions = flatten_list_of_dicts( + [DeclareDataFactory.get_field_dimensions(dim) for dim in parsed["content"]["Declare"]] + ) + directives = extract_directive(parsed["directives"], self.directive_cls) + return self.create_stencil_data( + parsed, field_dimensions, directives, self.directive_cls, self.dtype + ) + def create_stencil_data( self, parsed: ts.ParsedDict, @@ -394,29 +406,11 @@ class StartFusedStencilDataFactory(StartStencilDataFactoryBase): directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.StartFusedStencil dtype: Type[StartFusedStencilData] = StartFusedStencilData - def __call__(self, parsed: ts.ParsedDict) -> list[StartFusedStencilData]: - field_dimensions = flatten_list_of_dicts( - [DeclareDataFactory.get_field_dimensions(dim) for dim in parsed["content"]["Declare"]] - ) - directives = extract_directive(parsed["directives"], self.directive_cls) - return self.create_stencil_data( - parsed, field_dimensions, directives, self.directive_cls, self.dtype - ) - class StartStencilDataFactory(StartStencilDataFactoryBase): directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.StartStencil dtype: Type[StartStencilData] = StartStencilData - def __call__(self, parsed: ts.ParsedDict) -> list[StartStencilData]: - field_dimensions = flatten_list_of_dicts( - [DeclareDataFactory.get_field_dimensions(dim) for dim in parsed["content"]["Declare"]] - ) - directives = extract_directive(parsed["directives"], self.directive_cls) - return self.create_stencil_data( - parsed, field_dimensions, directives, self.directive_cls, self.dtype - ) - class InsertDataFactory(DataFactoryBase): directive_cls: Type[ts.ParsedDirective] = icon4pytools.liskov.parsing.parse.Insert From 3a9cec0bd092361c1eb9cbe0300557aeb29542bc Mon Sep 17 00:00:00 2001 From: samkellerhals Date: Wed, 9 Aug 2023 16:46:28 +0200 Subject: [PATCH 49/49] Run precommit --- .../liskov/codegen/integration/template.py | 2 +- tools/src/icon4pytools/liskov/external/gt4py.py | 2 +- tools/src/icon4pytools/liskov/parsing/transform.py | 11 ++++++----- tools/tests/liskov/test_cli.py | 4 ++-- tools/tests/liskov/test_directives_deserialiser.py | 5 ++--- tools/tests/liskov/test_parser.py | 7 ++++--- tools/tests/liskov/test_validation.py | 8 +++----- 7 files changed, 19 insertions(+), 20 deletions(-) diff --git a/tools/src/icon4pytools/liskov/codegen/integration/template.py b/tools/src/icon4pytools/liskov/codegen/integration/template.py index b8ec74b7c9..2f154f82d2 100644 --- a/tools/src/icon4pytools/liskov/codegen/integration/template.py +++ b/tools/src/icon4pytools/liskov/codegen/integration/template.py @@ -21,8 +21,8 @@ from icon4pytools.liskov.codegen.integration.exceptions import UndeclaredFieldError from icon4pytools.liskov.codegen.integration.interface import ( - DeclareData, BaseStartStencilData, + DeclareData, StartFusedStencilData, StartStencilData, ) diff --git a/tools/src/icon4pytools/liskov/external/gt4py.py b/tools/src/icon4pytools/liskov/external/gt4py.py index d24d46ebe5..53b815173f 100644 --- a/tools/src/icon4pytools/liskov/external/gt4py.py +++ b/tools/src/icon4pytools/liskov/external/gt4py.py @@ -21,8 +21,8 @@ from icon4pytools.common.logger import setup_logger from icon4pytools.icon4pygen.metadata import get_stencil_info from icon4pytools.liskov.codegen.integration.interface import ( - IntegrationCodeInterface, BaseStartStencilData, + IntegrationCodeInterface, ) from icon4pytools.liskov.external.exceptions import IncompatibleFieldError, UnknownStencilError from icon4pytools.liskov.pipeline.definition import Step diff --git a/tools/src/icon4pytools/liskov/parsing/transform.py b/tools/src/icon4pytools/liskov/parsing/transform.py index 8f8c5300c4..12fd9c0478 100644 --- a/tools/src/icon4pytools/liskov/parsing/transform.py +++ b/tools/src/icon4pytools/liskov/parsing/transform.py @@ -15,17 +15,18 @@ import icon4pytools.liskov.parsing.types as ts from icon4pytools.common.logger import setup_logger from icon4pytools.liskov.codegen.integration.interface import ( + EndDeleteData, + EndFusedStencilData, + EndStencilData, IntegrationCodeInterface, + StartDeleteData, + StartFusedStencilData, StartStencilData, UnusedDirective, - EndStencilData, - StartFusedStencilData, - EndFusedStencilData, - StartDeleteData, - EndDeleteData, ) from icon4pytools.liskov.pipeline.definition import Step + logger = setup_logger(__name__) diff --git a/tools/tests/liskov/test_cli.py b/tools/tests/liskov/test_cli.py index 7bab42fb1c..292374ba7f 100644 --- a/tools/tests/liskov/test_cli.py +++ b/tools/tests/liskov/test_cli.py @@ -21,12 +21,12 @@ from .fortran_samples import ( CONSECUTIVE_STENCIL, FREE_FORM_STENCIL, + MULTIPLE_FUSED, MULTIPLE_STENCILS, NO_DIRECTIVES_STENCIL, + SINGLE_FUSED, SINGLE_STENCIL, SINGLE_STENCIL_WITH_COMMENTS, - SINGLE_FUSED, - MULTIPLE_FUSED, ) diff --git a/tools/tests/liskov/test_directives_deserialiser.py b/tools/tests/liskov/test_directives_deserialiser.py index 9e8f1ee992..56fb66e57f 100644 --- a/tools/tests/liskov/test_directives_deserialiser.py +++ b/tools/tests/liskov/test_directives_deserialiser.py @@ -19,25 +19,24 @@ from icon4pytools.liskov.codegen.integration.deserialise import ( DeclareDataFactory, EndCreateDataFactory, + EndDeleteDataFactory, EndIfDataFactory, EndProfileDataFactory, EndStencilDataFactory, - EndDeleteDataFactory, ImportsDataFactory, InsertDataFactory, StartCreateDataFactory, StartProfileDataFactory, StartStencilDataFactory, - StartFusedStencilDataFactory, ) from icon4pytools.liskov.codegen.integration.interface import ( BoundsData, DeclareData, EndCreateData, + EndDeleteData, EndIfData, EndProfileData, EndStencilData, - EndDeleteData, FieldAssociationData, ImportsData, InsertData, diff --git a/tools/tests/liskov/test_parser.py b/tools/tests/liskov/test_parser.py index deb2c4a468..0b7094d56c 100644 --- a/tools/tests/liskov/test_parser.py +++ b/tools/tests/liskov/test_parser.py @@ -14,19 +14,20 @@ from collections import defaultdict -import icon4pytools.liskov.parsing.parse as ts import pytest +from pytest import mark + +import icon4pytools.liskov.parsing.parse as ts from icon4pytools.liskov.parsing.exceptions import UnsupportedDirectiveError from icon4pytools.liskov.parsing.parse import DirectivesParser -from pytest import mark from .conftest import insert_new_lines, scan_for_directives from .fortran_samples import ( MULTIPLE_STENCILS, NO_DIRECTIVES_STENCIL, + SINGLE_FUSED, SINGLE_STENCIL, SINGLE_STENCIL_WITH_COMMENTS, - SINGLE_FUSED, ) diff --git a/tools/tests/liskov/test_validation.py b/tools/tests/liskov/test_validation.py index e287f6ac9f..ffa2be1bef 100644 --- a/tools/tests/liskov/test_validation.py +++ b/tools/tests/liskov/test_validation.py @@ -12,6 +12,8 @@ # SPDX-License-Identifier: GPL-3.0-or-later import pytest +from pytest import mark + from icon4pytools.liskov.parsing.exceptions import ( DirectiveSyntaxError, RepeatedDirectiveError, @@ -20,13 +22,9 @@ ) from icon4pytools.liskov.parsing.parse import Declare, DirectivesParser, Imports, StartStencil from icon4pytools.liskov.parsing.validation import DirectiveSyntaxValidator -from pytest import mark from .conftest import insert_new_lines, scan_for_directives -from .fortran_samples import ( - MULTIPLE_STENCILS, - SINGLE_STENCIL_WITH_COMMENTS, -) +from .fortran_samples import MULTIPLE_STENCILS, SINGLE_STENCIL_WITH_COMMENTS @mark.parametrize(