transformations: New varith-fuse-repeated-operands pass

tests/filecheck/transforms/varith-fuse-repeated-operands.mlir

@@ -0,0 +1,43 @@
+// RUN: xdsl-opt --split-input-file -p varith-fuse-repeated-operands %s | filecheck %s
+
+func.func @test_addi() {
+    %a, %b, %c = "test.op"() : () -> (i32, i32, i32)
+    %1, %2, %3 = "test.op"() : () -> (i32, i32, i32)
+
+    %r = varith.add %a, %b, %a, %a, %b, %c : i32
+
+    "test.op"(%r) : (i32) -> ()
+
+    return
+
+    // CHECK-LABEL: @test_addi
+    // CHECK-NEXT:   %a, %b, %c = "test.op"() : () -> (i32, i32, i32)
+    // CHECK-NEXT:   %0, %1, %2 = "test.op"() : () -> (i32, i32, i32)
+    // CHECK-NEXT:   %3 = arith.constant 3 : i32
+    // CHECK-NEXT:   %4 = arith.constant 2 : i32
+    // CHECK-NEXT:   %5 = arith.muli %3, %a : i32
+    // CHECK-NEXT:   %6 = arith.muli %4, %b : i32
+    // CHECK-NEXT:   %r = varith.add %5, %6, %c : i32
+    // CHECK-NEXT:   "test.op"(%r) : (i32) -> ()
+}
+
+func.func @test_addf() {
+    %a, %b, %c = "test.op"() : () -> (f32, f32, f32)
+    %1, %2, %3 = "test.op"() : () -> (f32, f32, f32)
+
+    %r = varith.add %a, %b, %a, %a, %b, %c : f32
+
+    "test.op"(%r) : (f32) -> ()
+
+    return
+
+    // CHECK-LABEL: @test_addf
+    // CHECK-NEXT:   %a, %b, %c = "test.op"() : () -> (f32, f32, f32)
+    // CHECK-NEXT:   %0, %1, %2 = "test.op"() : () -> (f32, f32, f32)
+    // CHECK-NEXT:   %3 = arith.constant 3.000000e+00 : f32
+    // CHECK-NEXT:   %4 = arith.constant 2.000000e+00 : f32
+    // CHECK-NEXT:   %5 = arith.mulf %3, %a : f32
+    // CHECK-NEXT:   %6 = arith.mulf %4, %b : f32
+    // CHECK-NEXT:   %r = varith.add %5, %6, %c : f32
+    // CHECK-NEXT:   "test.op"(%r) : (f32) -> ()
+}

xdsl/tools/command_line_tool.py

@@ -440,6 +440,11 @@ def get_stencil_shape_minimize():
 
         return stencil_shape_minimize.StencilShapeMinimize
 
+    def get_varith_fuse_repeated_operands():
+        from xdsl.transforms import varith_transformations
+
+        return varith_transformations.VarithFuseRepeatedOperandsPass
+
     return {
         "arith-add-fastmath": get_arith_add_fastmath,
         "loop-hoist-memref": get_loop_hoist_memref,
@@ -525,6 +530,7 @@ def get_stencil_shape_minimize():
         "stencil-shape-minimize": get_stencil_shape_minimize,
         "test-lower-linalg-to-snitch": get_test_lower_linalg_to_snitch,
         "eqsat-create-eclasses": get_eqsat_create_eclasses,
+        "varith-fuse-repeated-operands": get_varith_fuse_repeated_operands,
     }
 
 

xdsl/transforms/varith_transformations.py

@@ -1,3 +1,5 @@
+import collections
+from dataclasses import dataclass
 from typing import Literal, cast
 
 from xdsl.context import MLContext
@@ -11,6 +13,7 @@
     RewritePattern,
     op_type_rewrite_pattern,
 )
+from xdsl.rewriter import InsertPoint
 from xdsl.utils.hints import isa
 
 # map the arith operation to the right varith op:
@@ -179,6 +182,57 @@ def is_integer_like_type(t: Attribute) -> bool:
     return False
 
 
+@dataclass
+class FuseRepeatedAddArgsPattern(RewritePattern):
+    """
+    Prefer `operand * count(operand)` over repeated addition of `operand`.
+
+    The minimum count to trigger this rewrite can be specified in `min_reps`.
+    """
+
+    min_reps: int
+    """Minimum repetitions of operand to trigger fusion."""
+
+    @op_type_rewrite_pattern
+    def match_and_rewrite(self, op: varith.VarithAddOp, rewriter: PatternRewriter, /):
+        elem_t = op.res.type
+        if isinstance(elem_t, builtin.ContainerType):
+            elem_t = cast(builtin.ContainerType[Attribute], elem_t).get_element_type()
+
+        assert isinstance(
+            elem_t, builtin.IntegerType | builtin.IndexType | builtin.AnyFloat
+        )
+
+        consts: list[arith.Constant] = []
+        fusions: list[Operation] = []
+        new_args: list[Operation | SSAValue] = []
+        for arg, count in collections.Counter(op.args).items():
+            if count >= self.min_reps:
+                c, f = self.fuse(arg, count, elem_t)
+                consts.append(c)
+                fusions.append(f)
+                new_args.append(f)
+            else:
+                new_args.append(arg)
+        if fusions:
+            rewriter.insert_op([*consts, *fusions], InsertPoint.before(op))
+            rewriter.replace_matched_op(varith.VarithAddOp(*new_args))
+
+    @staticmethod
+    def fuse(
+        arg: SSAValue,
+        count: int,
+        t: builtin.IntegerType | builtin.IndexType | builtin.AnyFloat,
+    ):
+        if isinstance(t, builtin.IntegerType | builtin.IndexType):
+            c = arith.Constant(builtin.IntegerAttr(count, t))
+            f = arith.Muli
+        else:
+            c = arith.Constant(builtin.FloatAttr(count, t))
+            f = arith.Mulf
+        return c, f(c, arg)
+
+
 class ConvertArithToVarithPass(ModulePass):
     """
     Convert chains of arith.{add|mul}{i,f} operations into a single long variadic add or mul operation.
@@ -215,3 +269,20 @@ def apply(self, ctx: MLContext, op: builtin.ModuleOp) -> None:
             VarithToArithPattern(),
             apply_recursively=False,
         ).rewrite_op(op)
+
+
+class VarithFuseRepeatedOperandsPass(ModulePass):
+    """
+    Fuses several occurrences of the same operand into one.
+    """
+
+    name = "varith-fuse-repeated-operands"
+
+    min_reps: int = 2
+    """The minimum number of times an operand needs to be repeated before being fused."""
+
+    def apply(self, ctx: MLContext, op: builtin.ModuleOp) -> None:
+        PatternRewriteWalker(
+            FuseRepeatedAddArgsPattern(self.min_reps),
+            apply_recursively=False,
+        ).rewrite_op(op)