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

tcp.custom_op support for torch.aten.arange with dynamic input #100

Merged
merged 2 commits into from
Oct 7, 2024

Conversation

zezhang
Copy link
Contributor

@zezhang zezhang commented Oct 7, 2024

tcp.custom_op support for torch.aten.arange with dynamic input. Static case will be handled through TOSA dialect.

To test:

(in docker)
bazel test //...

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.

Mostly LGTM. One minor comment inline.

Comment on lines +544 to +550
// convert constant int to fp if possible
int64_t constIntVal;
if (matchPattern(value, torch::Torch::m_TorchConstantInt(&constIntVal))) {
attrs.push_back(rewriter.getNamedAttr(
attrName, rewriter.getF64FloatAttr(static_cast<double>(constIntVal))));
return true;
}
Copy link
Collaborator

Choose a reason for hiding this comment

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

Why do we have to convert int to fp here? Why not convert them to IntAttr?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

aten.arange.start_step can take either int or fp for start and step arguments. However, trt-mlir only accepts F64Attr. Converting all constants to FPAttr saves us some time for type checking.

I can also add the conversion in tcp-to-trtmlir. LMKWYT

Copy link
Collaborator

Choose a reason for hiding this comment

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

Thanks for the clarification. Can you add a comment in the code to explain this?

@zezhang zezhang merged commit ee58041 into cruise-automation:main Oct 7, 2024
1 check passed
@zezhang zezhang deleted the zezhang/torch_arange branch October 7, 2024 20:35
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