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

update custom op conversions #92

Merged
merged 2 commits into from
Aug 30, 2024

Conversation

zezhang
Copy link
Contributor

@zezhang zezhang commented Aug 30, 2024

  • Fix a minor issue with broadcast op (fold the broadcast if axes attr is empty)
  • Add following ops to tcp.custom_op:
    --aten.sort
    --aten.cumsum
    --aten.min.dim
    --aten.view(dynamic shape only)
    --aten.topk

To test:
bazel test //... (in docker)

@zezhang zezhang requested a review from navahgar August 30, 2024 03:21
Copy link
Collaborator

@navahgar navahgar left a comment

Choose a reason for hiding this comment

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

LGTM. Some minor comments inline.

lib/Conversion/TorchToTcp/Misc.cpp Show resolved Hide resolved
lib/Conversion/TorchToTcp/TcpCustomOp.cpp Outdated Show resolved Hide resolved
lib/Conversion/TorchToTcp/TcpCustomOp.cpp Outdated Show resolved Hide resolved
lib/Conversion/TorchToTcp/TcpCustomOp.cpp Show resolved Hide resolved
@@ -211,6 +214,118 @@ class ConvertAtenFakeQuantizePerChannelAffineOp
}
};

class ConvertAtenSortOp : public OpConversionPattern<AtenSortOp> {
Copy link
Contributor

Choose a reason for hiding this comment

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

Most likely the sort is coming from a top-k. I would prefer that we keep top-k instead of changing it to a sort as trt supports top-k already, and if we add sort then it just ends up getting changed back to top-k internally with additional performance penalties due to any extra ops that get added

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Good idea. We can add an internal patch here to disable the topk-->sort. I will add a new custom_op to handle the topk then.

@zezhang zezhang merged commit 57d5e00 into cruise-automation:main Aug 30, 2024
1 check passed
@zezhang zezhang deleted the zezhang/update_custom_op branch August 30, 2024 19:56
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.

3 participants