transformations: (canonicalize) Arith const reassociation

tests/filecheck/dialects/arith/canonicalize.mlir

@@ -15,3 +15,50 @@
 // CHECK-NEXT:   %addf = arith.addf %lhsf32, %rhsf32 : f32
 // CHECK-NEXT:   %addf_vector = arith.addf %lhsvec, %rhsvec : vector<4xf32>
 // CHECK-NEXT:   "test.op"(%addf, %addf_vector) : (f32, vector<4xf32>) -> ()
+
+func.func @test_const_const() {
+    %a = arith.constant 2.9979 : f32
+    %b = arith.constant 3.1415 : f32
+    %1 = arith.addf %a, %b  : f32
+    %2 = arith.subf %a, %b  : f32
+    %3 = arith.mulf %a, %b  : f32
+    %4 = arith.divf %a, %b  : f32
+    "test.op"(%1, %2, %3, %4) : (f32, f32, f32, f32) -> ()
+
+    return
+
+    // CHECK-LABEL: @test_const_const
+    // CHECK-NEXT:   %0 = arith.constant 6.139400e+00 : f32
+    // CHECK-NEXT:   %1 = arith.constant -1.436000e-01 : f32
+    // CHECK-NEXT:   %2 = arith.constant 9.417903e+00 : f32
+    // 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

@@ -250,6 +250,29 @@ def __init__(
         )
 
 
+class FloatingPointLikeBinaryOpHasCanonicalizationPatternsTrait(
+    HasCanonicalizationPatternsTrait
+):
+    @classmethod
+    def get_canonicalization_patterns(cls) -> tuple[RewritePattern, ...]:
+        from xdsl.transforms.canonicalization_patterns.arith import FoldConstConstOp
+
+        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."""
 
@@ -902,28 +925,42 @@ def print(self, printer: Printer):
 class Addf(FloatingPointLikeBinaryOperation):
     name = "arith.addf"
 
-    traits = frozenset([Pure()])
+    traits = frozenset(
+        [
+            Pure(),
+            FloatingPointLikeBinaryOpHasFastReassociativeCanonicalizationPatternsTrait(),
+        ]
+    )
 
 
 @irdl_op_definition
 class Subf(FloatingPointLikeBinaryOperation):
     name = "arith.subf"
 
-    traits = frozenset([Pure()])
+    traits = frozenset(
+        [Pure(), FloatingPointLikeBinaryOpHasCanonicalizationPatternsTrait()]
+    )
 
 
 @irdl_op_definition
 class Mulf(FloatingPointLikeBinaryOperation):
     name = "arith.mulf"
 
-    traits = frozenset([Pure()])
+    traits = frozenset(
+        [
+            Pure(),
+            FloatingPointLikeBinaryOpHasFastReassociativeCanonicalizationPatternsTrait(),
+        ]
+    )
 
 
 @irdl_op_definition
 class Divf(FloatingPointLikeBinaryOperation):
     name = "arith.divf"
 
-    traits = frozenset([Pure()])
+    traits = frozenset(
+        [Pure(), FloatingPointLikeBinaryOpHasCanonicalizationPatternsTrait()]
+    )
 
 
 @irdl_op_definition

xdsl/transforms/canonicalization_patterns/arith.py

@@ -1,10 +1,11 @@
-from xdsl.dialects import arith
+from xdsl.dialects import arith, builtin
 from xdsl.dialects.builtin import IntegerAttr
 from xdsl.pattern_rewriter import (
     PatternRewriter,
     RewritePattern,
     op_type_rewrite_pattern,
 )
+from xdsl.utils.hints import isa
 
 
 class AddImmediateZero(RewritePattern):
@@ -16,3 +17,84 @@ def match_and_rewrite(self, op: arith.Addi, rewriter: PatternRewriter) -> None:
             and value.value.data == 0
         ):
             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:
+            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, /
+    ):
+        if (
+            isinstance(op.lhs.owner, arith.Constant)
+            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)
+        ):
+            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])