[Mlir-commits] [mlir] 3bb2563 - [mlir][vector] Fix crash in `vector.insert` canonicalization (#97801)

[llvmlistbot at llvm.org]

Author: Matthias Springer
Date: 2024-07-05T17:32:51+02:00
New Revision: 3bb25636414ee5b5eaf99c0bdcc191052c9d7ffb

URL: https://github.com/llvm/llvm-project/commit/3bb25636414ee5b5eaf99c0bdcc191052c9d7ffb
DIFF: https://github.com/llvm/llvm-project/commit/3bb25636414ee5b5eaf99c0bdcc191052c9d7ffb.diff

LOG: [mlir][vector] Fix crash in `vector.insert` canonicalization (#97801)

The `InsertOpConstantFolder` assumed that whenever the destination can
be folded to a constant attribute, that attribute must be a
`DenseElementsAttr`. That is is not necessarily the case.

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/IR/VectorOps.cpp
    mlir/test/Dialect/Vector/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 53a6648de014c..55bace2e35f44 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -2851,6 +2851,9 @@ class InsertOpConstantFolder final : public OpRewritePattern<InsertOp> {
     Attribute vectorDestCst;
     if (!matchPattern(destVector, m_Constant(&vectorDestCst)))
       return failure();
+    auto denseDest = llvm::dyn_cast<DenseElementsAttr>(vectorDestCst);
+    if (!denseDest)
+      return failure();
 
     VectorType destTy = destVector.getType();
     if (destTy.isScalable())
@@ -2861,8 +2864,6 @@ class InsertOpConstantFolder final : public OpRewritePattern<InsertOp> {
         !destVector.hasOneUse())
       return failure();
 
-    auto denseDest = llvm::cast<DenseElementsAttr>(vectorDestCst);
-
     Value sourceValue = op.getSource();
     Attribute sourceCst;
     if (!matchPattern(sourceValue, m_Constant(&sourceCst)))

diff  --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 1a674d715ca61..e71a6eb02ea46 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -2729,3 +2729,14 @@ func.func @fold_vector_step_to_constant() -> vector<4xindex> {
   %0 = vector.step : vector<4xindex>
   return %0 : vector<4xindex>
 }
+
+// -----
+
+// CHECK-LABEL: func @vector_insert_const_regression(
+//       CHECK:   llvm.mlir.undef
+//       CHECK:   vector.insert
+func.func @vector_insert_const_regression(%arg0: i8) -> vector<4xi8> {
+  %0 = llvm.mlir.undef : vector<4xi8>
+  %1 = vector.insert %arg0, %0 [0] : i8 into vector<4xi8>
+  return %1 : vector<4xi8>
+}