transformations: (canonicalize) Arith const reassociation

tests/filecheck/dialects/arith/canonicalize.mlir

@@ -34,3 +34,31 @@ func.func @test_const_const() {
     // CHECK-NEXT:   %3 = arith.constant 9.542894e-01 : f32
     // CHECK-NEXT:   "test.op"(%0, %1, %2, %3) : (f32, f32, f32, f32) -> ()
 }
+
+func.func @test_const_var_const() {
+    %0, %1 = "test.op"() : () -> (f32, f32)
+    %a = arith.constant 2.9979 : f32
+    %b = arith.constant 3.1415 : f32
+    %c = arith.constant 4.1415 : f32
+    %d = arith.constant 5.1415 : f32
+
+    %2 = arith.mulf %0, %a : f32
+    %3 = arith.mulf %2, %b : f32
+
+    %4 = arith.mulf %0, %c fastmath<reassoc> : f32
+    %5 = arith.mulf %4, %d fastmath<fast> : f32
+
+    "test.op"(%3, %5) : (f32, f32) -> ()
+
+    return
+
+    // CHECK-LABEL: @test_const_var_const
+    // CHECK-NEXT:   %0, %1 = "test.op"() : () -> (f32, f32)
+    // CHECK-NEXT:   %a = arith.constant 2.997900e+00 : f32
+    // CHECK-NEXT:   %b = arith.constant 3.141500e+00 : f32
+    // CHECK-NEXT:   %2 = arith.mulf %0, %a : f32
+    // CHECK-NEXT:   %3 = arith.mulf %2, %b : f32
+    // CHECK-NEXT:   %4 = arith.constant 2.129352e+01 : f32
+    // CHECK-NEXT:   %5 = arith.mulf %4, %0 fastmath<fast> : f32
+    // CHECK-NEXT:   "test.op"(%3, %5) : (f32, f32) -> ()
+}

xdsl/dialects/arith.py

@@ -260,6 +260,19 @@ def get_canonicalization_patterns(cls) -> tuple[RewritePattern, ...]:
         return (FoldConstConstOp(),)
 
 
+class FloatingPointLikeBinaryOpHasFastReassociativeCanonicalizationPatternsTrait(
+    HasCanonicalizationPatternsTrait
+):
+    @classmethod
+    def get_canonicalization_patterns(cls) -> tuple[RewritePattern, ...]:
+        from xdsl.transforms.canonicalization_patterns.arith import (
+            FoldConstConstOp,
+            FoldConstsByReassociation,
+        )
+
+        return FoldConstsByReassociation(), FoldConstConstOp()
+
+
 class FloatingPointLikeBinaryOperation(IRDLOperation, abc.ABC):
     """A generic base class for arith's binary operations on floats."""
 
@@ -913,7 +926,10 @@ class Addf(FloatingPointLikeBinaryOperation):
     name = "arith.addf"
 
     traits = frozenset(
-        [Pure(), FloatingPointLikeBinaryOpHasCanonicalizationPatternsTrait()]
+        [
+            Pure(),
+            FloatingPointLikeBinaryOpHasFastReassociativeCanonicalizationPatternsTrait(),
+        ]
     )
 
 
@@ -931,7 +947,10 @@ class Mulf(FloatingPointLikeBinaryOperation):
     name = "arith.mulf"
 
     traits = frozenset(
-        [Pure(), FloatingPointLikeBinaryOpHasCanonicalizationPatternsTrait()]
+        [
+            Pure(),
+            FloatingPointLikeBinaryOpHasFastReassociativeCanonicalizationPatternsTrait(),
+        ]
     )
 
 

xdsl/transforms/canonicalization_patterns/arith.py

@@ -19,7 +19,34 @@ def match_and_rewrite(self, op: arith.Addi, rewriter: PatternRewriter) -> None:
             rewriter.replace_matched_op([], [op.rhs])
 
 
+def _fold_const_operation(
+    op_t: type[arith.FloatingPointLikeBinaryOperation],
+    lhs: builtin.AnyFloatAttr,
+    rhs: builtin.AnyFloatAttr,
+) -> arith.Constant | None:
+    match op_t:
+        case arith.Addf:
+            val = lhs.value.data + rhs.value.data
+        case arith.Subf:
+            val = lhs.value.data - rhs.value.data
+        case arith.Mulf:
+            val = lhs.value.data * rhs.value.data
+        case arith.Divf:
+            if rhs.value.data == 0.0:
+                # this mirrors what mlir does
+                val = float("inf")
+            else:
+                val = lhs.value.data / rhs.value.data
+        case _:
+            return
+    return arith.Constant(builtin.FloatAttr(val, lhs.type))
+
+
 class FoldConstConstOp(RewritePattern):
+    """
+    Folds a floating point binary op whose operands are both `arith.constant`s.
+    """
+
     @op_type_rewrite_pattern
     def match_and_rewrite(
         self, op: arith.FloatingPointLikeBinaryOperation, rewriter: PatternRewriter, /
@@ -29,20 +56,49 @@ def match_and_rewrite(
             and isinstance(op.rhs.owner, arith.Constant)
             and isa(l := op.lhs.owner.value, builtin.AnyFloatAttr)
             and isa(r := op.rhs.owner.value, builtin.AnyFloatAttr)
+            and (cnst := _fold_const_operation(type(op), l, r))
+        ):
+            rewriter.replace_matched_op(cnst)
+
+
+class FoldConstsByReassociation(RewritePattern):
+    """
+    Rewrites a chain of
+        `(const1 <op> var) <op> const2`
+    as
+        `folded_const <op> val`
+
+    The op must be associative and have the `fastmath<reassoc>` flag set.
+    """
+
+    @op_type_rewrite_pattern
+    def match_and_rewrite(
+        self, op: arith.Addf | arith.Mulf, rewriter: PatternRewriter, /
+    ):
+        if isinstance(op.lhs.owner, arith.Constant):
+            const1, val = op.lhs.owner, op.rhs
+        else:
+            const1, val = op.rhs.owner, op.lhs
+
+        if (
+            not isinstance(const1, arith.Constant)
+            or len(op.result.uses) != 1
+            or not isinstance(u := list(op.result.uses)[0].operation, type(op))
+            or not isinstance(
+                const2 := u.lhs.owner if u.rhs == op.result else u.rhs.owner,
+                arith.Constant,
+            )
+            or op.fastmath is None
+            or u.fastmath is None
+            or arith.FastMathFlag.REASSOC not in op.fastmath.flags
+            or arith.FastMathFlag.REASSOC not in u.fastmath.flags
+            or not isa(c1 := const1.value, builtin.AnyFloatAttr)
+            or not isa(c2 := const2.value, builtin.AnyFloatAttr)
         ):
-            match type(op):
-                case arith.Addf:
-                    val = l.value.data + r.value.data
-                case arith.Subf:
-                    val = l.value.data - r.value.data
-                case arith.Mulf:
-                    val = l.value.data * r.value.data
-                case arith.Divf:
-                    if r.value.data == 0.0:
-                        # this mirrors what mlir does
-                        val = float("inf")
-                    else:
-                        val = l.value.data / r.value.data
-                case _:
-                    return
-            rewriter.replace_matched_op(arith.Constant(builtin.FloatAttr(val, l.type)))
+            return
+
+        if cnsts := _fold_const_operation(type(op), c1, c2):
+            flags = arith.FastMathFlagsAttr(list(op.fastmath.flags | u.fastmath.flags))
+            rebuild = type(op)(cnsts, val, flags)
+            rewriter.replace_matched_op([cnsts, rebuild])
+            rewriter.replace_op(u, [], [rebuild.result])