From 00c6b4dabd66d51ac672db710d16674ab4fa0c89 Mon Sep 17 00:00:00 2001 From: Mehdi Amini Date: Sat, 28 Mar 2026 11:08:23 +0100 Subject: [PATCH] [MLIR][Vector] Fix crash in foldDenseElementsAttrDestInsertOp on poison index (#188508) 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 bd419f2ba93e..939816262a2b 100644 --- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp +++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp @@ -3764,6 +3764,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 e2045fc526ec..acde86d20d34 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 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 +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