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

[Hugo Trachino]

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

Is there a usage of collapsing a broadcast op which does not lead to `dimensions = [0]`, and hence avoid all this ? Collapsing a broadcast seems illegal to me unless the input is 0D. 

I have been considering the case where we collapse only the input dimensions e.g. `broadcast <2x3> - > <2x4x3>` collapsed to `broadcast <6> -> <6x4>` but that still isnt legal as the dimension order is not kept. I think, the only valid collapsing is on contiguous input dimensions which are contiguous in the output as well. e.g. `broadcast <2x3> - > <2x3x4>` collapsed to `broadcast <6> -> <6x4>`. That is not implemented for FlattenElementwise, is it used somewhere else ?

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