Repository: GridTools/gt4py
Pull request #1410 (2024-01-09 08:57:55)

fix[cartesian]: DaCe array access in tasklet
Found some incompatible tasklet represention while upgrading to dace v0.15.1. Array access inside tasklet with partial index subset worked in v0.14.1, although not valid.

Here is an example of tasklet with partial index subset:
![Screenshot 2024-01-09 at 10 52 24](https://github.com/GridTools/gt4py/assets/25990065/5ad8663e-3013-48c5-9b0d-b8db6610cb14)

In case of a mapped tasklet, the memlet's closest source/endpoint is a map entry/exit node, like in this other example:
![Screenshot 2024-01-09 at 13 15 55](https://github.com/GridTools/gt4py/assets/25990065/9c0133c8-acf2-42c7-9ac7-517eb2d943d7)

The fix consists of modifying the memlets to pass the full array shape to such tasklet, and use all explicit indices inside the tasklet to access the array. This is the right representation in DaCe SDFG, as discussed with the DaCe developers.

The two tasklet cases shown above, with this fix, become:
<img width="1368" alt="Screenshot 2024-01-12 at 23 10 11" src="https://github.com/GridTools/gt4py/assets/25990065/f9970429-2b09-45c7-945a-b98da5fe1ee7">

and:
<img width="1537" alt="Screenshot 2024-01-12 at 23 15 30" src="https://github.com/GridTools/gt4py/assets/25990065/88e38b9a-b32a-4e22-9c77-c5d7a77ae379">

This fix is covered by the following test cases, which would otherwise fail with dace v0.15.1:
```
tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py::test_higher_dimensional_fields
tests/cartesian_tests/integration_tests/multi_feature_tests/test_suites.py::TestNon3DFields::test_generation
tests/cartesian_tests/integration_tests/multi_feature_tests/test_suites.py::TestTypedTemporary::test_generation
tests/cartesian_tests/integration_tests/multi_feature_tests/test_suites.py::TestMatrixAssignment::test_generation
tests/cartesian_tests/integration_tests/multi_feature_tests/test_suites.py::TestMatmul::test_generation
```

edopao: cscs-ci run
edopao: I have found one case for which this solution won't work. If you consider this test:
`tests/cartesian_tests/integration_tests/multi_feature_tests/test_code_generation.py::test_higher_dimensional_fields`

and this part of the SDFG generated by baseline:
![Screenshot 2024-01-09 at 14 44 45](https://github.com/GridTools/gt4py/assets/25990065/d47d7533-62a0-4922-9918-002e6eba6ad0)

it is not equivalent to the SDFG generated by the fix:
![Screenshot 2024-01-09 at 14 46 39](https://github.com/GridTools/gt4py/assets/25990065/dc2fc18a-1df1-4c85-ab7b-be9e915e605d)

The memlet from the transient `__mat_field_t` is copying 4 elements (subset [2,2]) to `mat_field`, while the original intention is to only write 2 elements.

I got a clue from the compiler warnings:
```
.../copy.h:153:40: warning: '__mat_field_t' may be used uninitialized [-Wmaybe-uninitialized]
  153 |                 *(vec<T, VECLEN> *)dst = *(vec<T, VECLEN> *)src;
      |                 ~~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~
.../computation.hpp: In function '_Z26__program_stencil_internalP15stencil_state_tPdS1_S1_iiiiiiiiiiiiiii._omp_fn.1':
.../computation.hpp:109:12: note: '__mat_field_t' declared here
  109 |     double __mat_field_t[4] DACE_ALIGN(64);
      |            ^~~~~~~~~~~~~
```
edopao: I have resolved the issue I reported yesterday, by avoiding use of transients. However, I have some dace compilation errors now which look like issues in dace itself.
edopao: cscs-ci run
edopao: All problems have been fixed. Coding done, ready for review.

Patch:
diff --git a/src/gt4py/cartesian/gtc/dace/expansion/daceir_builder.py b/src/gt4py/cartesian/gtc/dace/expansion/daceir_builder.py
index db276a48..48b129fa 100644
--- a/src/gt4py/cartesian/gtc/dace/expansion/daceir_builder.py
+++ b/src/gt4py/cartesian/gtc/dace/expansion/daceir_builder.py
@@ -30,6 +30,7 @@ from gt4py.cartesian.gtc.dace.utils import (
     compute_dcir_access_infos,
     flatten_list,
     get_tasklet_symbol,
+    make_dace_subset,
     union_inout_memlets,
     union_node_grid_subsets,
     untile_memlets,
@@ -458,6 +459,40 @@ class DaCeIRBuilder(eve.NodeTranslator):
             write_memlets=write_memlets,
         )
 
+        for memlet in [*read_memlets, *write_memlets]:
+            """
+            This loop handles the special case of a tasklet performing array access.
+            The memlet should pass the full array shape (no tiling) and
+            the tasklet expression for array access should use all explicit indexes.
+            """
+            array_ndims = len(global_ctx.arrays[memlet.field].shape)
+            field_decl = global_ctx.library_node.field_decls[memlet.field]
+            # calculate array subset on original memlet
+            memlet_subset = make_dace_subset(
+                global_ctx.library_node.access_infos[memlet.field],
+                memlet.access_info,
+                field_decl.data_dims,
+            )
+            # select index values for single-point grid access
+            memlet_data_index = [
+                dcir.Literal(value=str(dim_range[0]), dtype=common.DataType.INT32)
+                for dim_range, dim_size in zip(memlet_subset, memlet_subset.size())
+                if dim_size == 1
+            ]
+            if len(memlet_data_index) < array_ndims:
+                reshape_memlet = False
+                for access_node in dcir_node.walk_values().if_isinstance(dcir.IndexAccess):
+                    if access_node.data_index and access_node.name == memlet.connector:
+                        access_node.data_index = memlet_data_index + access_node.data_index
+                        assert len(access_node.data_index) == array_ndims
+                        reshape_memlet = True
+                if reshape_memlet:
+                    # ensure that memlet symbols used for array indexing are defined in context
+                    for sym in memlet.access_info.grid_subset.free_symbols:
+                        symbol_collector.add_symbol(sym)
+                    # set full shape on memlet
+                    memlet.access_info = global_ctx.library_node.access_infos[memlet.field]
+
         for item in reversed(expansion_items):
             iteration_ctx = iteration_ctx.pop()
             dcir_node = self._process_iteration_item(
diff --git a/src/gt4py/cartesian/gtc/daceir.py b/src/gt4py/cartesian/gtc/daceir.py
index 28ebc8cd..03663173 100644
--- a/src/gt4py/cartesian/gtc/daceir.py
+++ b/src/gt4py/cartesian/gtc/daceir.py
@@ -536,7 +536,7 @@ class GridSubset(eve.Node):
             else:
                 assert (
                     isinstance(interval2, (TileInterval, DomainInterval))
-                    and isinstance(interval1, IndexWithExtent)
+                    and isinstance(interval1, (IndexWithExtent, DomainInterval))
                 ) or (
                     isinstance(interval1, (TileInterval, DomainInterval))
                     and isinstance(interval2, IndexWithExtent)
@@ -573,7 +573,7 @@ class FieldAccessInfo(eve.Node):
     def apply_iteration(self, grid_subset: GridSubset):
         res_intervals = dict(self.grid_subset.intervals)
         for axis, field_interval in self.grid_subset.intervals.items():
-            if axis in grid_subset.intervals:
+            if axis in grid_subset.intervals and not isinstance(field_interval, DomainInterval):
                 grid_interval = grid_subset.intervals[axis]
                 assert isinstance(field_interval, IndexWithExtent)
                 extent = field_interval.extent
diff --git a/src/gt4py/next/iterator/transforms/power_unrolling.py b/src/gt4py/next/iterator/transforms/power_unrolling.py
new file mode 100644
index 00000000..ac71f274
--- /dev/null
+++ b/src/gt4py/next/iterator/transforms/power_unrolling.py
@@ -0,0 +1,84 @@
+# GT4Py - GridTools Framework
+#
+# Copyright (c) 2014-2023, ETH Zurich
+# All rights reserved.
+#
+# This file is part of the GT4Py project and the GridTools framework.
+# GT4Py 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 <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+import dataclasses
+import math
+
+from gt4py.eve import NodeTranslator
+from gt4py.next.iterator import ir
+from gt4py.next.iterator.ir_utils import ir_makers as im
+from gt4py.next.iterator.transforms.inline_lambdas import InlineLambdas
+
+
+def _is_power_call(
+    node: ir.FunCall,
+) -> bool:
+    """Match expressions of the form `power(base, integral_literal)`."""
+    return (
+        isinstance(node.fun, ir.SymRef)
+        and node.fun.id == "power"
+        and isinstance(node.args[1], ir.Literal)
+        and float(node.args[1].value) == int(node.args[1].value)
+        and node.args[1].value >= im.literal_from_value(0).value
+    )
+
+
+def _compute_integer_power_of_two(exp: int) -> int:
+    return math.floor(math.log2(exp))
+
+
+@dataclasses.dataclass
+class PowerUnrolling(NodeTranslator):
+    max_unroll: int
+
+    @classmethod
+    def apply(cls, node: ir.Node, max_unroll: int = 5) -> ir.Node:
+        return cls(max_unroll=max_unroll).visit(node)
+
+    def visit_FunCall(self, node: ir.FunCall):
+        new_node = self.generic_visit(node)
+
+        if _is_power_call(new_node):
+            assert len(new_node.args) == 2
+            # Check if unroll should be performed or if exponent is too large
+            base, exponent = new_node.args[0], int(new_node.args[1].value)
+            if 1 <= exponent <= self.max_unroll:
+                # Calculate and store powers of two of the base as long as they are smaller than the exponent.
+                # Do the same (using the stored values) with the remainder and multiply computed values.
+                pow_cur = _compute_integer_power_of_two(exponent)
+                pow_max = pow_cur
+                remainder = exponent
+
+                # Build target expression
+                ret = im.ref(f"power_{2 ** pow_max}")
+                remainder -= 2**pow_cur
+                while remainder > 0:
+                    pow_cur = _compute_integer_power_of_two(remainder)
+                    remainder -= 2**pow_cur
+
+                    ret = im.multiplies_(ret, f"power_{2 ** pow_cur}")
+
+                # Nest target expression to avoid multiple redundant evaluations
+                for i in range(pow_max, 0, -1):
+                    ret = im.let(
+                        f"power_{2 ** i}",
+                        im.multiplies_(f"power_{2**(i-1)}", f"power_{2**(i-1)}"),
+                    )(ret)
+                ret = im.let("power_1", base)(ret)
+
+                # Simplify expression in case of SymRef by resolving let statements
+                if isinstance(base, ir.SymRef):
+                    return InlineLambdas.apply(ret, opcount_preserving=True)
+                else:
+                    return ret
+        return new_node


Test patch:
diff --git a/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_power_unrolling.py b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_power_unrolling.py
new file mode 100644
index 00000000..ae23becb
--- /dev/null
+++ b/tests/next_tests/unit_tests/iterator_tests/transforms_tests/test_power_unrolling.py
@@ -0,0 +1,161 @@
+# GT4Py - GridTools Framework
+#
+# Copyright (c) 2014-2023, ETH Zurich
+# All rights reserved.
+#
+# This file is part of the GT4Py project and the GridTools framework.
+# GT4Py 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 <https://www.gnu.org/licenses/>.
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import pytest
+
+from gt4py.eve import SymbolRef
+from gt4py.next.iterator import ir
+from gt4py.next.iterator.ir_utils import ir_makers as im
+from gt4py.next.iterator.transforms.power_unrolling import PowerUnrolling
+
+
+def test_power_unrolling_zero():
+    pytest.xfail(
+        "Not implementeds we don't have an easy way to determine the type of the one literal (type inference is to expensive)."
+    )
+    testee = im.call("power")("x", 0)
+    expected = im.literal_from_value(1)
+
+    actual = PowerUnrolling.apply(testee)
+    assert actual == expected
+
+
+def test_power_unrolling_one():
+    testee = im.call("power")("x", 1)
+    expected = ir.SymRef(id=SymbolRef("x"))
+
+    actual = PowerUnrolling.apply(testee)
+    assert actual == expected
+
+
+def test_power_unrolling_two():
+    testee = im.call("power")("x", 2)
+    expected = im.multiplies_("x", "x")
+
+    actual = PowerUnrolling.apply(testee)
+    assert actual == expected
+
+
+def test_power_unrolling_two_x_plus_two():
+    testee = im.call("power")(im.plus("x", 2), 2)
+    expected = im.let("power_1", im.plus("x", 2))(
+        im.let("power_2", im.multiplies_("power_1", "power_1"))("power_2")
+    )
+
+    actual = PowerUnrolling.apply(testee)
+    assert actual == expected
+
+
+def test_power_unrolling_two_x_plus_one_times_three():
+    testee = im.call("power")(im.multiplies_(im.plus("x", 1), 3), 2)
+    expected = im.let("power_1", im.multiplies_(im.plus("x", 1), 3))(
+        im.let("power_2", im.multiplies_("power_1", "power_1"))("power_2")
+    )
+
+    actual = PowerUnrolling.apply(testee)
+    assert actual == expected
+
+
+def test_power_unrolling_three():
+    testee = im.call("power")("x", 3)
+    expected = im.multiplies_(im.multiplies_("x", "x"), "x")
+
+    actual = PowerUnrolling.apply(testee)
+    assert actual == expected
+
+
+def test_power_unrolling_four():
+    testee = im.call("power")("x", 4)
+    expected = im.let("power_2", im.multiplies_("x", "x"))(im.multiplies_("power_2", "power_2"))
+
+    actual = PowerUnrolling.apply(testee)
+    assert actual == expected
+
+
+def test_power_unrolling_five():
+    testee = im.call("power")("x", 5)
+    tmp2 = im.multiplies_("x", "x")
+    expected = im.multiplies_(im.multiplies_(tmp2, tmp2), "x")
+    expected = im.let("power_2", im.multiplies_("x", "x"))(
+        im.multiplies_(im.multiplies_("power_2", "power_2"), "x")
+    )
+
+    actual = PowerUnrolling.apply(testee)
+    assert actual == expected
+
+
+def test_power_unrolling_seven():
+    testee = im.call("power")("x", 7)
+    expected = im.call("power")("x", 7)
+
+    actual = PowerUnrolling.apply(testee, max_unroll=5)
+    assert actual == expected
+
+
+def test_power_unrolling_seven_unrolled():
+    testee = im.call("power")("x", 7)
+    expected = im.let("power_2", im.multiplies_("x", "x"))(
+        im.multiplies_(im.multiplies_(im.multiplies_("power_2", "power_2"), "power_2"), "x")
+    )
+
+    actual = PowerUnrolling.apply(testee, max_unroll=7)
+    assert actual == expected
+
+
+def test_power_unrolling_seven_x_plus_one_unrolled():
+    testee = im.call("power")(im.plus("x", 1), 7)
+    expected = im.let("power_1", im.plus("x", 1))(
+        im.let("power_2", im.multiplies_("power_1", "power_1"))(
+            im.let("power_4", im.multiplies_("power_2", "power_2"))(
+                im.multiplies_(im.multiplies_("power_4", "power_2"), "power_1")
+            )
+        )
+    )
+
+    actual = PowerUnrolling.apply(testee, max_unroll=7)
+    assert actual == expected
+
+
+def test_power_unrolling_eight():
+    testee = im.call("power")("x", 8)
+    expected = im.call("power")("x", 8)
+
+    actual = PowerUnrolling.apply(testee, max_unroll=5)
+    assert actual == expected
+
+
+def test_power_unrolling_eight_unrolled():
+    testee = im.call("power")("x", 8)
+    expected = im.let("power_2", im.multiplies_("x", "x"))(
+        im.let("power_4", im.multiplies_("power_2", "power_2"))(
+            im.multiplies_("power_4", "power_4")
+        )
+    )
+
+    actual = PowerUnrolling.apply(testee, max_unroll=8)
+    assert actual == expected
+
+
+def test_power_unrolling_eight_x_plus_one_unrolled():
+    testee = im.call("power")(im.plus("x", 1), 8)
+    expected = im.let("power_1", im.plus("x", 1))(
+        im.let("power_2", im.multiplies_("power_1", "power_1"))(
+            im.let("power_4", im.multiplies_("power_2", "power_2"))(
+                im.let("power_8", im.multiplies_("power_4", "power_4"))("power_8")
+            )
+        )
+    )
+
+    actual = PowerUnrolling.apply(testee, max_unroll=8)
+    assert actual == expected