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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/184626.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp (+11) 
- (modified) mlir/test/Conversion/ConvertToSPIRV/func-signature-vector-unroll.mlir (+12) 


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

``````````

</details>


https://github.com/llvm/llvm-project/pull/184626