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

[Chibuoyim Ogbonna]

================
@@ -1809,12 +1809,34 @@ GenericOp cloneToCollapsedOp<GenericOp>(RewriterBase &rewriter,
   return collapsedOp;
 }
 
+/// Collapse a `BroadcastOp`. Flattening leaves a single dimension, so a 0-D
+/// input broadcasts into it (`dimensions = [0]`) and any other input adds none.
+template <>
+BroadcastOp
+cloneToCollapsedOp<BroadcastOp>(RewriterBase &rewriter, BroadcastOp origOp,
+                                const CollapsingInfo &collapsingInfo) {
+  SmallVector<Value> inputOperands, outputOperands;
+  SmallVector<Type> resultTypes;
+  collapseOperandsAndResults(origOp, collapsingInfo, rewriter, inputOperands,
+                             outputOperands, resultTypes);
+
+  SmallVector<int64_t> newDimensions;
+  if (origOp.getInput().getType().getRank() == 0)
----------------
bruteforceboy wrote:

That makes sense. The identity case is the only example I can think of. I have now updated the broadcast collapse to handle only rank zero. With this, the identity case will remain as it is.

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