[mlir] Rectify mishandling in InsertOpConstantFolder causing crash with assertion when using mlir-opt --canonicalize

[Rajveer100]

Resolves #74236

[llvmbot]

@llvm/pr-subscribers-mlir-llvm
@llvm/pr-subscribers-mlir-vector
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Rajveer Singh Bharadwaj (Rajveer100)

Changes

Resolves #74236

In BuiltinAttributes when we recursively handle complex types and simplify them into integer, float, string, etc., index type needs to be handled for vector-based operations before making an assertion.

---

Full diff: https://github.com/llvm/llvm-project/pull/88314.diff

2 Files Affected:

• (modified) mlir/lib/IR/BuiltinAttributes.cpp (+2)
• (added) mlir/test/mlir-opt/issue-74236.mlir (+10)

diff --git a/mlir/lib/IR/BuiltinAttributes.cpp b/mlir/lib/IR/BuiltinAttributes.cpp
index 89b1ed67f5d067..adca7af7a91336 100644
--- a/mlir/lib/IR/BuiltinAttributes.cpp
+++ b/mlir/lib/IR/BuiltinAttributes.cpp
@@ -958,6 +958,8 @@ DenseElementsAttr DenseElementsAttr::get(ShapedType type,
       intVal = floatAttr.getValue().bitcastToAPInt();
     } else {
       auto intAttr = llvm::cast<IntegerAttr>(values[i]);
+      if (intAttr.getType().isIndex())
+        continue;
       assert(intAttr.getType() == eltType &&
              "expected integer attribute type to equal element type");
       intVal = intAttr.getValue();
diff --git a/mlir/test/mlir-opt/issue-74236.mlir b/mlir/test/mlir-opt/issue-74236.mlir
new file mode 100644
index 00000000000000..559daead0cf9c7
--- /dev/null
+++ b/mlir/test/mlir-opt/issue-74236.mlir
@@ -0,0 +1,10 @@
+// RUN: mlir-opt -split-input-file -verify-diagnostics %s
+
+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
+}

[Rajveer100]

@joker-eph
This currently breaks few tests since I have just put a continue, I think we probably want to check if there are two arguments only then we do so.

Let me know the direction for this.

[Rajveer100]

cc @kazutakahirata

[stellaraccident]

Thank you for the patch and repro. There is something subtle going on here that I can't just spot in a code review. If others can't either, it will need someone familiar to step through it and fix properly vs just giving advice. In that case, probably needs an issue vs PR so we don't lose it.

[joker-eph]

It's not clear to me that the bug is in DenseElementsAttr::get actually, I would push the fix on the API call site instead.

However digging a bit more, the underlying problem seems that our folding hook tolerates that an attribute can be returned for an SSA Value with a type mismatch. This happens only for llvm.constant in practice, because it accepts index types: %1 = llvm.mlir.constant(1 : index) : i64.
If we continue to accept this, then every folder should check that the attribute they get from matching a constant actually correspond to the expected SSA Type (in this case the fix would be in InsertOpConstantFolder for example).

[Rajveer100]

@stellaraccident @joker-eph

[joker-eph]

LG with a few comments