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

[Andrzej Warzyński]

================
@@ -1992,11 +2002,15 @@ vectorizeScalableVectorPrecondition(Operation *op,
     scalableFlags.pop_back();
   }
 
-  // TODO: Support scalable vectorisation for reduction dims
-  if (iterators.back() == utils::IteratorType::reduction)
-    return failure();
+  if (iterators.back() == utils::IteratorType::reduction) {
+    if (iterators.size() != inputVectorSizes.size()) {
+      LDBG("Non-trailing reduction dim requested for scalable "
+           "vectorization\n");
+      return failure();
+    }
+  }
 
-  // If this is not the _last_ parallel dim, 1. above is not met
+  // If this is not the _last_ parallel dim, 1. or 3. above is not met
   if (seenParalell)
     return failure();
----------------
banach-space wrote:

There are two cases here. Should we turn this into a switch statement to combine this somehow?
```cpp  
switch (iterators.back()) {
  case utils::IteratorType::reduction: {
    // Check 3. above is met.
    if (iterators.size() != inputVectorSizes.size()) {
      LDBG("Non-trailing reduction dim requested for scalable "
           "vectorization\n");
      return failure();
      break;
    }
  }
  case utils::IteratorType::parallel: {
    // Check 1. and 2. above are met.
    if (seenParalell) {
      LDBG("Inner parallel dim requested for scalable "
           "vectorization\n");
      return failure();
    }
    break;
  }
```

WDYT? I'm open to suggestion :)

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