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

[Matthias Springer]

https://github.com/matthias-springer updated https://github.com/llvm/llvm-project/pull/97801

>From 9db3421389cbbd0bfef39c088420d3323349376a Mon Sep 17 00:00:00 2001
From: Matthias Springer <mspringer at nvidia.com>
Date: Fri, 5 Jul 2024 11:22:12 +0200
Subject: [PATCH] [mlir][vector] Fix crash in `vector.insert` canonicalization

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.
---
 mlir/lib/Dialect/Vector/IR/VectorOps.cpp   |  5 +++--
 mlir/test/Dialect/Vector/canonicalize.mlir | 11 +++++++++++
 2 files changed, 14 insertions(+), 2 deletions(-)

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