[mlir][vector] Fix cast incompatible type bug in ShuffleOp::fold (#150037)

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

Author: CoTinker

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];
   }
 

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.