fixup! [mlir][Linalg] Refine how broadcast dims are treated

Addressing PR comments

Author: banach-space

mlir/include/mlir/IR/AffineMap.h

@@ -354,9 +354,23 @@ 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;
+  /// Returns the number of "zero" results (constant values == 0) in this map.
+  ///
+  /// Example:
+  ///   * For `(d0, d1) -> (d0, d1, 0)` returns 1
+  ///   * For `(d0, d1, d2) -> (d0, d1)` returns 0
+  ///   * For `(d0, d1, d2) -> (d0, 0, d1, 0, d2)` returns 2
+  size_t getNumOfZeroResults() const;
 
-  AffineMap dropZeros();
+  /// Returns the AffineMap resulting from removing "zero" results (constant
+  /// values == 0) from this map.
+  ///
+  /// Example:
+  ///   * For `(d0, d1) -> (d0, d1, 0)` returns `(d0, d1) -> (d0, d1)`
+  ///   * For `(d0, d1, d2) -> (d0, d1)` returns `(d0, d1, d2) -> (d0, d1)`
+  ///   * For `(d0, d1, d2) -> (d0, 0, d1, 0, d2)` returns
+  ///     `(d0, d1, d2) -> (d0, d1, d2)`
+  AffineMap dropZeroResults();
 
   /// Returns true if the AffineMap represents a subset (i.e. a projection) of a
   /// symbol-less permutation map. `allowZeroInResults` allows projected

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

@@ -476,7 +476,8 @@ static AffineMap reindexIndexingMap(AffineMap map) {
   assert(map.isProjectedPermutation(/*allowZeroInResults=*/true) &&
          "expected projected permutation");
   auto res = compressUnusedDims(map);
-  assert(res.getNumDims() == (res.getNumResults() - res.numOfZeroResults()) &&
+  assert(res.getNumDims() ==
+             (res.getNumResults() - res.getNumOfZeroResults()) &&
          "expected reindexed map with same number of dims and results");
   return res;
 }
@@ -651,7 +652,7 @@ static Value buildVectorWrite(RewriterBase &rewriter, Value value,
   //    vector<1x16x16xty>
   // rather than:
   //    vector<1x16x16x0xty>
-  auto opOperantMapWithoutZeros = opOperandMap.dropZeros();
+  AffineMap opOperantMapWithoutZeros = opOperandMap.dropZeroResults();
   write =
       state.maskOperation(rewriter, write, linalgOp, opOperantMapWithoutZeros);
 

mlir/lib/IR/AffineMap.cpp

@@ -553,18 +553,6 @@ AffineMap AffineMap::dropResults(const llvm::SmallBitVector &positions) const {
   return AffineMap::get(getNumDims(), getNumSymbols(), exprs, getContext());
 }
 
-AffineMap AffineMap::dropZeros() {
-  auto exprs = llvm::to_vector<4>(getResults());
-  SmallVector<AffineExpr, 8> newExprs;
-
-  for (auto expr : getResults()) {
-    auto constExpr = dyn_cast<AffineConstantExpr>(expr);
-    if (!constExpr)
-      newExprs.push_back(expr);
-  }
-  return AffineMap::get(getNumDims(), getNumSymbols(), newExprs, getContext());
-}
-
 AffineMap AffineMap::compose(AffineMap map) const {
   assert(getNumDims() == map.getNumResults() && "Number of results mismatch");
   // Prepare `map` by concatenating the symbols and rewriting its exprs.
@@ -604,7 +592,7 @@ SmallVector<int64_t, 4> AffineMap::compose(ArrayRef<int64_t> values) const {
   return res;
 }
 
-size_t AffineMap::numOfZeroResults() const {
+size_t AffineMap::getNumOfZeroResults() const {
   size_t res = 0;
   for (auto expr : getResults()) {
     auto constExpr = dyn_cast<AffineConstantExpr>(expr);
@@ -615,6 +603,18 @@ size_t AffineMap::numOfZeroResults() const {
   return res;
 }
 
+AffineMap AffineMap::dropZeroResults() {
+  auto exprs = llvm::to_vector(getResults());
+  SmallVector<AffineExpr> newExprs;
+
+  for (auto expr : getResults()) {
+    auto constExpr = dyn_cast<AffineConstantExpr>(expr);
+    if (!constExpr || constExpr.getValue() != 0)
+      newExprs.push_back(expr);
+  }
+  return AffineMap::get(getNumDims(), getNumSymbols(), newExprs, getContext());
+}
+
 bool AffineMap::isProjectedPermutation(bool allowZeroInResults) const {
   if (getNumSymbols() > 0)
     return false;