[Mlir-commits] [mlir] [mlir] [linalg] Fold broadcast/transpose into linalg.generic (PR #212415)

[Renato Golin]

================
@@ -49,34 +50,55 @@ struct ElementwiseOpFolder {
 };
 
 template <typename... ProducerOps>
-struct FoldIntoElementwisePattern : public OpRewritePattern<ElementwiseOp> {
-  using OpRewritePattern<ElementwiseOp>::OpRewritePattern;
+static bool foldInputOperands(LinalgOp op, SmallVector<Value> &newIns,
+                              SmallVector<AffineMap> &newMaps) {
+  bool changed = false;
+  for (OpOperand *operand : op.getDpsInputOperands()) {
+    AffineMap consumerMap = op.getMatchingIndexingMap(operand);
+    const bool folded = (ElementwiseOpFolder<ProducerOps>::fold(
+                             operand, consumerMap, newIns, newMaps) ||
+                         ...);
+    if (folded) {
+      changed = true;
+    } else {
+      newIns.push_back(operand->get());
+      newMaps.push_back(consumerMap);
+    }
+  }
+  return changed;
+}
+
+template <typename... ProducerOps>
+struct FoldIntoElementwisePattern : public OpInterfaceRewritePattern<LinalgOp> {
+  using OpInterfaceRewritePattern<LinalgOp>::OpInterfaceRewritePattern;
 
-  LogicalResult matchAndRewrite(ElementwiseOp op,
+  LogicalResult matchAndRewrite(LinalgOp op,
                                 PatternRewriter &rewriter) const override {
-    bool changed = false;
+    if (!isa<GenericOp, ElementwiseOp>(op.getOperation()) || !isElementwise(op))
+      return failure();
+
     SmallVector<Value> newIns;
     SmallVector<AffineMap> newMaps;
-    for (OpOperand *operand : op.getDpsInputOperands()) {
-      AffineMap consumerMap = op.getMatchingIndexingMap(operand);
-      const bool folded = (ElementwiseOpFolder<ProducerOps>::fold(
-                               operand, consumerMap, newIns, newMaps) ||
-                           ...);
-      if (folded) {
-        changed = true;
-      } else {
-        // push in original operand and its map.
-        newIns.push_back(operand->get());
-        newMaps.push_back(consumerMap);
-      }
-    }
-    if (!changed)
+    if (!foldInputOperands<ProducerOps...>(op, newIns, newMaps))
----------------
rengolin wrote:

you don't need to outline this now

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