diff --git a/src/gt4py/cartesian/gtc/dace/oir_to_treeir.py b/src/gt4py/cartesian/gtc/dace/oir_to_treeir.py index c578fea24e..5300585131 100644 --- a/src/gt4py/cartesian/gtc/dace/oir_to_treeir.py +++ b/src/gt4py/cartesian/gtc/dace/oir_to_treeir.py @@ -238,8 +238,16 @@ def visit_While(self, node: oir.While, ctx: tir.Context) -> None: def visit_AxisBound(self, node: common.AxisBound, axis_start: str, axis_end: str) -> str: if node.level == common.LevelMarker.START: + if axis_start == "0": + return f"({node.offset})" + if node.offset == 0: + return f"({axis_start})" return f"({axis_start}) + ({node.offset})" + if axis_end == "0": + return f"({node.offset})" + if node.offset == 0: + return f"({axis_end})" return f"({axis_end}) + ({node.offset})" def visit_RuntimeAxisBound(self, node: common.RuntimeAxisBound, **kwargs: Any) -> None: diff --git a/tests/cartesian_tests/unit_tests/test_gtc/dace/test_oir_to_treeir.py b/tests/cartesian_tests/unit_tests/test_gtc/dace/test_oir_to_treeir.py new file mode 100644 index 0000000000..c4d7f57eb8 --- /dev/null +++ b/tests/cartesian_tests/unit_tests/test_gtc/dace/test_oir_to_treeir.py @@ -0,0 +1,45 @@ +# GT4Py - GridTools Framework +# +# Copyright (c) 2014-2024, ETH Zurich +# All rights reserved. +# +# Please, refer to the LICENSE file in the root directory. +# SPDX-License-Identifier: BSD-3-Clause + +import pytest + +from gt4py.cartesian.gtc.dace import oir_to_treeir +from gt4py.cartesian.gtc import common +from gt4py.cartesian.stencil_builder import StencilBuilder +from gt4py.cartesian.gtscript import PARALLEL, Field, computation, interval + +# Because "dace tests" filter by `requires_dace`, we still need to add the marker. +# This global variable adds the marker to all test functions in this module. +pytestmark = pytest.mark.requires_dace + + +def ignore_me(field: Field[float]): # type: ignore + with computation(PARALLEL), interval(...): + field += 1 + + +@pytest.mark.parametrize( + "node,axis_start,axis_end,expected", + [ + (common.AxisBound.from_start(1), "2", "n/a", "(2) + (1)"), + (common.AxisBound.from_start(0), "2", "n/a", "(2)"), + (common.AxisBound.from_start(1), "0", "n/a", "(1)"), + (common.AxisBound.from_start(0), "0", "n/a", "(0)"), + (common.AxisBound.from_end(1), "n/a", "2", "(2) + (1)"), + (common.AxisBound.from_end(0), "n/a", "2", "(2)"), + (common.AxisBound.from_end(1), "n/a", "0", "(1)"), + (common.AxisBound.from_end(0), "n/a", "0", "(0)"), + ], +) +def test_visit_AxisBound( + node: common.AxisBound, axis_start: str, axis_end: str, expected: str +) -> None: + builder = StencilBuilder(ignore_me) + visitor = oir_to_treeir.OIRToTreeIR(builder) + + assert visitor.visit_AxisBound(node, axis_start, axis_end) == expected