[Mlir-commits] [mlir] [MLIR][Vector]Add constraints to vector.shape_cast(constant) -> constant (PR #147691)

[Mengmeng Sun]

https://github.com/MengmSun created https://github.com/llvm/llvm-project/pull/147691

We have the case that after `ConvertToLLVMPass` it looks like:

```
...
%4 = llvm.mlir.constant(dense<0.000000e+00> : vector<192xf8E4M3FN>) : vector<192xi8>
%8 = vector.shape_cast %4 : vector<192xi8> to vector<1x192xi8>
%10 = vector.extract %8[0] : vector<192xi8> from vector<1x192xi8>
...
```

Our next pass is `Canonicalizer`. Several months ago everything went smoothly. However recently we met problem that 

```
mlir::DenseElementsAttr mlir::DenseElementsAttr::reshape(mlir::ShapedType): Assertion `newType.getElementType() == curType.getElementType() && "expected the same element type"' failed.
```
and we found that's because a `reshape` operation is added for `vector.shape_cast(constant) -> constant`. This operation will fail if the element type of the source attribute and return type are different. 

So we want to add the constraints that only when **the element type of the source attribute and return type are the same** it will return `reshape` operation to make our case work as before and will not influence other cases.

>From 5e9aa6b6b1fe332f5e0c92b14835484558c00732 Mon Sep 17 00:00:00 2001
From: MengmengSun <mengmengs at nvidia.com>
Date: Wed, 9 Jul 2025 02:53:36 -0700
Subject: [PATCH] [mlir][Vector]Add constraints to vector.shape_cast(constant)
 -> constant

---
 mlir/lib/Dialect/Vector/IR/VectorOps.cpp   |  9 ++++++---
 mlir/test/Dialect/Vector/canonicalize.mlir | 12 ++++++++++++
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 214d2ba7e1b8e..5bbe6704aac48 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -5922,10 +5922,13 @@ OpFoldResult ShapeCastOp::fold(FoldAdaptor adaptor) {
       return bcastOp.getSource();
   }
 
-  // shape_cast(constant) -> constant
+  // shape_cast(constant) -> constant,
+  // if element type of the source and result are the same
   if (auto splatAttr =
-          llvm::dyn_cast_if_present<SplatElementsAttr>(adaptor.getSource()))
-    return splatAttr.reshape(getType());
+          llvm::dyn_cast_if_present<SplatElementsAttr>(adaptor.getSource())) {
+    if (splatAttr.getElementType() == resultType.getElementType())
+      return splatAttr.reshape(getType());
+  }
 
   // shape_cast(poison) -> poison
   if (llvm::dyn_cast_if_present<ub::PoisonAttr>(adaptor.getSource())) {
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 8a9e27378df61..69da8a31d2c9b 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -1002,6 +1002,18 @@ func.func @fold_broadcast_shapecast(%arg0: vector<4xf32>) -> vector<4xf32> {
 
 // -----
 
+// CHECK-LABEL: func @canonicalize_extract_shapecast_different_element_type
+func.func @canonicalize_extract_shapecast_different_element_type()->vector<12xi8> {
+  %0 = llvm.mlir.constant(dense<0.000000e+00> : vector<12xf8E4M3FN>) : vector<12xi8>
+  // CHECK-NOT: vector.shape_cast
+  %1 = vector.shape_cast %0 : vector<12xi8> to vector<1x12xi8>
+  // CHECK-NOT: vector.extract
+  %2 = vector.extract %1[0] : vector<12xi8> from vector<1x12xi8>
+  return %2 : vector<12xi8>
+}
+
+// -----
+
 // CHECK-LABEL: func @canonicalize_broadcast_shapecast_scalar
 //       CHECK:   vector.broadcast
 //   CHECK-NOT:   vector.shape_cast