[Mlir-commits] [mlir] [mlir][Linalg] Refine how broadcast dims are treated (PR #99015)

[Andrzej Warzyński]

================
@@ -346,6 +346,10 @@ class AffineMap {
   /// returns the resulting values. `this` must be symbol-less.
   SmallVector<int64_t, 4> compose(ArrayRef<int64_t> values) const;
 
+  size_t numOfZeroResults() const;
+
+  AffineMap dropZeros();
----------------
banach-space wrote:

I saw your comment on the other PR and that immediately made me think about this :)

I couldn't find anything that would help for these specific case. However, we do have:
* [getBroadcastDims()](https://github.com/llvm/llvm-project/blob/867ff2d4268ca7eed89a24f100b67b68f5329439/mlir/include/mlir/IR/AffineMap.h#L149-L155)
  * NOTE: to be removed in https://github.com/llvm/llvm-project/pull/99341/,
* [isMinorIdentityWithBroadcasting()](https://github.com/llvm/llvm-project/blob/867ff2d4268ca7eed89a24f100b67b68f5329439/mlir/include/mlir/IR/AffineMap.h#L157-L164),
* [dropResults()](https://github.com/llvm/llvm-project/blob/867ff2d4268ca7eed89a24f100b67b68f5329439/mlir/include/mlir/IR/AffineMap.h#L297-L307).

So, there's scope for re-use. 

Here's what I suggest:
*  Introduce `SmallVector<unsigned> getZeroResults()`
   * It would return the positions of **zero**/**bcast** results.
* Simplify `isMinorIdentityWithBroadcasting()`
  * It would use the newly introduced `getZeroResults()`.
* Replace `dropZeros()` (proposed in this PR) with a call to `dropResults(getZeroResults())`.

I would still have to introduce `getZeroResults()` (i.e. yet another hook in this file), but would avoid introducing `dropZeros()` :) WDYT?

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