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

[Julian Oppermann]

================
@@ -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> {
----------------
jopperm wrote:

I think the logic in the transpose and broadcast pattern is identical -- could you make the op class a template parameter, or alternatively handle both in the same pattern?

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