[Mlir-commits] [mlir] [mlir][linalg][elementwise] Fold broadcast into new elementwise (PR #167626)

[llvmlistbot at llvm.org]

================
@@ -41,16 +41,54 @@ struct FoldTransposePattern : public OpRewritePattern<ElementwiseOp> {
       AffineMap map = op.getMatchingIndexingMap(operand);
       auto transposeOp = operand->get().getDefiningOp<TransposeOp>();
 
-      if (!map.isIdentity() || !transposeOp) {
+      if (!transposeOp) {
         // push in original operand and its map.
         newIns.push_back(operand->get());
         newMaps.push_back(map);
         continue;
       }
       newIns.push_back(transposeOp.getInput());
-      // push in transposeOp's inverse permutation map.
-      newMaps.push_back(transposeOp.getMatchingIndexingMap(
-          transposeOp.getDpsInputOperand(0)));
+      // push in composed affine map.
+      newMaps.push_back(
+          transposeOp.getMatchingIndexingMap(transposeOp.getDpsInputOperand(0))
+              .compose(map));
+      changed = true;
+    }
+    if (!changed)
+      return failure();
+    newMaps.push_back(op.getIndexingMapsArray().back());
+
+    rewriter.replaceOpWithNewOp<ElementwiseOp>(
+        op, newIns, op.getDpsInits()[0], op.getKindAttr(),
+        rewriter.getAffineMapArrayAttr(newMaps));
+    return success();
+  }
+};
+
+struct FoldBroadcastPattern : public OpRewritePattern<ElementwiseOp> {
----------------
someoneinjd wrote:

Thanks for the feedback! I've refactored the code to merge the patterns using variadic templates and a strategy pattern.

I introduced a generic `FoldIntoElementwisePattern<Ops...>` that delegates the specific folding logic to a `ElementwiseOpFolder` struct. This design offers two key advantages:

1. Extensibility: We can easily support new ops by adding them to the template list. Simple ops use the default `ElementwiseOpFolder`, while complex ops can define custom logic via template specialization.

2. Efficiency: Unlike separate patterns that require multiple scans, this approach performs a single traversal over the operands to check for all specified producer types simultaneously.

https://github.com/llvm/llvm-project/pull/167626