[Mlir-commits] [mlir] [mlir][vector] Restrict DropInnerMostUnitDimsTransferRead (PR #94904)
Andrzej Warzyński
llvmlistbot at llvm.org
Wed Jun 12 00:08:09 PDT 2024
https://github.com/banach-space updated https://github.com/llvm/llvm-project/pull/94904
>From 545c8cf77a3383f9ae9be89879c41b7f3a559286 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Sun, 9 Jun 2024 16:38:18 +0100
Subject: [PATCH 1/3] [mlir][vector] Restrict
`DropInnerMostUnitDimsTransferRead`
Restrict `DropInnerMostUnitDimsTransferRead` so that it fails when one
of the indices to be dropped could be != 0, e.g.
```
func.func @negative_example(%A: memref<16x1xf32>, %i:index, %j:index) -> (vector<8x1xf32>) {
%f0 = arith.constant 0.0 : f32
%1 = vector.transfer_read %A[%i, %j], %f0 : memref<16x1xf32>, vector<8x1xf32>
return %1 : vector<8x1xf32>
}
```
This is an edge case that could represent an out-of-bounds access,
though that will depend on the actual value of `%j`.
NOTE: This PR is limited to tests for `vector.transfer_read`.
Depends on: #94490, #94604
---
.../Vector/Transforms/VectorTransforms.cpp | 15 ++++++++++++
...tor-transfer-collapse-inner-most-dims.mlir | 24 +++++++++++++------
2 files changed, 32 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index f29eba90c3ceb..caf1506e0db93 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -1293,6 +1293,21 @@ class DropInnerMostUnitDimsTransferRead
if (dimsToDrop == 0)
return failure();
+ // Make sure that the indixes to be dropped are equal 0.
+ // TODO: Deal with cases when the indices are not 0.
+ auto isZeroIdx = [](Value idx) {
+ Attribute attr;
+ APInt value;
+ if (!matchPattern(idx, m_Constant(&attr)))
+ return false;
+ if (matchPattern(attr, m_ConstantInt(&value)))
+ if (!value.isZero())
+ return false;
+ return true;
+ };
+ if (!llvm::all_of(readOp.getIndices().take_back(dimsToDrop), isZeroIdx))
+ return failure();
+
auto resultTargetVecType =
VectorType::get(targetType.getShape().drop_back(dimsToDrop),
targetType.getElementType(),
diff --git a/mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir b/mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir
index a50c01898c62e..bb37d5b45520c 100644
--- a/mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir
+++ b/mlir/test/Dialect/Vector/vector-transfer-collapse-inner-most-dims.mlir
@@ -111,31 +111,41 @@ func.func @contiguous_inner_most_outer_dim_dyn_scalable_inner_dim(%a: index, %b:
// -----
-func.func @contiguous_inner_most_dim_non_zero_idxs(%A: memref<16x1xf32>, %i:index, %j:index) -> (vector<8x1xf32>) {
+func.func @contiguous_inner_most_dim_non_zero_idx(%A: memref<16x1xf32>, %i:index) -> (vector<8x1xf32>) {
%c0 = arith.constant 0 : index
%f0 = arith.constant 0.0 : f32
- %1 = vector.transfer_read %A[%i, %j], %f0 : memref<16x1xf32>, vector<8x1xf32>
+ %1 = vector.transfer_read %A[%i, %c0], %f0 : memref<16x1xf32>, vector<8x1xf32>
return %1 : vector<8x1xf32>
}
-// CHECK: func @contiguous_inner_most_dim_non_zero_idxs(%[[SRC:.+]]: memref<16x1xf32>, %[[I:.+]]: index, %[[J:.+]]: index) -> vector<8x1xf32>
+// CHECK: func @contiguous_inner_most_dim_non_zero_idx(%[[SRC:.+]]: memref<16x1xf32>, %[[I:.+]]: index) -> vector<8x1xf32>
// CHECK: %[[SRC_0:.+]] = memref.subview %[[SRC]]
// CHECK-SAME: memref<16x1xf32> to memref<16xf32, strided<[1]>>
// CHECK: %[[V:.+]] = vector.transfer_read %[[SRC_0]]
// CHECK: %[[RESULT:.+]] = vector.shape_cast %[[V]] : vector<8xf32> to vector<8x1xf32>
// CHECK: return %[[RESULT]]
+// The index to be dropped is != 0 - this is currently not supported.
+func.func @negative_contiguous_inner_most_dim_non_zero_idxs(%A: memref<16x1xf32>, %i:index) -> (vector<8x1xf32>) {
+ %f0 = arith.constant 0.0 : f32
+ %1 = vector.transfer_read %A[%i, %i], %f0 : memref<16x1xf32>, vector<8x1xf32>
+ return %1 : vector<8x1xf32>
+}
+// CHECK-LABEL: func @negative_contiguous_inner_most_dim_non_zero_idxs
+// CHECK-NOT: memref.subview
+// CHECK: vector.transfer_read
+
// Same as the top example within this split, but with the outer vector
// dim scalable. Note that this example only makes sense when "8 = [8]" (i.e.
// vscale = 1). This is assumed (implicitly) via the `in_bounds` attribute.
-func.func @contiguous_inner_most_dim_non_zero_idxs_scalable_inner_dim(%A: memref<16x1xf32>, %i:index, %j:index) -> (vector<[8]x1xf32>) {
+func.func @contiguous_inner_most_dim_non_zero_idx_scalable_inner_dim(%A: memref<16x1xf32>, %i:index) -> (vector<[8]x1xf32>) {
%c0 = arith.constant 0 : index
%f0 = arith.constant 0.0 : f32
- %1 = vector.transfer_read %A[%i, %j], %f0 : memref<16x1xf32>, vector<[8]x1xf32>
+ %1 = vector.transfer_read %A[%i, %c0], %f0 : memref<16x1xf32>, vector<[8]x1xf32>
return %1 : vector<[8]x1xf32>
}
-// CHECK-LABEL: func @contiguous_inner_most_dim_non_zero_idxs_scalable_inner_dim(
-// CHECK-SAME: %[[SRC:.+]]: memref<16x1xf32>, %[[I:.+]]: index, %[[J:.+]]: index) -> vector<[8]x1xf32>
+// CHECK-LABEL: func @contiguous_inner_most_dim_non_zero_idx_scalable_inner_dim(
+// CHECK-SAME: %[[SRC:.+]]: memref<16x1xf32>, %[[I:.+]]: index) -> vector<[8]x1xf32>
// CHECK: %[[SRC_0:.+]] = memref.subview %[[SRC]]
// CHECK-SAME: memref<16x1xf32> to memref<16xf32, strided<[1]>>
// CHECK: %[[V:.+]] = vector.transfer_read %[[SRC_0]]
>From a9651e1c427357a3abf8220c86201ef03f8b6029 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Tue, 11 Jun 2024 15:32:46 +0100
Subject: [PATCH 2/3] fixup! [mlir][vector] Restrict
`DropInnerMostUnitDimsTransferRead`
Switch to using isZeroIndex from StaticValueUtils.h
---
.../Dialect/Vector/Transforms/VectorTransforms.cpp | 12 +-----------
1 file changed, 1 insertion(+), 11 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index caf1506e0db93..ecb8e8060aed0 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -1295,17 +1295,7 @@ class DropInnerMostUnitDimsTransferRead
// Make sure that the indixes to be dropped are equal 0.
// TODO: Deal with cases when the indices are not 0.
- auto isZeroIdx = [](Value idx) {
- Attribute attr;
- APInt value;
- if (!matchPattern(idx, m_Constant(&attr)))
- return false;
- if (matchPattern(attr, m_ConstantInt(&value)))
- if (!value.isZero())
- return false;
- return true;
- };
- if (!llvm::all_of(readOp.getIndices().take_back(dimsToDrop), isZeroIdx))
+ if (!llvm::all_of(readOp.getIndices().take_back(dimsToDrop), isZeroIndex))
return failure();
auto resultTargetVecType =
>From c55adde64eec9087611bcf12da793584d69bce34 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Tue, 11 Jun 2024 19:10:04 +0100
Subject: [PATCH 3/3] fixup! fixup! [mlir][vector] Restrict
`DropInnerMostUnitDimsTransferRead`
Fix typo
---
mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index ecb8e8060aed0..ea4a02f2f2e77 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -1293,7 +1293,7 @@ class DropInnerMostUnitDimsTransferRead
if (dimsToDrop == 0)
return failure();
- // Make sure that the indixes to be dropped are equal 0.
+ // Make sure that the indices to be dropped are equal 0.
// TODO: Deal with cases when the indices are not 0.
if (!llvm::all_of(readOp.getIndices().take_back(dimsToDrop), isZeroIndex))
return failure();
More information about the Mlir-commits
mailing list