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

[llvmlistbot at llvm.org]

https://github.com/lijinpei-amd updated https://github.com/llvm/llvm-project/pull/203291

>From 284588bc96e6caba36a3c4268ecf0cd53ee0be0b Mon Sep 17 00:00:00 2001
From: Li Jinpei <jinpli at amd.com>
Date: Thu, 11 Jun 2026 22:21:22 +0800
Subject: [PATCH] [mlir][spirv] Fix crash on 0-D vectors in vector unrolling

`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
---
 .../SPIRV/Transforms/SPIRVConversion.cpp      |  6 +++++
 .../ConvertToSPIRV/vector-unroll.mlir         | 22 +++++++++++++++++++
 2 files changed, 28 insertions(+)

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>
+}