[Mlir-commits] [mlir] [mlir][vector] Fix cast incompatible type bug in `ShuffleOp::fold` (PR #150037)

[Longsheng Mou]

https://github.com/CoTinker created https://github.com/llvm/llvm-project/pull/150037

This PR uses `dyn_cast` instead of `cast` to avoid a crash when the constant attribute is not a `DenseElementsAttr`. Fixes #149325.

>From 3e68a703e55f3d3a658f52dfb2e6809dd7260baa Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Tue, 22 Jul 2025 22:45:52 +0800
Subject: [PATCH] [mlir][vector] Fix cast incompatible type bug in
 `ShuffleOp::fold`

This PR uses `dyn_cast` instead of `cast` to avoid a crash when the
constant attribute is not a `DenseElementsAttr`.
---
 mlir/lib/Dialect/Vector/IR/VectorOps.cpp   | 12 ++++++++----
 mlir/test/Dialect/Vector/canonicalize.mlir | 13 +++++++++++++
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 56f748fbbe1d6..d2e2f7c125598 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -3057,13 +3057,17 @@ OpFoldResult vector::ShuffleOp::fold(FoldAdaptor adaptor) {
   SmallVector<Attribute> v1Elements, v2Elements;
   Attribute poisonElement;
   if (!isV2Poison) {
-    v2Elements =
-        to_vector(cast<DenseElementsAttr>(v2Attr).getValues<Attribute>());
+    auto v2DenseAttr = dyn_cast<DenseElementsAttr>(v2Attr);
+    if (!v2DenseAttr)
+      return {};
+    v2Elements = to_vector(v2DenseAttr.getValues<Attribute>());
     poisonElement = v2Elements[0];
   }
   if (!isV1Poison) {
-    v1Elements =
-        to_vector(cast<DenseElementsAttr>(v1Attr).getValues<Attribute>());
+    auto v1DenseAttr = dyn_cast<DenseElementsAttr>(v1Attr);
+    if (!v1DenseAttr)
+      return {};
+    v1Elements = to_vector(v1DenseAttr.getValues<Attribute>());
     poisonElement = v1Elements[0];
   }
 
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 6809122974545..1461c30162c5f 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -2347,6 +2347,19 @@ func.func @shuffle_1d() -> vector<4xi32> {
 
 // -----
 
+// Ensure shuffle dense resource elements not crash.
+
+// CHECK-LABEL: func.func @shuffle_1d_dense_resource
+//       CHECK:    vector.shuffle
+func.func @shuffle_1d_dense_resource() -> vector<4xi32> {
+  %v0 = arith.constant dense_resource<__elided__> : vector<3xi32>
+  %v1 = arith.constant dense_resource<__elided__> : vector<3xi32>
+  %shuffle = vector.shuffle %v0, %v1 [3, 2, 5, 1] : vector<3xi32>, vector<3xi32>
+  return %shuffle : vector<4xi32>
+}
+
+// -----
+
 // Check that poison indices pick the first element of the first non-poison
 // input vector. That is, %v[0] (i.e., 5) in this test.