[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:

if you recall the identity case above, which is now removed, for something like `linalg.broadcast ins(%a : tensor<4x8xf32>) outs(%b : tensor<4x8xf32>) dimensions = []`, it's not always the case that the rank is 0. If we reject non-zero here, then we will make flatten to fail on this broadcast. 

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