feat[next]: Extend and refactor constant folding

src/gt4py/next/ffront/fbuiltins.py

@@ -10,6 +10,7 @@
 import functools
 import inspect
 import math
+import operator
 from builtins import bool, float, int, tuple  # noqa: A004 shadowing a Python built-in
 from typing import Any, Callable, Final, Generic, ParamSpec, Tuple, TypeAlias, TypeVar, Union, cast
 
@@ -203,7 +204,7 @@ def astype(
     return core_defs.dtype(type_).scalar_type(value)
 
 
-_UNARY_MATH_NUMBER_BUILTIN_IMPL: Final = {"abs": abs}
+_UNARY_MATH_NUMBER_BUILTIN_IMPL: Final = {"abs": abs, "neg": operator.neg}
 UNARY_MATH_NUMBER_BUILTIN_NAMES: Final = [*_UNARY_MATH_NUMBER_BUILTIN_IMPL.keys()]
 
 _UNARY_MATH_FP_BUILTIN_IMPL: Final = {

src/gt4py/next/ffront/foast_to_gtir.py

@@ -241,12 +241,12 @@ def visit_UnaryOp(self, node: foast.UnaryOp, **kwargs: Any) -> itir.Expr:
             if dtype.kind != ts.ScalarKind.BOOL:
                 raise NotImplementedError(f"'{node.op}' is only supported on 'bool' arguments.")
             return self._lower_and_map("not_", node.operand)
-
-        return self._lower_and_map(
-            node.op.value,
-            foast.Constant(value="0", type=dtype, location=node.location),
-            node.operand,
-        )
+        if node.op in [dialect_ast_enums.UnaryOperator.USUB]:
+            return self._lower_and_map("neg", node.operand)
+        if node.op in [dialect_ast_enums.UnaryOperator.UADD]:
+            return self.visit(node.operand)
+        else:
+            raise NotImplementedError(f"Unary operator '{node.op}' is not supported.")
 
     def visit_BinOp(self, node: foast.BinOp, **kwargs: Any) -> itir.FunCall:
         return self._lower_and_map(node.op.value, node.left, node.right)

src/gt4py/next/iterator/builtins.py

@@ -292,6 +292,11 @@ def trunc(*args):
     raise BackendNotSelectedError()
 
 
+@builtin_dispatch
+def neg(*args):
+    raise BackendNotSelectedError()
+
+
 @builtin_dispatch
 def isfinite(*args):
     raise BackendNotSelectedError()
@@ -397,7 +402,7 @@ def bool(*args):  # noqa: A001 [builtin-variable-shadowing]
     raise BackendNotSelectedError()
 
 
-UNARY_MATH_NUMBER_BUILTINS = {"abs"}
+UNARY_MATH_NUMBER_BUILTINS = {"abs", "neg"}
 UNARY_LOGICAL_BUILTINS = {"not_"}
 UNARY_MATH_FP_BUILTINS = {
     "sin",
@@ -420,6 +425,7 @@ def bool(*args):  # noqa: A001 [builtin-variable-shadowing]
     "floor",
     "ceil",
     "trunc",
+    "neg",
 }
 UNARY_MATH_FP_PREDICATE_BUILTINS = {"isfinite", "isinf", "isnan"}
 BINARY_MATH_NUMBER_BUILTINS = {

src/gt4py/next/iterator/embedded.py

@@ -392,6 +392,13 @@ def not_(a):
     return not a
 
 
+@builtins.neg.register(EMBEDDED)
+def neg(a):
+    if isinstance(a, Column):
+        return np.negative(a)
+    return np.negative(a)
+
+
 @builtins.gamma.register(EMBEDDED)
 def gamma(a):
     gamma_ = np.vectorize(math.gamma)
@@ -538,10 +545,11 @@ def promote_scalars(val: CompositeOfScalarOrField):
         "and_": operator.and_,
         "or_": operator.or_,
         "xor_": operator.xor,
+        "neg": operator.neg,
     }
     decorator = getattr(builtins, math_builtin_name).register(EMBEDDED)
     impl: Callable
-    if math_builtin_name in ["gamma", "not_"]:
+    if math_builtin_name in ["gamma", "not_", "neg"]:
         continue  # treated explicitly
     elif math_builtin_name in python_builtins:
         # TODO: Should potentially use numpy fixed size types to be consistent

src/gt4py/next/iterator/transforms/collapse_tuple.py

@@ -14,7 +14,6 @@
 import operator
 from typing import Optional
 
-from gt4py import eve
 from gt4py.eve import utils as eve_utils
 from gt4py.next import common
 from gt4py.next.iterator import ir
@@ -23,6 +22,7 @@
     ir_makers as im,
     misc as ir_misc,
 )
+from gt4py.next.iterator.transforms.fixed_point_transform import FixedPointTransform
 from gt4py.next.iterator.transforms.inline_lambdas import InlineLambdas, inline_lambda
 from gt4py.next.iterator.type_system import inference as itir_type_inference
 from gt4py.next.type_system import type_info, type_specifications as ts
@@ -87,7 +87,7 @@ def _is_trivial_or_tuple_thereof_expr(node: ir.Node) -> bool:
 #  reads a little convoluted and is also different to how we write other transformations. We
 #  should revisit the pattern here and try to find a more general mechanism.
 @dataclasses.dataclass(frozen=True)
-class CollapseTuple(eve.PreserveLocationVisitor, eve.NodeTranslator):
+class CollapseTuple(FixedPointTransform):
     """
     Simplifies `make_tuple`, `tuple_get` calls.
 
@@ -217,32 +217,6 @@ def visit_FunCall(self, node: ir.FunCall, **kwargs) -> ir.Node:
         node = self.generic_visit(node, **kwargs)
         return self.fp_transform(node, **kwargs)
 
-    def fp_transform(self, node: ir.Node, **kwargs) -> ir.Node:
-        while True:
-            new_node = self.transform(node, **kwargs)
-            if new_node is None:
-                break
-            assert new_node != node
-            node = new_node
-        return node
-
-    def transform(self, node: ir.Node, **kwargs) -> Optional[ir.Node]:
-        if not isinstance(node, ir.FunCall):
-            return None
-
-        for transformation in self.Flag:
-            if self.flags & transformation:
-                assert isinstance(transformation.name, str)
-                method = getattr(self, f"transform_{transformation.name.lower()}")
-                result = method(node, **kwargs)
-                if result is not None:
-                    assert (
-                        result is not node
-                    )  # transformation should have returned None, since nothing changed
-                    itir_type_inference.reinfer(result)
-                    return result
-        return None
-
     def transform_collapse_make_tuple_tuple_get(
         self, node: ir.FunCall, **kwargs
     ) -> Optional[ir.Node]: