Skip to content

Commit 237d5c6

Browse files
committed
[mlir] Fix vector::InsertElementOp::fold on invalid input
The IR is valid, but UB: there is an out-of-bound index for the position to insert inside the vector. We should just ignore this in the folder. Fixes #70884
1 parent ad1b7a4 commit 237d5c6

File tree

2 files changed

+15
-1
lines changed

2 files changed

+15
-1
lines changed

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2507,7 +2507,8 @@ OpFoldResult vector::InsertElementOp::fold(FoldAdaptor adaptor) {
25072507

25082508
auto attr = llvm::dyn_cast<IntegerAttr>(pos);
25092509
uint64_t posIdx = attr.getInt();
2510-
2510+
if (posIdx >= results.size())
2511+
return {};
25112512
results[posIdx] = src;
25122513

25132514
return DenseElementsAttr::get(getDestVectorType(), results);

mlir/test/Dialect/Vector/canonicalize.mlir

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2016,6 +2016,19 @@ func.func @insert_element_fold() -> vector<4xi32> {
20162016

20172017
// -----
20182018

2019+
// CHECK-LABEL: func @insert_element_invalid_fold
2020+
func.func @insert_element_invalid_fold() -> vector<1xf32> {
2021+
// Out-of-bound index here.
2022+
%c26 = arith.constant 26 : index
2023+
%cst_2 = arith.constant 1.60215309E+9 : f32
2024+
%cst_20 = arith.constant dense<1.60215309E+9> : vector<1xf32>
2025+
// CHECK: vector.insertelement
2026+
%46 = vector.insertelement %cst_2, %cst_20[%c26 : index] : vector<1xf32>
2027+
return %46 : vector<1xf32>
2028+
}
2029+
2030+
// -----
2031+
20192032
// CHECK-LABEL: func @extract_element_fold
20202033
// CHECK: %[[C:.+]] = arith.constant 5 : i32
20212034
// CHECK: return %[[C]]

0 commit comments

Comments
 (0)