Repository: GridTools/gt4py
Pull request #1409 (2024-01-05 10:04:25)

feature[next]: Add power unrolling functionality and respective unit tests.



## PR Title: Add power unrolling functionality and respective unit tests.

    <type>:
        - powers of SymRefs and FunCalls until order 5 are unrolled per default
        - powers with higher exponent are still computed by the built-in functions.
        - test: add respective tests

    <scope>: next 




Patch:
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