From 43d682062b2d669322f798d255b6d544c843a4b8 Mon Sep 17 00:00:00 2001 From: Magdalena Date: Wed, 20 Mar 2024 13:22:18 +0100 Subject: [PATCH] remove column_axis from StencilInfo and generate_gtheader (#418) * remove column_axis from StencilInfo and generate_gtheader * fix test_backend.py * set column_axis=KDim for scan --- tools/src/icon4pytools/icon4pygen/backend.py | 9 +++------ tools/src/icon4pytools/icon4pygen/metadata.py | 4 +--- tools/tests/icon4pygen/test_backend.py | 2 +- 3 files changed, 5 insertions(+), 10 deletions(-) diff --git a/tools/src/icon4pytools/icon4pygen/backend.py b/tools/src/icon4pytools/icon4pygen/backend.py index b401cdfdb0..7c7e833178 100644 --- a/tools/src/icon4pytools/icon4pygen/backend.py +++ b/tools/src/icon4pytools/icon4pygen/backend.py @@ -12,14 +12,13 @@ # SPDX-License-Identifier: GPL-3.0-or-later import warnings from pathlib import Path -from typing import Any, Iterable, List, Optional +from typing import Any, Iterable, List -from gt4py.next import common from gt4py.next.common import Connectivity, Dimension from gt4py.next.iterator import ir as itir from gt4py.next.iterator.transforms import LiftMode from gt4py.next.program_processors.codegens.gtfn import gtfn_module -from icon4py.model.common.dimension import Koff +from icon4py.model.common.dimension import KDim, Koff from icon4pytools.icon4pygen.bindings.utils import write_string from icon4pytools.icon4pygen.metadata import StencilInfo @@ -118,7 +117,6 @@ def check_for_domain_bounds(fencil: itir.FencilDefinition) -> None: def generate_gtheader( fencil: itir.FencilDefinition, offset_provider: dict[str, Connectivity | Dimension], - column_axis: Optional[common.Dimension], imperative: bool, temporaries: bool, **kwargs: Any, @@ -146,7 +144,7 @@ def generate_gtheader( return translation.generate_stencil_source( transformed_fencil, offset_provider=offset_provider, - column_axis=column_axis, + column_axis=KDim, # only used for ScanOperator **kwargs, ) @@ -162,7 +160,6 @@ def __call__(self, outpath: Path, imperative: bool, temporaries: bool) -> None: gtheader = generate_gtheader( fencil=self.stencil_info.itir, offset_provider=self.stencil_info.offset_provider, - column_axis=self.stencil_info.column_axis, imperative=imperative, temporaries=temporaries, ) diff --git a/tools/src/icon4pytools/icon4pygen/metadata.py b/tools/src/icon4pytools/icon4pygen/metadata.py index c90eb4fb24..5c3321871e 100644 --- a/tools/src/icon4pytools/icon4pygen/metadata.py +++ b/tools/src/icon4pytools/icon4pygen/metadata.py @@ -44,7 +44,6 @@ class StencilInfo: fields: dict[str, FieldInfo] connectivity_chains: list[eve.concepts.SymbolRef] offset_provider: dict - column_axis: Optional[Dimension] @dataclass(frozen=True) @@ -236,10 +235,9 @@ def get_stencil_info( offsets = scan_for_offsets(fvprog) itir = fvprog.itir fields = _get_field_infos(fvprog) - column_axis = fvprog._column_axis offset_provider = {} for offset in offsets: offset_provider[offset] = provide_offset(offset, is_global) connectivity_chains = [offset for offset in offsets if offset != Koff.value] - return StencilInfo(itir, fields, connectivity_chains, offset_provider, column_axis) + return StencilInfo(itir, fields, connectivity_chains, offset_provider) diff --git a/tools/tests/icon4pygen/test_backend.py b/tools/tests/icon4pygen/test_backend.py index 4667c5e831..a6c84b798d 100644 --- a/tools/tests/icon4pygen/test_backend.py +++ b/tools/tests/icon4pygen/test_backend.py @@ -65,5 +65,5 @@ def testee_prog( fencil = testee_prog.itir # validate the grid sizes appear in the generated code - gtheader = generate_gtheader(fencil, offset_provider, None, temporaries, imperative) + gtheader = generate_gtheader(fencil, offset_provider, temporaries, imperative) assert search_for_grid_sizes(gtheader)