[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());
+    // The maps of the rewritten op must still determine bounds for every loop
+    // dimension. Folding a broadcast can otherwise drop the only map result
+    // that covers a dimension.
+    // See `generic_broadcast_not_folded_non_invertible` in
+    // mlir/test/Dialect/Linalg/elementwise/fold.mlir for an example.
+    if (!inversePermutation(concatAffineMaps(newMaps, op.getContext())))
+      return failure();
+    auto newOp =
----------------
ChuanqiXu9 wrote:

Done by using `modifyOpInPlace`.

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