Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[GPU] Bail out in GPUReduceBankConflicts if we have collapse_shape user #18863

Merged
merged 4 commits into from
Oct 24, 2024
Merged

[GPU] Bail out in GPUReduceBankConflicts if we have collapse_shape user #18863

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
5c82292
[GPU] Bail out in GPUReduceBankConflicts Pass if we have collapse_sha…
nirvedhmeshram Oct 21, 2024
4890f77
Address reviewer comments
nirvedhmeshram Oct 21, 2024
379040c
add comment about utility function
nirvedhmeshram Oct 21, 2024
c9141bb
use constructor
nirvedhmeshram Oct 21, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions compiler/src/iree/compiler/Codegen/Common/GPU/GPUReduceBankConflicts.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,23 @@ namespace mlir::iree_compiler {

namespace {

/// Check if AllocOp has a CollapseShapeOp user.
static bool hasCollapseShapeUser(memref::AllocOp allocOp) {
SmallVector<Operation *> users(allocOp->getUsers());
while (!users.empty()) {
auto user = users.pop_back_val();
if (isa<memref::CollapseShapeOp>(user)) {
return true;
}
if (isa<ViewLikeOpInterface>(user)) {
for (auto u : user->getUsers()) {
Comment on lines +29 to +30
Copy link
Contributor

@Max191 Max191 Oct 22, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should probably add arith::SelectOp to this too. And ops with a block that has an argument tied to this use should add the block arg users (I think we mainly need to worry about scf.for).

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think both of these cases should be okay with the pass, as it will add a strided subview with the same shape as the pre padded alloc and replace the type in all the users, the problem with the collapse shape case is it does not support strides that dont add up to the shape.

If there is some IR case you are worried about, happy to add that as a test.

users.push_back(u);
}
}
}
return false;
}

/// Pad out the inner dimension of the `memref.alloc` op in order reduce the
/// chances to have bank conflicts when reading 2D shapes within shared memory.
static void padAlloc(MLIRContext *context, memref::AllocOp allocOp,
Expand All @@ -28,6 +45,12 @@ static void padAlloc(MLIRContext *context, memref::AllocOp allocOp,
int64_t innerDim = allocOpShape.back();
if (ShapedType::isDynamic(innerDim))
return;

// Return if we have CollapseShape op as an user as padding in that case is
// unsupported.
if (hasCollapseShapeUser(allocOp))
return;

Type elType = allocOp.getType().getElementType();
unsigned bitwidth =
mlir::DataLayout::closest(allocOp).getTypeSizeInBits(elType);
Expand Down
60 changes: 60 additions & 0 deletions compiler/src/iree/compiler/Codegen/Common/GPU/test/reduce_bank_conflicts.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,66 @@ func.func @pad_alloc_expand_shape(%a: memref<1024x1024xf32>) {
return
}

// -----
// CHECK-LABEL: func.func @no_pad_alloc_collapse_shape
// CHECK: %[[A:.*]] = memref.alloc() : memref<4x2x16x8x8xf32, #gpu.address_space<workgroup>>
// CHECK: %[[C:.*]] = memref.collapse_shape %[[A]] {{\[}}[0], [1, 2], [3, 4]]
// CHECK-SAME: memref<4x2x16x8x8xf32, #gpu.address_space<workgroup>> into
// CHECK-SAME: memref<4x32x64xf32, #gpu.address_space<workgroup>>
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: %[[CST_0:.*]] = arith.constant 0.000000e+00 : f32
// CHECK: %[[VEC_READ:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]]], %[[CST_0]] {in_bounds = [true]} :
// CHECK-SAME: memref<1024x1024xf32>, vector<4xf32>
// CHECK: vector.transfer_write %[[VEC_READ]], %[[C]][%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true]} :
// CHECK-SAME: vector<4xf32>, memref<4x32x64xf32, #gpu.address_space<workgroup>>


func.func @no_pad_alloc_collapse_shape(%a: memref<1024x1024xf32>) {
%0 = memref.alloc() : memref<4x2x16x8x8xf32, #gpu.address_space<workgroup>>
%1 = memref.collapse_shape %0 [[0], [1, 2], [3, 4]]
: memref<4x2x16x8x8xf32, #gpu.address_space<workgroup>> into memref<4x32x64xf32, #gpu.address_space<workgroup>>
%c0 = arith.constant 0 : index
%cst_0 = arith.constant 0.000000e+00 : f32
%3 = vector.transfer_read %a[%c0, %c0], %cst_0 {in_bounds = [true]} :
memref<1024x1024xf32>, vector<4xf32>
vector.transfer_write %3, %1[%c0, %c0, %c0] {in_bounds = [true]} :
vector<4xf32>, memref<4x32x64xf32, #gpu.address_space<workgroup>>
return
}

// -----

// CHECK-LABEL: func.func @no_pad_alloc_collapse_shape_throughsubview
// CHECK: %[[A:.*]] = memref.alloc() : memref<4x2x16x8x8xf32, #gpu.address_space<workgroup>>
// CHECK: %[[S:.*]] = memref.subview %[[A]][0, 0, 0, 0, 0] [4, 2, 16, 8, 8] [1, 1, 1, 1, 1] :
// CHECK-SAME: memref<4x2x16x8x8xf32, #gpu.address_space<workgroup>> to
// CHECK-SAME: memref<4x2x16x8x8xf32, #gpu.address_space<workgroup>>
// CHECK: %[[C:.*]] = memref.collapse_shape %[[S]] {{\[}}[0], [1, 2], [3, 4]]
// CHECK-SAME: memref<4x2x16x8x8xf32, #gpu.address_space<workgroup>> into
// CHECK-SAME: memref<4x32x64xf32, #gpu.address_space<workgroup>>
// CHECK: %[[C0:.*]] = arith.constant 0 : index
// CHECK: %[[CST:.*]] = arith.constant 0.000000e+00 : f32
// CHECK: %[[VEC_READ:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]]], %[[CST]] {in_bounds = [true]} :
// CHECK-SAME: memref<1024x1024xf32>, vector<4xf32>
// CHECK: vector.transfer_write %[[VEC_READ]], %[[C]][%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true]} :
// CHECK-SAME: vector<4xf32>, memref<4x32x64xf32, #gpu.address_space<workgroup>>


func.func @no_pad_alloc_collapse_shape_throughsubview(%a: memref<1024x1024xf32>) {
%0 = memref.alloc() : memref<4x2x16x8x8xf32, #gpu.address_space<workgroup>>
%subview = memref.subview %0[0, 0, 0, 0, 0] [4, 2, 16, 8, 8] [1, 1, 1, 1, 1]
: memref<4x2x16x8x8xf32, #gpu.address_space<workgroup>> to memref<4x2x16x8x8xf32, #gpu.address_space<workgroup>>
%1 = memref.collapse_shape %subview [[0], [1, 2], [3, 4]]
: memref<4x2x16x8x8xf32, #gpu.address_space<workgroup>> into memref<4x32x64xf32, #gpu.address_space<workgroup>>
%c0 = arith.constant 0 : index
%cst_0 = arith.constant 0.000000e+00 : f32
%3 = vector.transfer_read %a[%c0, %c0], %cst_0 {in_bounds = [true]} :
memref<1024x1024xf32>, vector<4xf32>
vector.transfer_write %3, %1[%c0, %c0, %c0] {in_bounds = [true]} :
vector<4xf32>, memref<4x32x64xf32, #gpu.address_space<workgroup>>
return
}

// -----

// CHECK-LABEL: func.func @pad_alloc_negative
Expand Down
Loading