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

[Andrzej Warzyński]

================
@@ -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)
----------------
banach-space wrote:

That one was a bad example :) Basically, that was a NOP and IMO we shouldn't optimise/tweak for that case - it should be optimised away. 

Is there any other example that we should consider? I am still thinking that non-zero ranks should be removed.

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