[Mlir-commits] [mlir] [mlir][spirv] Fix crash in FuncOpVectorUnroll for dynamic-shaped tensor args (PR #184626)

[Mehdi Amini]

https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/184626

>From ecb94fc9ec79fd2b281f88c83295f5e54148c245 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Wed, 4 Mar 2026 06:42:21 -0800
Subject: [PATCH] [mlir][spirv] Fix crash in FuncOpVectorUnroll for
 dynamic-shaped tensor args

FuncOpVectorUnroll::matchAndRewrite iterates over all function arguments
and creates a zero-constant placeholder for any non-vector argument via
rewriter.getZeroAttr(type). For dynamically-shaped tensor types such as
tensor<?x2xi8>, getZeroAttr ultimately calls DenseIntOrFPElementsAttr::getRaw
which asserts that the type has a static shape.

Add an early pre-check that returns failure() for any non-vector argument
with a dynamically-shaped type. This must be done before any IR modifications
are made, since returning failure() after inlining the function body would
leave the IR in an inconsistent state.

Fixes #148369
---
 .../lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp |  9 +++++++++
 .../ConvertToSPIRV/func-signature-vector-unroll.mlir | 12 ++++++++++++
 2 files changed, 21 insertions(+)

diff --git a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
index 973c16e62bb12..2c9e9c040d460 100644
--- a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
+++ b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
@@ -1050,6 +1050,15 @@ struct FuncOpVectorUnroll final : OpRewritePattern<func::FuncOp> {
       return failure();
     }
 
+    // Bail out early for dynamically-shaped argument types: getZeroAttr
+    // requires a statically-shaped type. VectorType is always statically
+    // shaped, so this correctly skips it without a special-case guard.
+    if (llvm::any_of(fnType.getInputs(), [](Type argType) {
+          auto shapedType = dyn_cast<ShapedType>(argType);
+          return shapedType && !shapedType.hasStaticShape();
+        }))
+      return failure();
+
     // Create a new func op with the original type and copy the function body.
     auto newFuncOp = func::FuncOp::create(rewriter, funcOp.getLoc(),
                                           funcOp.getName(), fnType);
diff --git a/mlir/test/Conversion/ConvertToSPIRV/func-signature-vector-unroll.mlir b/mlir/test/Conversion/ConvertToSPIRV/func-signature-vector-unroll.mlir
index 211d6c90243bd..912cd308a8fbe 100644
--- a/mlir/test/Conversion/ConvertToSPIRV/func-signature-vector-unroll.mlir
+++ b/mlir/test/Conversion/ConvertToSPIRV/func-signature-vector-unroll.mlir
@@ -262,3 +262,15 @@ func.func @legal_params_for_loop(%arg0: i32, %arg1: i32, %arg2: i32) -> i32 {
   }
   return %result : i32
 }
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/148369
+// Functions with dynamically-shaped tensor arguments should not crash.
+// The pattern should bail out and leave such functions unchanged.
+
+// CHECK-LABEL: @dynamic_tensor_arg
+// CHECK-SAME: (%[[ARG0:.+]]: tensor<?x2xi8>)
+func.func @dynamic_tensor_arg(%arg0: tensor<?x2xi8>) {
+  return
+}