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.

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

Conversation

srinathava
Copy link
Contributor

The current fusion algorithm does not handle tcp.bind_symbolic_shape ops very well because these ops cause an op to have multiple downstream uses. This change handles these ops "specially". We collect all the uses of an op which only specify the shape of the output into a separate use category. We then move the bind shape ops along with the original op if it was moved to another region.

@sjain-stanford sjain-stanford self-requested a review July 15, 2024 19:34
Copy link
Collaborator

@sjain-stanford sjain-stanford left a comment

Choose a reason for hiding this comment

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

Could you please add some lit tests?

@@ -58,6 +65,10 @@ GenericBottomUpFuser::matchAndRewrite(Operation *op,
uses.push_back(use.getOwner());
}

// All its uses are tcp.bind_symbolic_shpae ops.
Copy link
Collaborator

Choose a reason for hiding this comment

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

nit: typo

Srinath Avadhanula added 2 commits July 15, 2024 12:40
// 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?

Copy link
Collaborator

@sjain-stanford sjain-stanford left a comment

Choose a reason for hiding this comment

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

Thanks! LGTM besides one question.

@srinathava srinathava merged commit a4fba88 into cruise-automation:main Jul 15, 2024
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Development

Successfully merging this pull request may close these issues.

2 participants