[Mlir-commits] [mlir] d74caa4 - [mlir][spirv] Fix crash on 0-D vectors in vector unrolling (#203291)

[llvmlistbot at llvm.org]

Author: lijinpei-amd
Date: 2026-06-17T16:04:07+08:00
New Revision: d74caa40ec86ceb2c3e6e18d45adeab8fb40b4a4

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

LOG: [mlir][spirv] Fix crash on 0-D vectors in vector unrolling (#203291)

`getTargetShape` and `getNativeVectorShape` called `getShape().back()`
without checking for rank-0 vectors, whose shape is empty. This crashed
when the SPIR-V vector unrolling pass processed a function returning a
0-D vector (e.g. `vector<f32>`) or a 0-D elementwise op.

0-D vectors have no dimension to unroll along and are not SPIR-V vector
types, so bail out and leave them unchanged in both paths.

Fixes https://github.com/llvm/llvm-project/issues/203220

Added: 
    

Modified: 
    mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
    mlir/test/Conversion/ConvertToSPIRV/vector-unroll.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
index 2c9e9c040d460..858b89728d432 100644
--- a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
+++ b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
@@ -53,6 +53,10 @@ static std::optional<SmallVector<int64_t>> getTargetShape(VectorType vecType) {
                << "--scalable vectors are not supported -> BAIL\n");
     return std::nullopt;
   }
+  if (vecType.getRank() == 0) {
+    LLVM_DEBUG(llvm::dbgs() << "--0-D vectors are not supported -> BAIL\n");
+    return std::nullopt;
+  }
   SmallVector<int64_t> unrollShape = llvm::to_vector<4>(vecType.getShape());
   std::optional<SmallVector<int64_t>> targetShape = SmallVector<int64_t>(
       1, mlir::spirv::getComputeVectorSize(vecType.getShape().back()));
@@ -1471,6 +1475,8 @@ std::optional<SmallVector<int64_t>>
 mlir::spirv::getNativeVectorShape(Operation *op) {
   if (OpTrait::hasElementwiseMappableTraits(op) && op->getNumResults() == 1) {
     if (auto vecType = dyn_cast<VectorType>(op->getResultTypes()[0])) {
+      if (vecType.getRank() == 0)
+        return std::nullopt;
       SmallVector<int64_t> nativeSize(vecType.getRank(), 1);
       nativeSize.back() =
           mlir::spirv::getComputeVectorSize(vecType.getShape().back());

diff  --git a/mlir/test/Conversion/ConvertToSPIRV/vector-unroll.mlir b/mlir/test/Conversion/ConvertToSPIRV/vector-unroll.mlir
index 0957f67690b97..6b013f96feb4d 100644
--- a/mlir/test/Conversion/ConvertToSPIRV/vector-unroll.mlir
+++ b/mlir/test/Conversion/ConvertToSPIRV/vector-unroll.mlir
@@ -140,3 +140,25 @@ func.func @unroll_from_elements_2d(%arg0: f32, %arg1: f32, %arg2: f32, %arg3: f3
   // return %[[RES0]], %%[[RES1]] : vector<2xf32>, vector<2xf32>
   return %1 : vector<2x2xf32>
 }
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/203220
+
+// CHECK-LABEL: @zero_d_vector_return
+// CHECK-SAME: (%[[ARG0:.+]]: vector<f32>)
+func.func @zero_d_vector_return(%arg0: vector<f32>) -> vector<f32> {
+  // CHECK: return %[[ARG0]] : vector<f32>
+  return %arg0 : vector<f32>
+}
+
+// -----
+
+// CHECK-LABEL: @zero_d_vector_elementwise
+// CHECK-SAME: (%[[ARG0:.+]]: vector<f32>, %[[ARG1:.+]]: vector<f32>)
+func.func @zero_d_vector_elementwise(%arg0: vector<f32>, %arg1: vector<f32>) -> vector<f32> {
+  // CHECK: %[[ADD:.*]] = arith.addf %[[ARG0]], %[[ARG1]] : vector<f32>
+  // CHECK: return %[[ADD]] : vector<f32>
+  %0 = arith.addf %arg0, %arg1 : vector<f32>
+  return %0 : vector<f32>
+}