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

[Chuanqi Xu]

================
@@ -81,6 +85,46 @@ struct FoldIntoElementwisePattern : public OpRewritePattern<ElementwiseOp> {
   }
 };
 
+template <typename... ProducerOps>
+struct FoldIntoGenericPattern : public OpRewritePattern<GenericOp> {
+  using OpRewritePattern<GenericOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(GenericOp op,
+                                PatternRewriter &rewriter) const override {
+    // Restrict this pattern to elementwise-like generic ops.
+    // It may be safe to do so reduction dimensions in some cases. But we try
+    // to focus on simple cases here.
+    if (!op.isAllParallelLoops())
+      return failure();
+
+    SmallVector<Value> newIns;
+    SmallVector<AffineMap> newMaps;
+    if (!foldInputOperands<GenericOp, ProducerOps...>(op, newIns, newMaps))
+      return failure();
+
+    // Keep all output operands and their maps unchanged. The body is cloned
+    // so that the block arguments continue to correspond to the new operand
+    // list.
+    SmallVector<AffineMap> allMaps = op.getIndexingMapsArray();
+    newMaps.append(allMaps.begin() + op.getNumDpsInputs(), allMaps.end());
----------------
ChuanqiXu9 wrote:

`foldInputOperands` only handles `inputs` and we're handling outputs here. We can move this logic into `foldInputOperands`. But I feel the naming is not so good.

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