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

[Chibuoyim Ogbonna]

================
@@ -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],
----------------
bruteforceboy wrote:

yeah, this makes sense. I guess there are some other passes besides FlattenElementwise where this kind of collapse may be useful, but I've now simplified the collapse here, since this is enough for our use case.

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