[Mlir-commits] [mlir] 4e2efea - [mlir][vector] Add all view-like ops to transfer flow opt (#110521)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Oct 1 21:20:47 PDT 2024
Author: Quinn Dawkins
Date: 2024-10-02T00:20:44-04:00
New Revision: 4e2efea5e8e55b26dd7ac90c6cd1ab7bf6775650
URL: https://github.com/llvm/llvm-project/commit/4e2efea5e8e55b26dd7ac90c6cd1ab7bf6775650
DIFF: https://github.com/llvm/llvm-project/commit/4e2efea5e8e55b26dd7ac90c6cd1ab7bf6775650.diff
LOG: [mlir][vector] Add all view-like ops to transfer flow opt (#110521)
`vector.transfer_*` folding and forwarding currently does not take into
account reshaping view-like memref ops (expand and collapse shape),
leading to potentially invalid store folding or value forwarding. This
patch adds tracking for those (and other) view-like ops. It is still
possible to design operations that alias memrefs without being a view
(e.g. memref in the iter_args of an `scf.for`), so these patterns may
still need revisiting in the future.
Added:
Modified:
mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h
mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp
mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
mlir/test/Dialect/Vector/vector-transferop-opt.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h b/mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h
index 1c094c1c1328a9..ca3326dbbef519 100644
--- a/mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h
+++ b/mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h
@@ -106,9 +106,9 @@ inline bool isSameViewOrTrivialAlias(MemrefValue a, MemrefValue b) {
return skipFullyAliasingOperations(a) == skipFullyAliasingOperations(b);
}
-/// Walk up the source chain until something an op other than a `memref.subview`
-/// or `memref.cast` is found.
-MemrefValue skipSubViewsAndCasts(MemrefValue source);
+/// Walk up the source chain until we find an operation that is not a view of
+/// the source memref (i.e. implements ViewLikeOpInterface).
+MemrefValue skipViewLikeOps(MemrefValue source);
} // namespace memref
} // namespace mlir
diff --git a/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp b/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp
index 68edd45448ee5f..7321b19068016c 100644
--- a/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp
+++ b/mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp
@@ -15,6 +15,7 @@
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
+#include "mlir/Interfaces/ViewLikeInterface.h"
#include "llvm/ADT/STLExtras.h"
namespace mlir {
@@ -193,15 +194,13 @@ MemrefValue skipFullyAliasingOperations(MemrefValue source) {
return source;
}
-MemrefValue skipSubViewsAndCasts(MemrefValue source) {
+MemrefValue skipViewLikeOps(MemrefValue source) {
while (auto op = source.getDefiningOp()) {
- if (auto subView = dyn_cast<memref::SubViewOp>(op)) {
- source = cast<MemrefValue>(subView.getSource());
- } else if (auto cast = dyn_cast<memref::CastOp>(op)) {
- source = cast.getSource();
- } else {
- return source;
+ if (auto viewLike = dyn_cast<ViewLikeOpInterface>(op)) {
+ source = cast<MemrefValue>(viewLike.getViewSource());
+ continue;
}
+ return source;
}
return source;
}
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
index 4c93d3841bf878..e05c801121ffc4 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
@@ -105,8 +105,7 @@ void TransferOptimization::deadStoreOp(vector::TransferWriteOp write) {
<< "\n");
llvm::SmallVector<Operation *, 8> blockingAccesses;
Operation *firstOverwriteCandidate = nullptr;
- Value source =
- memref::skipSubViewsAndCasts(cast<MemrefValue>(write.getSource()));
+ Value source = memref::skipViewLikeOps(cast<MemrefValue>(write.getSource()));
llvm::SmallVector<Operation *, 32> users(source.getUsers().begin(),
source.getUsers().end());
llvm::SmallDenseSet<Operation *, 32> processed;
@@ -115,7 +114,7 @@ void TransferOptimization::deadStoreOp(vector::TransferWriteOp write) {
// If the user has already been processed skip.
if (!processed.insert(user).second)
continue;
- if (isa<memref::SubViewOp, memref::CastOp>(user)) {
+ if (isa<ViewLikeOpInterface>(user)) {
users.append(user->getUsers().begin(), user->getUsers().end());
continue;
}
@@ -192,8 +191,7 @@ void TransferOptimization::storeToLoadForwarding(vector::TransferReadOp read) {
<< "\n");
SmallVector<Operation *, 8> blockingWrites;
vector::TransferWriteOp lastwrite = nullptr;
- Value source =
- memref::skipSubViewsAndCasts(cast<MemrefValue>(read.getSource()));
+ Value source = memref::skipViewLikeOps(cast<MemrefValue>(read.getSource()));
llvm::SmallVector<Operation *, 32> users(source.getUsers().begin(),
source.getUsers().end());
llvm::SmallDenseSet<Operation *, 32> processed;
@@ -202,7 +200,7 @@ void TransferOptimization::storeToLoadForwarding(vector::TransferReadOp read) {
// If the user has already been processed skip.
if (!processed.insert(user).second)
continue;
- if (isa<memref::SubViewOp, memref::CollapseShapeOp, memref::CastOp>(user)) {
+ if (isa<ViewLikeOpInterface>(user)) {
users.append(user->getUsers().begin(), user->getUsers().end());
continue;
}
diff --git a/mlir/test/Dialect/Vector/vector-transferop-opt.mlir b/mlir/test/Dialect/Vector/vector-transferop-opt.mlir
index 07e6647533c6fe..f4f7fb1ba03045 100644
--- a/mlir/test/Dialect/Vector/vector-transferop-opt.mlir
+++ b/mlir/test/Dialect/Vector/vector-transferop-opt.mlir
@@ -229,7 +229,7 @@ func.func @forward_dead_store_negative(%arg0: i1, %arg1 : memref<4x4xf32>,
// final `vector.transfer_write` should be preserved as:
// vector.transfer_write %2, %subview
-// CHECK-LABEL: func.func @collapse_shape
+// CHECK-LABEL: func.func @collapse_shape_and_read_from_source
// CHECK: scf.for {{.*}} {
// CHECK: vector.transfer_read
// CHECK: vector.transfer_write
@@ -237,7 +237,7 @@ func.func @forward_dead_store_negative(%arg0: i1, %arg1 : memref<4x4xf32>,
// CHECK: vector.transfer_read
// CHECK: vector.transfer_write
-func.func @collapse_shape(%in_0: memref<1x20x1xi32>, %vec: vector<4xi32>) {
+func.func @collapse_shape_and_read_from_source(%in_0: memref<1x20x1xi32>, %vec: vector<4xi32>) {
%c0_i32 = arith.constant 0 : i32
%c0 = arith.constant 0 : index
%c4 = arith.constant 4 : index
@@ -257,6 +257,98 @@ func.func @collapse_shape(%in_0: memref<1x20x1xi32>, %vec: vector<4xi32>) {
return
}
+// The same regression test for expand_shape.
+
+// CHECK-LABEL: func.func @expand_shape_and_read_from_source
+// CHECK: scf.for {{.*}} {
+// CHECK: vector.transfer_read
+// CHECK: vector.transfer_write
+// CHECK: vector.transfer_write
+// CHECK: vector.transfer_read
+// CHECK: vector.transfer_write
+
+func.func @expand_shape_and_read_from_source(%in_0: memref<20xi32>, %vec: vector<1x4x1xi32>) {
+ %c0_i32 = arith.constant 0 : i32
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %c20 = arith.constant 20 : index
+
+ %alloca = memref.alloca() {alignment = 64 : i64} : memref<4xi32>
+ %expand_shape = memref.expand_shape %alloca [[0, 1, 2]] output_shape [1, 4, 1] : memref<4xi32> into memref<1x4x1xi32>
+ scf.for %arg0 = %c0 to %c20 step %c4 {
+ %subview = memref.subview %in_0[%arg0] [4] [1] : memref<20xi32> to memref<4xi32, strided<[1], offset: ?>>
+ %1 = vector.transfer_read %subview[%c0], %c0_i32 {in_bounds = [true]} : memref<4xi32, strided<[1], offset: ?>>, vector<4xi32>
+ // $alloca and $expand_shape alias
+ vector.transfer_write %1, %alloca[%c0] {in_bounds = [true]} : vector<4xi32>, memref<4xi32>
+ vector.transfer_write %vec, %expand_shape[%c0, %c0, %c0] {in_bounds = [true, true, true]} : vector<1x4x1xi32>, memref<1x4x1xi32>
+ %2 = vector.transfer_read %alloca[%c0], %c0_i32 {in_bounds = [true]} : memref<4xi32>, vector<4xi32>
+ vector.transfer_write %2, %subview[%c0] {in_bounds = [true]} : vector<4xi32>, memref<4xi32, strided<[1], offset: ?>>
+ }
+ return
+}
+
+// The same regression test, but the initial write is to the collapsed memref,
+// and the subsequent unforwardable read is from the collapse shape.
+
+// CHECK-LABEL: func.func @collapse_shape_and_read_from_collapse
+// CHECK: scf.for {{.*}} {
+// CHECK: vector.transfer_read
+// CHECK: vector.transfer_write
+// CHECK: vector.transfer_write
+// CHECK: vector.transfer_read
+// CHECK: vector.transfer_write
+
+func.func @collapse_shape_and_read_from_collapse(%in_0: memref<20xi32>, %vec: vector<1x4x1xi32>) {
+ %c0_i32 = arith.constant 0 : i32
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %c20 = arith.constant 20 : index
+
+ %alloca = memref.alloca() {alignment = 64 : i64} : memref<1x4x1xi32>
+ %collapse_shape = memref.collapse_shape %alloca [[0, 1, 2]] : memref<1x4x1xi32> into memref<4xi32>
+ scf.for %arg0 = %c0 to %c20 step %c4 {
+ %subview = memref.subview %in_0[%arg0] [4] [1] : memref<20xi32> to memref<4xi32, strided<[1], offset: ?>>
+ %1 = vector.transfer_read %subview[%c0], %c0_i32 {in_bounds = [true]} : memref<4xi32, strided<[1], offset: ?>>, vector<4xi32>
+ vector.transfer_write %1, %collapse_shape[%c0] {in_bounds = [true]} : vector<4xi32>, memref<4xi32>
+ // $alloca and $collapse_shape alias
+ vector.transfer_write %vec, %alloca[%c0, %c0, %c0] {in_bounds = [true, true, true]} : vector<1x4x1xi32>, memref<1x4x1xi32>
+ %2 = vector.transfer_read %collapse_shape[%c0], %c0_i32 {in_bounds = [true]} : memref<4xi32>, vector<4xi32>
+ vector.transfer_write %2, %subview[%c0] {in_bounds = [true]} : vector<4xi32>, memref<4xi32, strided<[1], offset: ?>>
+ }
+ return
+}
+
+// The same test except writing to the expanded source first (same as the
+// previous collapse test but for expand).
+
+// CHECK-LABEL: func.func @expand_shape_and_read_from_expand
+// CHECK: scf.for {{.*}} {
+// CHECK: vector.transfer_read
+// CHECK: vector.transfer_write
+// CHECK: vector.transfer_write
+// CHECK: vector.transfer_read
+// CHECK: vector.transfer_write
+
+func.func @expand_shape_and_read_from_expand(%in_0: memref<1x20x1xi32>, %vec: vector<4xi32>) {
+ %c0_i32 = arith.constant 0 : i32
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %c20 = arith.constant 20 : index
+
+ %alloca = memref.alloca() {alignment = 64 : i64} : memref<4xi32>
+ %expand_shape = memref.expand_shape %alloca [[0, 1, 2]] output_shape [1, 4, 1] : memref<4xi32> into memref<1x4x1xi32>
+ scf.for %arg0 = %c0 to %c20 step %c4 {
+ %subview = memref.subview %in_0[0, %arg0, 0] [1, 4, 1] [1, 1, 1] : memref<1x20x1xi32> to memref<1x4x1xi32, strided<[20, 1, 1], offset: ?>>
+ %1 = vector.transfer_read %subview[%c0, %c0, %c0], %c0_i32 {in_bounds = [true, true, true]} : memref<1x4x1xi32, strided<[20, 1, 1], offset: ?>>, vector<1x4x1xi32>
+ vector.transfer_write %1, %expand_shape[%c0, %c0, %c0] {in_bounds = [true, true, true]} : vector<1x4x1xi32>, memref<1x4x1xi32>
+ // $alloca and $expand_shape alias
+ vector.transfer_write %vec, %alloca[%c0] {in_bounds = [true]} : vector<4xi32>, memref<4xi32>
+ %2 = vector.transfer_read %expand_shape[%c0, %c0, %c0], %c0_i32 {in_bounds = [true, true, true]} : memref<1x4x1xi32>, vector<1x4x1xi32>
+ vector.transfer_write %2, %subview[%c0, %c0, %c0] {in_bounds = [true, true, true]} : vector<1x4x1xi32>, memref<1x4x1xi32, strided<[20, 1, 1], offset: ?>>
+ }
+ return
+}
+
// CHECK-LABEL: func @forward_dead_store_dynamic_same_index
// CHECK-NOT: vector.transfer_write
// CHECK-NOT: vector.transfer_read
More information about the Mlir-commits
mailing list