[Mlir-commits] [mlir] [mlir] Rectify mishandling in `DenseElementsAttr::get(...)` causing crash with assertion when using `mlir-opt --canonicalize` (PR #88314)

[Rajveer Singh Bharadwaj]

https://github.com/Rajveer100 updated https://github.com/llvm/llvm-project/pull/88314

>From 8e449016f7f5b9836fdeb501b80e4b7e3e0f655e Mon Sep 17 00:00:00 2001
From: Rajveer <rajveer.developer at icloud.com>
Date: Thu, 11 Apr 2024 00:40:32 +0530
Subject: [PATCH] [mlir] Rectify mishandling in `InsertOpConstantFolder`
 causing crash with assertion when using `mlir-opt --canonicalize`

Resolves #74236
---
 mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 21 +++++++++++++++++----
 mlir/test/mlir-opt/issue-74236.mlir      | 10 ++++++++++
 2 files changed, 27 insertions(+), 4 deletions(-)
 create mode 100644 mlir/test/mlir-opt/issue-74236.mlir

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 44bd4aa76ffbd6..5df5e7a38caeda 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -2896,10 +2896,13 @@ class InsertOpConstantFolder final : public OpRewritePattern<InsertOp> {
         linearize(completePositions, computeStrides(destTy.getShape()));
 
     SmallVector<Attribute> insertedValues;
-    if (auto denseSource = llvm::dyn_cast<DenseElementsAttr>(sourceCst))
-      llvm::append_range(insertedValues, denseSource.getValues<Attribute>());
-    else
-      insertedValues.push_back(sourceCst);
+    Type destEltType = destTy.getElementType();
+    if (auto denseSource = llvm::dyn_cast<DenseElementsAttr>(sourceCst)) {
+      for (auto value : denseSource.getValues<Attribute>())
+        insertedValues.push_back(castOrConvertAttr(value, destEltType));
+    } else {
+      insertedValues.push_back(castOrConvertAttr(sourceCst, destEltType));
+    }
 
     auto allValues = llvm::to_vector(denseDest.getValues<Attribute>());
     copy(insertedValues, allValues.begin() + insertBeginPosition);
@@ -2908,6 +2911,16 @@ class InsertOpConstantFolder final : public OpRewritePattern<InsertOp> {
     rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, newAttr);
     return success();
   }
+
+private:
+  /// Helper function to cast or convert attribute to the expected type.
+  Attribute castOrConvertAttr(Attribute attr, Type expectedType) const {
+    if (auto intAttr = attr.dyn_cast<IntegerAttr>()) {
+      if (intAttr.getType() != expectedType)
+        return IntegerAttr::get(expectedType, intAttr.getInt());
+    }
+    return attr;
+  }
 };
 
 } // namespace
diff --git a/mlir/test/mlir-opt/issue-74236.mlir b/mlir/test/mlir-opt/issue-74236.mlir
new file mode 100644
index 00000000000000..411814bd18cc71
--- /dev/null
+++ b/mlir/test/mlir-opt/issue-74236.mlir
@@ -0,0 +1,10 @@
+// RUN: mlir-opt -split-input-file -verify-diagnostics %s -canonicalize
+
+llvm.func @malloc(i64) -> !llvm.ptr
+func.func @func2(%arg0: index, %arg1: memref<13x13xi64>, %arg2: index) {
+  %cst_7 = arith.constant dense<1526248407> : vector<1xi64>
+  %1 = llvm.mlir.constant(1 : index) : i64
+  %101 = vector.insert %1, %cst_7 [0] : i64 into vector<1xi64>
+  vector.print %101 : vector<1xi64>
+  return
+}