[Mlir-commits] [mlir] [MLIR][Linalg] Recompute linalg.broadcast dimensions when flattening (PR #213641)

[Andrzej Warzyński]

================
@@ -1809,12 +1809,49 @@ GenericOp cloneToCollapsedOp<GenericOp>(RewriterBase &rewriter,
   return collapsedOp;
 }
 
+/// Collapse a `BroadcastOp`, recomputing its `dimensions` for the collapsed
+/// iteration space. Returns null if the collapse is not expressible as a
+/// broadcast.
+template <>
+BroadcastOp
+cloneToCollapsedOp<BroadcastOp>(RewriterBase &rewriter, BroadcastOp origOp,
+                                const CollapsingInfo &collapsingInfo) {
+  ArrayRef<int64_t> broadcastDims = origOp.getDimensions();
+  SmallVector<int64_t> newDimensions;
+  for (auto [collapsedDim, foldedDims] :
+       llvm::enumerate(collapsingInfo.getCollapsedOpToOrigOpMapping())) {
+    size_t numBroadcast = llvm::count_if(foldedDims, [&](int64_t d) {
+      return llvm::is_contained(broadcastDims, d);
+    });
+    // A collapsed dimension is a broadcast dimension iff all the dimensions it
+    // folds are; a mix of the two cannot be represented as a broadcast.
+    if (numBroadcast != 0 && numBroadcast != foldedDims.size())
+      return nullptr;
+    if (numBroadcast != 0)
+      newDimensions.push_back(collapsedDim);
+  }
+
+  SmallVector<Value> inputOperands, outputOperands;
+  SmallVector<Type> resultTypes;
+  collapseOperandsAndResults(origOp, collapsingInfo, rewriter, inputOperands,
+                             outputOperands, resultTypes);
+
+  return BroadcastOp::create(rewriter, origOp.getLoc(), inputOperands[0],
----------------
banach-space wrote:

> Is there a usage of collapsing a broadcast op which does not lead to dimensions = [0], and hence avoid all this ?

+1, lets keep things simple

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