[Mlir-commits] [mlir] [mlir][linalg][elementwise] Fold transpose into new elementwise (PR #130207)

[Javed Absar]

================
@@ -4285,6 +4286,47 @@ Speculation::Speculatability ElementwiseOp::getSpeculatability() {
   return getGenericSpeculatabilityImpl(cast<LinalgOp>(getOperation()));
 }
 
+namespace {
+struct FoldTranspose : public OpRewritePattern<ElementwiseOp> {
+  using OpRewritePattern<ElementwiseOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ElementwiseOp op,
+                                PatternRewriter &rewriter) const override {
+    bool changed = false;
+    SmallVector<Value> newIns;
+    SmallVector<AffineMap> newMaps;
+    for (OpOperand *operand : op.getDpsInputOperands()) {
+      AffineMap map = op.getMatchingIndexingMap(operand);
+      auto transposeOp = operand->get().getDefiningOp<TransposeOp>();
+
+      if (!map.isIdentity() || !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)));
+      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();
+  }
+};
+} // namespace
+void ElementwiseOp::getCanonicalizationPatterns(RewritePatternSet &results,
----------------
javedabsar1 wrote:

Thanks @MaheshRavishankar / @rengolin . 
(a) I have redone the diff to use populate* approach as you suggested.
(b) But in addition to (a) I do still have to discuss this - When the transpose has a single-use, then folding it into elementwise should be a canonicalization step as that's a big point of new elementwise over say linalg.add. Other argument is that if the elementwise (from front-end) came with transpose folded in already we wouldn't consider unfolding it. So then it becomes that the IR is front-end dependent and we leave it as is. When transpose has not-single use I totally agree its a 'opt' and not a cononicalization.


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