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

[Renato Golin]

================
@@ -72,11 +76,24 @@ struct FoldIntoElementwisePattern : public OpRewritePattern<ElementwiseOp> {
     }
     if (!changed)
       return failure();
-    newMaps.push_back(op.getIndexingMapsArray().back());
 
-    rewriter.replaceOpWithNewOp<ElementwiseOp>(
-        op, newIns, op.getDpsInits()[0], op.getKindAttr(),
-        rewriter.getAffineMapArrayAttr(newMaps));
+    // Keep all output operands and their maps unchanged.
+    SmallVector<AffineMap> originalMaps = op.getIndexingMapsArray();
+    newMaps.append(originalMaps.begin() + op.getNumDpsInputs(),
+                   originalMaps.end());
+    // The maps of the rewritten op must still determine bounds for every loop
----------------
rengolin wrote:

nit: space before the comment

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