[Mlir-commits] [mlir] [MLIR][Linalg] Scalable Vectorization of Reduction (PR #97788)

[Andrzej Warzyński]

================
@@ -1947,13 +1956,30 @@ vectorizeScalableVectorPrecondition(Operation *op,
   if (inputVectorSizes.empty())
     return success();
 
+  auto linalgOp = dyn_cast<LinalgOp>(op);
+  if (linalgOp && isLinalgReduction(linalgOp)) {
+    LDBG("Checking reduce op dims for scalable vectorization\n");
+    auto iteratorTypes = linalgOp.getIteratorTypesArray();
+    assert(iteratorTypes.size() == inputScalableVecDims.size() &&
+           "Number of iterator types and input scalable dims mismatch");
+    // For now, only support scalable vectorization of a reduction on the
+    // trailing dim.
+    for (size_t i = 0; i < inputScalableVecDims.size() - 1; ++i) {
+      if (inputScalableVecDims[i] && isReductionIterator(iteratorTypes[i])) {
+        LDBG("Non-trailing reduction dim requested for scalable "
+             "vectorization\n");
+        return failure();
+      }
+    }
+    return success();
+  }
----------------
banach-space wrote:

Wouldn't this be sufficient?
```
// Only the trailing scalable dim is allowed to be scalable.
if (llvm::all_of(ArrayRef<bool>(inputScalableVecDims).drop_back(1), [](bool flag) {return flag == false;})
  return failure();
```
As in, we only need to make sure that all the flags except for the trailing one are `false`. Btw, I might have made a typo - let me know if my suggestion "doesn't work" for you :)

Also, we shouldn't return `success` until all pre-conditions are checked (there's more further down).

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