[Mlir-commits] [mlir] ebceb73 - [mlir][vector] Update the folder for vector.{insert|extract} (#136579)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Apr 24 02:43:28 PDT 2025
Author: Andrzej Warzyński
Date: 2025-04-24T10:43:25+01:00
New Revision: ebceb732d5afa7b97b73b34ce078a58c654ff0a0
URL: https://github.com/llvm/llvm-project/commit/ebceb732d5afa7b97b73b34ce078a58c654ff0a0
DIFF: https://github.com/llvm/llvm-project/commit/ebceb732d5afa7b97b73b34ce078a58c654ff0a0.diff
LOG: [mlir][vector] Update the folder for vector.{insert|extract} (#136579)
This is a minor follow-up to #135498. It ensures that operations like
the following are not treated as out-of-bounds accesses and can be
folded correctly (*):
```mlir
%c_neg_1 = arith.constant -1 : index
%0 = vector.insert %value_to_store, %dest[%c_neg_1] : vector<5xf32> into vector<4x5xf32>
%1 = vector.extract %src[%c_neg_1, 0] : f32 from vector<4x5xf32>
```
In addition to adding tests for the case above, this PR also relocates
the tests from #135498 to be alongside existing tests for the
`vector.{insert|extract}` folder, and reformats them to follow:
* https://mlir.llvm.org/getting_started/TestingGuide/
For example:
* The "no_fold" prefix is now used to label negative tests.
* Redundant check lines have been removed (e.g., CHECK: vector.insert
is sufficient to verify that folding did not occur).
(*) As per https://mlir.llvm.org/docs/Dialects/Vector/#vectorinsert-vectorinsertop,
these are poison values.
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 3fee1e949aeed..df56c68c65cf1 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -2045,8 +2045,9 @@ static Value extractInsertFoldConstantOp(OpType op, AdaptorType adaptor,
Value position = dynamicPosition[index++];
if (auto attr = mlir::dyn_cast_if_present<IntegerAttr>(positionAttr)) {
int64_t value = attr.getInt();
- // Do not fold if the value is out of bounds.
- if (value >= 0 && value < vectorShape[i]) {
+ // Do not fold if the value is out of bounds (-1 signifies a poison
+ // value rather than OOB index).
+ if (value >= -1 && value < vectorShape[i]) {
staticPosition[i] = attr.getInt();
opChange = true;
continue;
diff --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 2d365ac2b4287..29a11f47481c8 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -165,6 +165,33 @@ func.func @extract_scalar_poison_idx(%a: vector<4x5xf32>) -> f32 {
return %0 : f32
}
+// -----
+
+// Similar to the test above, but the index is not a static constant.
+
+// CHECK-LABEL: @extract_scalar_poison_idx_non_cst
+func.func @extract_scalar_poison_idx_non_cst(%a: vector<4x5xf32>) -> f32 {
+ // CHECK-NEXT: %[[UB:.*]] = ub.poison : f32
+ // CHECK-NOT: vector.extract
+ // CHECK-NEXT: return %[[UB]] : f32
+ %c_neg_1 = arith.constant -1 : index
+ %0 = vector.extract %a[%c_neg_1, 0] : f32 from vector<4x5xf32>
+ return %0 : f32
+}
+
+// -----
+
+// Similar to test above, but now the index is out-of-bounds.
+
+// CHECK-LABEL: @no_fold_extract_scalar_oob_idx
+func.func @no_fold_extract_scalar_oob_idx(%a: vector<4x5xf32>) -> f32 {
+ // CHECK: vector.extract
+ %c_neg_2 = arith.constant -2 : index
+ %0 = vector.extract %a[%c_neg_2, 0] : f32 from vector<4x5xf32>
+ return %0 : f32
+}
+
+
// -----
// CHECK-LABEL: @extract_vector_poison_idx
@@ -3062,6 +3089,34 @@ func.func @insert_vector_poison_idx(%a: vector<4x5xf32>, %b: vector<5xf32>)
// -----
+// Similar to the test above, but the index is not a static constant.
+
+// CHECK-LABEL: @insert_vector_poison_idx_non_cst
+func.func @insert_vector_poison_idx_non_cst(%a: vector<4x5xf32>, %b: vector<5xf32>)
+ -> vector<4x5xf32> {
+ // CHECK-NEXT: %[[UB:.*]] = ub.poison : vector<4x5xf32>
+ // CHECK-NOT: vector.insert
+ // CHECK-NEXT: return %[[UB]] : vector<4x5xf32>
+ %c_neg_1 = arith.constant -1 : index
+ %0 = vector.insert %b, %a[%c_neg_1] : vector<5xf32> into vector<4x5xf32>
+ return %0 : vector<4x5xf32>
+}
+
+// -----
+
+// Similar to test above, but now the index is out-of-bounds.
+
+// CHECK-LABEL: @no_fold_insert_scalar_idx_oob
+func.func @no_fold_insert_scalar_idx_oob(%a: vector<4x5xf32>, %b: vector<5xf32>)
+ -> vector<4x5xf32> {
+ // CHECK: vector.insert
+ %c_neg_2 = arith.constant -2 : index
+ %0 = vector.insert %b, %a[%c_neg_2] : vector<5xf32> into vector<4x5xf32>
+ return %0 : vector<4x5xf32>
+}
+
+// -----
+
// CHECK-LABEL: @insert_multiple_poison_idx
func.func @insert_multiple_poison_idx(%a: vector<4x5x8xf32>, %b: vector<8xf32>)
-> vector<4x5x8xf32> {
@@ -3311,41 +3366,3 @@ func.func @fold_insert_constant_indices(%arg : vector<4x1xi32>) -> vector<4x1xi3
%res = vector.insert %1, %arg[%0, %0] : i32 into vector<4x1xi32>
return %res : vector<4x1xi32>
}
-
-// -----
-
-// Check that out of bounds indices are not folded for vector.insert.
-
-// CHECK-LABEL: @fold_insert_oob
-// CHECK-SAME: %[[ARG:.*]]: vector<4x1x2xi32>) -> vector<4x1x2xi32> {
-// CHECK: %[[OOB1:.*]] = arith.constant -2 : index
-// CHECK: %[[OOB2:.*]] = arith.constant 2 : index
-// CHECK: %[[VAL:.*]] = arith.constant 1 : i32
-// CHECK: %[[RES:.*]] = vector.insert %[[VAL]], %[[ARG]] [0, %[[OOB1]], %[[OOB2]]] : i32 into vector<4x1x2xi32>
-// CHECK: return %[[RES]] : vector<4x1x2xi32>
-func.func @fold_insert_oob(%arg : vector<4x1x2xi32>) -> vector<4x1x2xi32> {
- %c0 = arith.constant 0 : index
- %c-2 = arith.constant -2 : index
- %c2 = arith.constant 2 : index
- %c1 = arith.constant 1 : i32
- %res = vector.insert %c1, %arg[%c0, %c-2, %c2] : i32 into vector<4x1x2xi32>
- return %res : vector<4x1x2xi32>
-}
-
-// -----
-
-// Check that out of bounds indices are not folded for vector.extract.
-
-// CHECK-LABEL: @fold_extract_oob
-// CHECK-SAME: %[[ARG:.*]]: vector<4x1x2xi32>) -> i32 {
-// CHECK: %[[OOB1:.*]] = arith.constant -2 : index
-// CHECK: %[[OOB2:.*]] = arith.constant 2 : index
-// CHECK: %[[RES:.*]] = vector.extract %[[ARG]][0, %[[OOB1]], %[[OOB2]]] : i32 from vector<4x1x2xi32>
-// CHECK: return %[[RES]] : i32
-func.func @fold_extract_oob(%arg : vector<4x1x2xi32>) -> i32 {
- %c0 = arith.constant 0 : index
- %c-2 = arith.constant -2 : index
- %c2 = arith.constant 2 : index
- %res = vector.extract %arg[%c0, %c-2, %c2] : i32 from vector<4x1x2xi32>
- return %res : i32
-}
More information about the Mlir-commits
mailing list