Repository: GridTools/gt4py (python)

Issue:
ProgramWithBoundArgs: pos_or_kw_args need to go to kwargs sometimes
in decorator.py:455, we need to do something like

```
        full_args = [*args]
        full_kwargs = {**kwargs}
        for index, param in enumerate(self.past_node.params):
            if param.id in self.bound_args.keys():
                if index < len(full_args):
                    full_args.insert(index, self.bound_args[param.id])
                else:
                    full_kwargs[param.id] = self.bound_args[param.id]

        return super()._process_args(tuple(full_args), full_kwargs)
```
in case we have a bound_arg which need to be provided as kwarg, because all previous pos_or_kw_args are passed as kwargs.

Test case would be something like
```python
@program
def foo(a, b, c):
    ...

bound = foo.with_bound_args(b = True)

bound(a=1,c=2)
```
currently we try to call foo as `foo(True, a=1,c=2)` which messes up the order.

Pull request:
bug[next]: Bound args kwargs edit
Fixes #1402 

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/ffront/decorator.py b/src/gt4py/next/ffront/decorator.py
index 147059b1..05cbe1c8 100644
--- a/src/gt4py/next/ffront/decorator.py
+++ b/src/gt4py/next/ffront/decorator.py
@@ -453,27 +453,32 @@ class ProgramWithBoundArgs(Program):
             ) from err
 
         full_args = [*args]
+        full_kwargs = {**kwargs}
         for index, param in enumerate(self.past_node.params):
             if param.id in self.bound_args.keys():
-                full_args.insert(index, self.bound_args[param.id])
+                if index < len(full_args):
+                    full_args.insert(index, self.bound_args[param.id])
+                else:
+                    full_kwargs[str(param.id)] = self.bound_args[param.id]
 
-        return super()._process_args(tuple(full_args), kwargs)
+        return super()._process_args(tuple(full_args), full_kwargs)
 
     @functools.cached_property
     def itir(self):
         new_itir = super().itir
         for new_clos in new_itir.closures:
-            for key in self.bound_args.keys():
+            new_args = [ref(inp.id) for inp in new_clos.inputs]
+            for key, value in self.bound_args.items():
                 index = next(
                     index
                     for index, closure_input in enumerate(new_clos.inputs)
                     if closure_input.id == key
                 )
+                new_args[new_args.index(new_clos.inputs[index])] = promote_to_const_iterator(
+                    literal_from_value(value)
+                )
                 new_clos.inputs.pop(index)
-            new_args = [ref(inp.id) for inp in new_clos.inputs]
             params = [sym(inp.id) for inp in new_clos.inputs]
-            for value in self.bound_args.values():
-                new_args.append(promote_to_const_iterator(literal_from_value(value)))
             expr = itir.FunCall(
                 fun=new_clos.stencil,
                 args=new_args,
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/integration_tests/feature_tests/ffront_tests/test_bound_args.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_bound_args.py
new file mode 100644
index 00000000..0de953d8
--- /dev/null
+++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_bound_args.py
@@ -0,0 +1,64 @@
+# -*- coding: utf-8 -*-
+# 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 numpy as np
+
+import gt4py.next as gtx
+from gt4py.next import int32
+
+from next_tests.integration_tests import cases
+from next_tests.integration_tests.cases import cartesian_case
+from next_tests.integration_tests.feature_tests.ffront_tests.ffront_test_utils import (
+    fieldview_backend,
+    reduction_setup,
+)
+
+
+def test_with_bound_args(cartesian_case):
+    @gtx.field_operator
+    def fieldop_bound_args(a: cases.IField, scalar: int32, condition: bool) -> cases.IField:
+        if not condition:
+            scalar = 0
+        return a + scalar
+
+    @gtx.program
+    def program_bound_args(a: cases.IField, scalar: int32, condition: bool, out: cases.IField):
+        fieldop_bound_args(a, scalar, condition, out=out)
+
+    a = cases.allocate(cartesian_case, program_bound_args, "a")()
+    scalar = int32(1)
+    ref = a + scalar
+    out = cases.allocate(cartesian_case, program_bound_args, "out")()
+
+    prog_bounds = program_bound_args.with_bound_args(scalar=scalar, condition=True)
+    cases.verify(cartesian_case, prog_bounds, a, out, inout=out, ref=ref)
+
+
+def test_with_bound_args_order_args(cartesian_case):
+    @gtx.field_operator
+    def fieldop_args(a: cases.IField, condition: bool, scalar: int32) -> cases.IField:
+        scalar = 0 if not condition else scalar
+        return a + scalar
+
+    @gtx.program(backend=cartesian_case.backend)
+    def program_args(a: cases.IField, condition: bool, scalar: int32, out: cases.IField):
+        fieldop_args(a, condition, scalar, out=out)
+
+    a = cases.allocate(cartesian_case, program_args, "a")()
+    out = cases.allocate(cartesian_case, program_args, "out")()
+
+    prog_bounds = program_args.with_bound_args(condition=True)
+    prog_bounds(a=a, scalar=int32(1), out=out, offset_provider={})
+    np.allclose(out.asnumpy(), a.asnumpy() + int32(1))
diff --git a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py
index a0893162..70c79d7b 100644
--- a/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py
+++ b/tests/next_tests/integration_tests/feature_tests/ffront_tests/test_execution.py
@@ -898,26 +898,6 @@ def test_docstring(cartesian_case):
     cases.verify(cartesian_case, test_docstring, a, inout=a, ref=a)
 
 
-def test_with_bound_args(cartesian_case):
-    @gtx.field_operator
-    def fieldop_bound_args(a: cases.IField, scalar: int32, condition: bool) -> cases.IField:
-        if not condition:
-            scalar = 0
-        return a + a + scalar
-
-    @gtx.program
-    def program_bound_args(a: cases.IField, scalar: int32, condition: bool, out: cases.IField):
-        fieldop_bound_args(a, scalar, condition, out=out)
-
-    a = cases.allocate(cartesian_case, program_bound_args, "a")()
-    scalar = int32(1)
-    ref = a + a + 1
-    out = cases.allocate(cartesian_case, program_bound_args, "out")()
-
-    prog_bounds = program_bound_args.with_bound_args(scalar=scalar, condition=True)
-    cases.verify(cartesian_case, prog_bounds, a, out, inout=out, ref=ref)
-
-
 def test_domain(cartesian_case):
     @gtx.field_operator
     def fieldop_domain(a: cases.IField) -> cases.IField:
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