This is an archive of the discontinued LLVM Phabricator instance.

Differential D128545

[mlir][vector] Relax transfer_write vector distribution pattern
ClosedPublic

Authored by ThomasRaoux on Jun 24 2022, 11:42 AM.

Details

Reviewers
nicolasvasilache
aartbik
Commits
rG7eba5cdf9ce4: [mlir][vector] Relax transfer_write vector distribution pattern
Summary

Small change to relax the pattern to support any vector containing a
single element.

Diff Detail

Event Timeline

ThomasRaoux created this revision.Jun 24 2022, 11:42 AM
Herald added a reviewer: aartbik. · View Herald TranscriptJun 24 2022, 11:42 AM
Herald added a project: Restricted Project. · View Herald Transcript
Herald added subscribers: bzcheeseman, sdasgup3, wenzhicui and 19 others. · View Herald Transcript
ThomasRaoux requested review of this revision.Jun 24 2022, 11:42 AM
Herald added a project: Restricted Project. · View Herald TranscriptJun 24 2022, 11:42 AM
Herald added a subscriber: stephenneuendorffer. · View Herald Transcript
nicolasvasilache accepted this revision.Jun 24 2022, 11:51 AM
This revision is now accepted and ready to land.Jun 24 2022, 11:51 AM
This revision was landed with ongoing or failed builds.Jun 24 2022, 12:03 PM
Closed by commit rG7eba5cdf9ce4: [mlir][vector] Relax transfer_write vector distribution pattern (authored by ThomasRaoux). · Explain Why
This revision was automatically updated to reflect the committed changes.
ThomasRaoux added a commit: rG7eba5cdf9ce4: [mlir][vector] Relax transfer_write vector distribution pattern.
Harbormaster completed remote builds in B171903: Diff 439829.Jun 24 2022, 12:07 PM

Revision Contents

PathSize
mlir/
lib/
Dialect/
Vector/
Transforms/
5 lines
test/
Dialect/
Vector/
19 lines

Diff 439829

mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp

Show First 20 Lines • Show All 340 Lines • ▼ Show 20 Linesstruct WarpOpTransferWrite : public OpRewritePattern<vector::TransferWriteOp> {
/// Extract TransferWriteOps of vector<1x> into a separate warp op. /// Extract TransferWriteOps of vector<1x> into a separate warp op.
LogicalResult tryExtractOp(RewriterBase &rewriter, LogicalResult tryExtractOp(RewriterBase &rewriter,
vector::TransferWriteOp writeOp, vector::TransferWriteOp writeOp,
WarpExecuteOnLane0Op warpOp) const { WarpExecuteOnLane0Op warpOp) const {
Location loc = writeOp.getLoc(); Location loc = writeOp.getLoc();
VectorType vecType = writeOp.getVectorType(); VectorType vecType = writeOp.getVectorType();
// Only vector<1x> is supported at the moment. // Only sink out vector of 1 element for now to not serialize large vector
if (vecType.getShape().size() != 1 || vecType.getShape()[0] != 1) // store. This can later be controlled by user.
if (vecType.getNumElements() != 1)
return failure(); return failure();
// Do not process warp ops that contain only TransferWriteOps. // Do not process warp ops that contain only TransferWriteOps.
if (llvm::all_of(warpOp.getOps(), [](Operation &op) { if (llvm::all_of(warpOp.getOps(), [](Operation &op) {
return isa<vector::TransferWriteOp, vector::YieldOp>(&op); return isa<vector::TransferWriteOp, vector::YieldOp>(&op);
})) }))
return failure(); return failure();
▲ Show 20 Lines • Show All 501 Lines • Show Last 20 Lines

mlir/test/Dialect/Vector/vector-warp-distribute.mlir

Show First 20 Lines • Show All 103 Lines • ▼ Show 20 Lines vector.warp_execute_on_lane_0(%laneid)[32] {
vector.transfer_write %7, %sc[%c32] : vector<64xf32>, memref<128xf32, #map0> vector.transfer_write %7, %sc[%c32] : vector<64xf32>, memref<128xf32, #map0>
} }
return return
} }
// ----- // -----
// CHECK-D-LABEL: func @warp_extract( // CHECK-D-LABEL: func @warp_extract(
// CHECK-D: %[[WARPOP:.*]] = vector.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<1xf32>) // CHECK-D: %[[WARPOP:.*]]:2 = vector.warp_execute_on_lane_0(%{{.*}})[32] -> (vector<1xf32>, vector<1x1xf32>)
// CHECK-D: "test.dummy_op" // CHECK-D: "test.dummy_op"
// CHECK-D: vector.yield %{{.*}} : vector<1xf32> // CHECK-D: "test.dummy_op"
// CHECK-D: vector.yield %{{.*}}, %{{.*}} : vector<1xf32>, vector<1x1xf32>
// CHECK-D: } // CHECK-D: }
// CHECK-D: vector.warp_execute_on_lane_0(%{{.*}})[32] { // CHECK-D: vector.warp_execute_on_lane_0(%{{.*}})[32] {
// CHECK-D: vector.transfer_write %[[WARPOP]], %{{.*}}[%{{.*}}] {{.*}} : vector<1xf32> // CHECK-D: vector.transfer_write %[[WARPOP]]#1, %{{.*}}[%{{.*}}] {{.*}} : vector<1x1xf32>
// CHECK-D: }
// CHECK-D: vector.warp_execute_on_lane_0(%{{.*}})[32] {
// CHECK-D: vector.transfer_write %[[WARPOP]]#0, %{{.*}}[%{{.*}}] {{.*}} : vector<1xf32>
// CHECK-D: } // CHECK-D: }
#map2 = affine_map<(d0)[s0] -> (d0 + s0)> func.func @warp_extract(%laneid: index, %arg1: memref<1024x1024xf32>, %gid : index) {
func.func @warp_extract(%laneid: index, %arg1: memref<1024xf32>, %gid : index) {
vector.warp_execute_on_lane_0(%laneid)[32] { vector.warp_execute_on_lane_0(%laneid)[32] {
%sa = memref.subview %arg1[%gid] [128] [1] : memref<1024xf32> to memref<128xf32, #map2>
%c0 = arith.constant 0 : index %c0 = arith.constant 0 : index
%v = "test.dummy_op"() : () -> (vector<1xf32>) %v = "test.dummy_op"() : () -> (vector<1xf32>)
vector.transfer_write %v, %sa[%c0] : vector<1xf32>, memref<128xf32, #map2> %v1 = "test.dummy_op"() : () -> (vector<1x1xf32>)
vector.transfer_write %v1, %arg1[%c0, %c0] : vector<1x1xf32>, memref<1024x1024xf32>
vector.transfer_write %v, %arg1[%c0, %c0] : vector<1xf32>, memref<1024x1024xf32>
} }
return return
} }
// ----- // -----
// CHECK-PROP-LABEL: func @warp_dead_result( // CHECK-PROP-LABEL: func @warp_dead_result(
func.func @warp_dead_result(%laneid: index) -> (vector<1xf32>) { func.func @warp_dead_result(%laneid: index) -> (vector<1xf32>) {
▲ Show 20 Lines • Show All 335 Lines • Show Last 20 Lines