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

[Mehdi Amini]

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

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

>From 5d9548a2c0bef82dbd3f0e58c1a81e77a406e295 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 | 11 +++++++++++
 .../ConvertToSPIRV/func-signature-vector-unroll.mlir | 12 ++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
index 973c16e62bb12..2a36bf8a67a90 100644
--- a/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
+++ b/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp
@@ -1050,6 +1050,17 @@ struct FuncOpVectorUnroll final : OpRewritePattern<func::FuncOp> {
       return failure();
     }
 
+    // Bail out early for non-vector argument types that cannot have a zero
+    // constant created (e.g., dynamically-shaped tensors), as getZeroAttr
+    // requires a statically-shaped type.
+    for (Type argType : fnType.getInputs()) {
+      if (dyn_cast<VectorType>(argType))
+        continue;
+      if (auto shapedType = dyn_cast<ShapedType>(argType))
+        if (!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
+}