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

@@ -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,
         )

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)

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)