[Mlir-commits] [mlir] [mlir] Rectify mishandling in `InsertOpConstantFolder` 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 1bf4fbedf467d1c7f7788bb96a8bc1b0b3b62789 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      | 26 ++++++++++++++++---
 .../test/Dialect/LLVMIR/constant-folding.mlir | 13 ++++++++++
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 44bd4aa76ffbd6..b62f1fa5992958 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -2896,10 +2896,17 @@ 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();
+
+    // The `convertIntegerAttr` method specifically handles the case
+    // for `llvm.mlir.constant` which can hold an attribute with a
+    // different type than the return type.
+    if (auto denseSource = llvm::dyn_cast<DenseElementsAttr>(sourceCst)) {
+      for (auto value : denseSource.getValues<Attribute>())
+        insertedValues.push_back(convertIntegerAttr(value, destEltType));
+    } else {
+      insertedValues.push_back(convertIntegerAttr(sourceCst, destEltType));
+    }
 
     auto allValues = llvm::to_vector(denseDest.getValues<Attribute>());
     copy(insertedValues, allValues.begin() + insertBeginPosition);
@@ -2908,6 +2915,17 @@ class InsertOpConstantFolder final : public OpRewritePattern<InsertOp> {
     rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, newAttr);
     return success();
   }
+
+private:
+  /// Converts the expected type to an IntegerAttr if there's
+  /// a mismatch.
+  Attribute convertIntegerAttr(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/Dialect/LLVMIR/constant-folding.mlir b/mlir/test/Dialect/LLVMIR/constant-folding.mlir
index 497d679a12a09f..a3570a11d4a2d0 100644
--- a/mlir/test/Dialect/LLVMIR/constant-folding.mlir
+++ b/mlir/test/Dialect/LLVMIR/constant-folding.mlir
@@ -169,3 +169,16 @@ llvm.func @null_pointer_select(%cond: i1) -> !llvm.ptr {
   // CHECK-NEXT: llvm.return %[[NULLPTR]]
   llvm.return %result : !llvm.ptr
 }
+
+// -----
+
+llvm.func @malloc(i64) -> !llvm.ptr
+
+// CHECK-LABEL: func.func @insert_op
+func.func @insert_op(%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
+}