[Mlir-commits] [mlir] [MLIR][Vector] Fix crash in foldDenseElementsAttrDestInsertOp on poison index (PR #188508)

[Mehdi Amini]

https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/188508

>From f3bd501aec27ef217db675ecaa8c63a78c4c5e7c Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Wed, 25 Mar 2026 08:16:30 -0700
Subject: [PATCH 1/2] [MLIR][Vector] Fix crash in
 foldDenseElementsAttrDestInsertOp on poison index

When a dynamic index of -1 (the kPoisonIndex sentinel) was folded into
the static position of a vector.insert op, foldDenseElementsAttrDestInsertOp
would proceed to call calculateInsertPosition, which returned -1. The
subsequent iterator arithmetic (allValues.begin() + (-1)) was undefined
behaviour, causing an assertion in DenseElementsAttr::get.

Fix by bailing out early in foldDenseElementsAttrDestInsertOp when any
static position equals kPoisonIndex, consistent with how
InsertChainFullyInitialized already guards this case.

Fixes #188404

Assisted-by: Claude Code
---
 mlir/lib/Dialect/Vector/IR/VectorOps.cpp   |  5 +++++
 mlir/test/Dialect/Vector/canonicalize.mlir | 19 +++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 73632875ca9e2..5c8c2417bfaa6 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -3763,6 +3763,11 @@ foldDenseElementsAttrDestInsertOp(InsertOp insertOp, Attribute srcAttr,
       !insertOp->hasOneUse())
     return {};
 
+  // Bail out on poison indices (kPoisonIndex = -1) to avoid computing an
+  // invalid (negative) linearized position which would cause UB below.
+  if (is_contained(insertOp.getStaticPosition(), InsertOp::kPoisonIndex))
+    return {};
+
   // Calculate the linearized position for inserting elements.
   int64_t insertBeginPosition =
       calculateInsertPosition(destTy, insertOp.getStaticPosition());
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index e2045fc526ec1..341f044b5112f 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -3807,6 +3807,25 @@ func.func @insert_vector_poison_idx_non_cst(%a: vector<4x5xf32>, %b: vector<5xf3
 
 // -----
 
+// Similar to the test above, but now the destination is a dense constant.
+// foldDenseElementsAttrDestInsertOp must not crash when the dynamic index
+// constant is -1 (the poison sentinel). Instead it should fold to ub.poison.
+// Regression test for https://github.com/llvm/llvm-project/issues/188404
+
+// CHECK-LABEL: @insert_scalar_poison_idx_dense_dest
+func.func @insert_scalar_poison_idx_dense_dest() -> vector<2xf16> {
+  // CHECK-NEXT: %[[UB:.*]] = ub.poison : vector<2xf16>
+  //  CHECK-NOT: vector.insert
+  // CHECK-NEXT: return %[[UB]] : vector<2xf16>
+  %cst = arith.constant dense<0.0> : vector<2xf16>
+  %idx = arith.constant -1 : index
+  %val = arith.constant 2.5 : f16
+  %0 = vector.insert %val, %cst [%idx] : f16 into vector<2xf16>
+  return %0 : vector<2xf16>
+}
+
+// -----
+
 // Similar to test above, but now the index is out-of-bounds.
 
 // CHECK-LABEL: @no_fold_insert_scalar_idx_oob

>From 194a5ef8500c769d032112e382a78c1d07332c58 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Sat, 28 Mar 2026 10:48:32 +0100
Subject: [PATCH 2/2] Update mlir/test/Dialect/Vector/canonicalize.mlir
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Andrzej Warzyński <andrzej.warzynski at gmail.com>
---
 mlir/test/Dialect/Vector/canonicalize.mlir | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 341f044b5112f..acde86d20d346 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -3807,9 +3807,9 @@ func.func @insert_vector_poison_idx_non_cst(%a: vector<4x5xf32>, %b: vector<5xf3
 
 // -----
 
-// Similar to the test above, but now the destination is a dense constant.
-// foldDenseElementsAttrDestInsertOp must not crash when the dynamic index
-// constant is -1 (the poison sentinel). Instead it should fold to ub.poison.
+// Similar to the test above, but with a dense constant destination. This exercises
+// foldDenseElementsAttrDestInsertOp, which must not crash when the dynamic index
+// constant is -1 (the poison sentinel). The IR should instead fold to ub.poison.
 // Regression test for https://github.com/llvm/llvm-project/issues/188404
 
 // CHECK-LABEL: @insert_scalar_poison_idx_dense_dest