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

Handle tcp.bind_symbolic_shape ops in fusion algorithm #82

Merged
merged 4 commits into from
Jul 15, 2024
Merged

Handle tcp.bind_symbolic_shape ops in fusion algorithm #82

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
f89227f
Handle `tcp.bind_symbolic_shape` ops in fusion algorithm
Jul 15, 2024
9c680d8
add lit test
Jul 15, 2024
5de76fe
typo
Jul 15, 2024
b060976
one more lit test
Jul 15, 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
17 changes: 17 additions & 0 deletions lib/Dialect/Transforms/FusionPatterns.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,8 +44,15 @@ GenericBottomUpFuser::matchAndRewrite(Operation *op,
Region *usesParentRegion = nullptr;
SmallVector<Operation *> uses;
llvm::DenseSet<Operation *> usesSet;
llvm::DenseSet<tcp::BindSymbolicShapeOp> bindShapeUses;

LLVM_DEBUG(llvm::dbgs() << "Processing op: " << *op << "\n");
for (auto &use : op->getUses()) {
if (auto bindShapeOp = dyn_cast<tcp::BindSymbolicShapeOp>(use.getOwner())) {
bindShapeUses.insert(bindShapeOp);
continue;
}

auto parentRegion = use.getOwner()->getParentRegion();
if (usesParentRegion && usesParentRegion != parentRegion)
return failure();
Expand All @@ -58,6 +65,10 @@ GenericBottomUpFuser::matchAndRewrite(Operation *op,
uses.push_back(use.getOwner());
}

// All its uses are tcp.bind_symbolic_shape ops.
if (uses.empty())
return failure();

// Sorting by dominance ensures that the first element of this vector is
// the first use of the def. Used below when we want to move the op into
// an existing group.
Expand Down Expand Up @@ -104,6 +115,9 @@ GenericBottomUpFuser::matchAndRewrite(Operation *op,
use->moveBefore(yieldOp);
}
op->moveBefore(*uses.begin());
for (auto bindShapeOp : bindShapeUses) {
bindShapeOp->moveAfter(op);
}
}

// We then replace all uses of the uses which lie outside the group
Expand Down Expand Up @@ -157,6 +171,9 @@ GenericBottomUpFuser::matchAndRewrite(Operation *op,
// reordered.
auto &firstOp = *usesParentRegion->getOps().begin();
op->moveBefore(&firstOp);
for (auto bindShapeOp : bindShapeUses) {
bindShapeOp->moveBefore(&firstOp);
}
} else {
op->emitError("Unhandled case during fusion");
llvm_unreachable("Unhandled case during fusion");
Expand Down
70 changes: 70 additions & 0 deletions test/Dialect/tcp_fusion.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,3 +97,73 @@ func.func @test_multi_use_fusion(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32
%3 = tcp.mul %1, %2 : tensor<?x?xf32>, tensor<?x?xf32> -> tensor<?x?xf32>
"func.return" (%2, %3) : (tensor<?x?xf32>, tensor<?x?xf32>) -> ()
}

// -----

// CHECK: func.func @test_fusion_with_symbolic_shape(%[[ARG0:.+]]: tensor<?x?xf32>, %[[ARG1:.+]]: tensor<?x?xf32>) -> tensor<?x?xf32> {
// CHECK: %[[V0:.+]] = tcp.symbolic_int "s0" {min_val = 2, max_val = 9223372036854775806} : i64
// CHECK: %[[V1:.+]] = tcp.symbolic_int "s1" {min_val = 2, max_val = 9223372036854775806} : i64
// CHECK: tcp.bind_symbolic_shape %[[ARG0]], [%[[V0]], %[[V1]]], affine_map<()[s0, s1] -> (s0, s1)> : tensor<?x?xf32>
// CHECK: %[[V2:.+]] = tcp.group {
// CHECK: %[[V3:.+]] = tcp.tanh %[[ARG0]] : tensor<?x?xf32> -> tensor<?x?xf32>
// CHECK: tcp.bind_symbolic_shape %[[V3]], [%[[V0]], %[[V1]]], affine_map<()[s0, s1] -> (s0, s1)> : tensor<?x?xf32>
// CHECK: %[[V4:.+]] = tcp.tanh %[[V3]] : tensor<?x?xf32> -> tensor<?x?xf32>
// CHECK: tcp.bind_symbolic_shape %[[V4]], [%[[V0]], %[[V1]]], affine_map<()[s0, s1] -> (s0, s1)> : tensor<?x?xf32>
// CHECK: %[[V5:.+]] = tcp.add %[[V4]], %[[V4]] : tensor<?x?xf32>, tensor<?x?xf32> -> tensor<?x?xf32>
// CHECK: tcp.yield %[[V5]] : tensor<?x?xf32>
// CHECK: } : tensor<?x?xf32>
// CHECK: tcp.bind_symbolic_shape %[[V2]], [%[[V0]], %[[V1]]], affine_map<()[s0, s1] -> (s0, s1)> : tensor<?x?xf32>
Copy link
Collaborator

Choose a reason for hiding this comment

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

Does it work the same way when the tcp.group returns more than one output? i.e. the tcp.bind_symbolic_shape ops are outside the groups?

// CHECK: return %[[V2]] : tensor<?x?xf32>
// CHECK: }
func.func @test_fusion_with_symbolic_shape(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>) -> tensor<?x?xf32> {
%s0 = tcp.symbolic_int "s0" {min_val = 2, max_val = 9223372036854775806} : i64
%s1 = tcp.symbolic_int "s1" {min_val = 2, max_val = 9223372036854775806} : i64

tcp.bind_symbolic_shape %arg0, [%s0, %s1], affine_map<()[s0, s1] -> (s0, s1)> : tensor<?x?xf32>

%0 = tcp.tanh %arg0 : tensor<?x?xf32> -> tensor<?x?xf32>
tcp.bind_symbolic_shape %0, [%s0, %s1], affine_map<()[s0, s1] -> (s0, s1)> : tensor<?x?xf32>

%1 = tcp.tanh %0 : tensor<?x?xf32> -> tensor<?x?xf32>
tcp.bind_symbolic_shape %1, [%s0, %s1], affine_map<()[s0, s1] -> (s0, s1)> : tensor<?x?xf32>

%2 = tcp.add %1, %1 : tensor<?x?xf32>, tensor<?x?xf32> -> tensor<?x?xf32>
tcp.bind_symbolic_shape %2, [%s0, %s1], affine_map<()[s0, s1] -> (s0, s1)> : tensor<?x?xf32>

return %2 : tensor<?x?xf32>
}

// -----

// CHECK: func.func @test_multi_use_fusion_with_sym_shapes(%[[ARG0:.+]]: tensor<?x?xf32>, %[[ARG1:.+]]: tensor<?x?xf32>) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
// CHECK: %[[V0:.+]] = tcp.symbolic_int "s0" {min_val = 2, max_val = 9223372036854775806} : i64
// CHECK: %[[V1:.+]] = tcp.symbolic_int "s1" {min_val = 2, max_val = 9223372036854775806} : i64
// CHECK: tcp.bind_symbolic_shape %[[ARG0]], [%[[V0]], %[[V1]]], affine_map<()[s0, s1] -> (s0, s1)> : tensor<?x?xf32>
// CHECK: %[[V2:.+]]:2 = tcp.group {
// CHECK: %[[V3:.+]] = tcp.tanh %[[ARG0]] : tensor<?x?xf32> -> tensor<?x?xf32>
// CHECK: tcp.bind_symbolic_shape %[[V3]], [%[[V0]], %[[V1]]], affine_map<()[s0, s1] -> (s0, s1)> : tensor<?x?xf32>
// CHECK: %[[V4:.+]] = tcp.add %[[V3]], %[[V3]] : tensor<?x?xf32>, tensor<?x?xf32> -> tensor<?x?xf32>
// CHECK: tcp.bind_symbolic_shape %[[V4]], [%[[V0]], %[[V1]]], affine_map<()[s0, s1] -> (s0, s1)> : tensor<?x?xf32>
// CHECK: %[[V5:.+]] = tcp.sub %[[V4]], %[[ARG1]] : tensor<?x?xf32>, tensor<?x?xf32> -> tensor<?x?xf32>
// CHECK: %[[V6:.+]] = tcp.mul %[[V4]], %[[V5]] : tensor<?x?xf32>, tensor<?x?xf32> -> tensor<?x?xf32>
// CHECK: tcp.yield %[[V5]], %[[V6]] : tensor<?x?xf32>, tensor<?x?xf32>
// CHECK: } : tensor<?x?xf32>, tensor<?x?xf32>
// CHECK: tcp.bind_symbolic_shape %[[V2]]#0, [%[[V0]], %[[V1]]], affine_map<()[s0, s1] -> (s0, s1)> : tensor<?x?xf32>
// CHECK: tcp.bind_symbolic_shape %[[V2]]#1, [%[[V0]], %[[V1]]], affine_map<()[s0, s1] -> (s0, s1)> : tensor<?x?xf32>
// CHECK: return %[[V2]]#0, %[[V2]]#1 : tensor<?x?xf32>, tensor<?x?xf32>
// CHECK: }
func.func @test_multi_use_fusion_with_sym_shapes(%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>) -> (tensor<?x?xf32>, tensor<?x?xf32>) {
%s0 = tcp.symbolic_int "s0" {min_val = 2, max_val = 9223372036854775806} : i64
%s1 = tcp.symbolic_int "s1" {min_val = 2, max_val = 9223372036854775806} : i64
tcp.bind_symbolic_shape %arg0, [%s0, %s1], affine_map<()[s0, s1] -> (s0, s1)> : tensor<?x?xf32>

%0 = tcp.tanh %arg0 : tensor<?x?xf32> -> tensor<?x?xf32>
tcp.bind_symbolic_shape %0, [%s0, %s1], affine_map<()[s0, s1] -> (s0, s1)> : tensor<?x?xf32>
%1 = tcp.add %0, %0 : tensor<?x?xf32>, tensor<?x?xf32> -> tensor<?x?xf32>
tcp.bind_symbolic_shape %1, [%s0, %s1], affine_map<()[s0, s1] -> (s0, s1)> : tensor<?x?xf32>
%2 = tcp.sub %1, %arg1 : tensor<?x?xf32>, tensor<?x?xf32> -> tensor<?x?xf32>
tcp.bind_symbolic_shape %2, [%s0, %s1], affine_map<()[s0, s1] -> (s0, s1)> : tensor<?x?xf32>
%3 = tcp.mul %1, %2 : tensor<?x?xf32>, tensor<?x?xf32> -> tensor<?x?xf32>
tcp.bind_symbolic_shape %3, [%s0, %s1], affine_map<()[s0, s1] -> (s0, s1)> : tensor<?x?xf32>
"func.return" (%2, %3) : (tensor<?x?xf32>, tensor<?x?xf32>) -> ()
}