[Relax][PyTorch] Add support for decomposed operators and fix IR of ops tests(2) #18403
Conversation
Summary of ChangesHello @tlopex, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request enhances the Relax PyTorch frontend by introducing support for decomposed operators. This change allows for a more granular representation of complex PyTorch operations within the Relax Intermediate Representation (IR), improving flexibility and optimization opportunities. It also includes specific fixes and parameter additions for operators like Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request enhances the PyTorch frontend in Relax by adding support for decomposed operators. It introduces a handler for scalar_tensor and adds a new mapping for _softmax.default. The majority of the changes are within the test suite, where tests for several operators (softmax, softsign, softshrink, tril, triu, div, batchnorm2d, addmm, baddbmm, bmm) have been updated to execute with run_ep_decomposition=True. This required updating the expected Relax IR to align with the output from the decomposed PyTorch operators. The changes are well-aligned with the PR's objectives. I have included a few suggestions to enhance the conciseness and readability of the updated test IRs.
| lv5: R.Tensor((1, 3, 10, 10), dtype="float32") = R.multiply( | ||
| input, R.const(0.0, "float32") | ||
| ) |
There was a problem hiding this comment.
Using R.zeros_like(input) would be more idiomatic and clearer to express the intent of creating a zero tensor with the same shape and dtype as input.
lv5: R.Tensor((1, 3, 10, 10), dtype="float32") = R.zeros_like(input)| lv2: R.Tensor((10,), dtype="int64") = R.arange( | ||
| R.prim_value(0), R.prim_value(10), R.prim_value(1), dtype="int64" | ||
| ) | ||
| lv3: R.Tensor((10, 1), dtype="int64") = R.expand_dims(lv2, axis=[-1]) |
There was a problem hiding this comment.
The arange call to create lv2 is redundant, as lv created on line 1197 is identical. You can reuse lv to create lv3 and remove the definition of lv2 to make the IR more concise and avoid a redundant operation. This same issue is present in expected_triu as well.
lv3: R.Tensor((10, 1), dtype="int64") = R.expand_dims(lv, axis=[-1])| lv2: R.Tensor((10,), dtype="int64") = R.arange( | ||
| R.prim_value(0), R.prim_value(10), R.prim_value(1), dtype="int64" | ||
| ) | ||
| lv3: R.Tensor((10, 1), dtype="int64") = R.expand_dims(lv2, axis=[-1]) |
There was a problem hiding this comment.
Similar to expected_tril, the arange call to create lv2 is redundant. You can reuse lv (defined on line 1227) to create lv3 and remove the definition of lv2.
lv3: R.Tensor((10, 1), dtype="int64") = R.expand_dims(lv, axis=[-1])
|
@tvm-bot rerun |
|
cc @mshr-h |
2 participants
This pr fixes
softmax,softsign,softshrink,tril,triu,div,batchnorm2d,addmm,baddbmm,bmm