[Mlir-commits] [mlir] 73ed9a9 - [mlir][Vector] Reject scalable reduction in warp distribution (#225267)

[llvmlistbot at llvm.org]

Author: Cheng-Chieh Liu
Date: 2026-10-01T16:58:45+01:00
New Revision: 73ed9a9f1b5dafa53cba0f68ea51131aaa79fed7

URL: https://github.com/llvm/llvm-project/commit/73ed9a9f1b5dafa53cba0f68ea51131aaa79fed7
DIFF: https://github.com/llvm/llvm-project/commit/73ed9a9f1b5dafa53cba0f68ea51131aaa79fed7.diff

LOG: [mlir][Vector] Reject scalable reduction in warp distribution (#225267)

`WarpOpReduction` distributes a `vector.reduction` across warp lanes by
computing `numElements = vectorType.getShape()[0] / warpSize` and
building the per-lane type as a plain `VectorType::get({numElements},
...)`, with no check for a scalable operand.

For a scalable reduction vector, e.g. `vector<[32]xf32>` with warp size
32, this silently drops the scalable marker: the per-lane type becomes
the fixed `vector<1xf32>` instead of `vector<[1]xf32>`, so only 32 total
elements get reduced across lanes instead of `32*vscale`.

Reject a scalable reduction operand before any IR is created, and add a
negative test.

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
    mlir/test/Dialect/Vector/vector-warp-distribute.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
index 522c555da378a..0340452c75d07 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
@@ -2335,6 +2335,9 @@ struct WarpOpReduction : public WarpDistributionPattern {
     if (vectorType.getRank() != 1)
       return rewriter.notifyMatchFailure(
           warpOp, "Only rank 1 reductions can be distributed.");
+    if (vectorType.isScalable())
+      return rewriter.notifyMatchFailure(
+          warpOp, "Scalable reductions cannot be distributed.");
     // Only warp_size-sized vectors supported.
     if (vectorType.getShape()[0] % warpOp.getWarpSize() != 0)
       return rewriter.notifyMatchFailure(

diff  --git a/mlir/test/Dialect/Vector/vector-warp-distribute.mlir b/mlir/test/Dialect/Vector/vector-warp-distribute.mlir
index e6f5cdcb036c0..896c577e14de2 100644
--- a/mlir/test/Dialect/Vector/vector-warp-distribute.mlir
+++ b/mlir/test/Dialect/Vector/vector-warp-distribute.mlir
@@ -763,6 +763,20 @@ func.func @vector_reduction(%laneid: index) -> (f32) {
 
 // -----
 
+// CHECK-PROP-LABEL: @negative_vector_reduction_scalable
+// CHECK-PROP-NOT: gpu.shuffle
+// CHECK-PROP: vector.reduction <add>, %{{.*}} : vector<[32]xf32> into f32
+func.func @negative_vector_reduction_scalable(%laneid: index) -> (f32) {
+  %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (f32) {
+    %0 = "test.some_def"() : () -> (vector<[32]xf32>)
+    %1 = vector.reduction <add>, %0 : vector<[32]xf32> into f32
+    gpu.yield %1 : f32
+  }
+  return %r : f32
+}
+
+// -----
+
 // CHECK-PROP-LABEL: func @warp_distribute(
 //  CHECK-PROP-SAME:    %[[ID:[a-zA-Z0-9]+]]
 //  CHECK-PROP-SAME:    %[[SRC:[a-zA-Z0-9]+]]