[Mlir-commits] [mlir] 237d5c6 - [mlir] Fix vector::InsertElementOp::fold on invalid input

[Mehdi Amini]

Author: Mehdi Amini
Date: 2023-10-31T21:16:28-07:00
New Revision: 237d5c646991b57aacd60b23b3dab509eca9e1aa

URL: https://github.com/llvm/llvm-project/commit/237d5c646991b57aacd60b23b3dab509eca9e1aa
DIFF: https://github.com/llvm/llvm-project/commit/237d5c646991b57aacd60b23b3dab509eca9e1aa.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/IR/VectorOps.cpp
    mlir/test/Dialect/Vector/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index f7b15f98e166543..4e34caa6d8aaba8 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -2507,7 +2507,8 @@ OpFoldResult vector::InsertElementOp::fold(FoldAdaptor adaptor) {
 
   auto attr = llvm::dyn_cast<IntegerAttr>(pos);
   uint64_t posIdx = attr.getInt();
-
+  if (posIdx >= results.size())
+    return {};
   results[posIdx] = src;
 
   return DenseElementsAttr::get(getDestVectorType(), results);

diff  --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index d866c14fcbf2543..f6bb42b1b249153 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -2016,6 +2016,19 @@ func.func @insert_element_fold() -> vector<4xi32> {
 
 // -----
 
+// CHECK-LABEL: func @insert_element_invalid_fold
+func.func @insert_element_invalid_fold() -> vector<1xf32> {
+  // Out-of-bound index here.
+  %c26 = arith.constant 26 : index
+  %cst_2 = arith.constant 1.60215309E+9 : f32
+  %cst_20 = arith.constant dense<1.60215309E+9> : vector<1xf32>
+// CHECK: vector.insertelement
+  %46 = vector.insertelement %cst_2, %cst_20[%c26 : index] : vector<1xf32>
+  return %46 : vector<1xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @extract_element_fold
 //       CHECK:   %[[C:.+]] = arith.constant 5 : i32
 //       CHECK:   return %[[C]]