Skip to content
New issue

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

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

Already on GitHub? Sign in to your account

remove column_axis from StencilInfo and generate_gtheader #418

Merged
merged 3 commits into from
Mar 20, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 3 additions & 6 deletions tools/src/icon4pytools/icon4pygen/backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
)

Expand All @@ -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,
)
Expand Down
4 changes: 1 addition & 3 deletions tools/src/icon4pytools/icon4pygen/metadata.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
2 changes: 1 addition & 1 deletion tools/tests/icon4pygen/test_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Loading