[Relax][PyTorch] Support basic range constraints#18429
Conversation
guan404ming
force-pushed
the
fix-relax-pytorch-constraints
branch
from
05fd731 to
77685df
Compare
|
cc @mshr-h |
mshr-h
left a comment
mshr-h
left a comment
There was a problem hiding this comment.
Thanks! Few changes are needed.
| else s | ||
| for s in torch_shape | ||
| ] | ||
| # Create TIR variables for symbolic dimensions |
There was a problem hiding this comment.
Looks like there's no functional changes. Any reason for the change?
There was a problem hiding this comment.
Sorry about this, it seems like legacy change in #17898. Let me revert this.
| if range_constraints: | ||
| if func_attrs is None: | ||
| func_attrs = {} | ||
| func_attrs["shape_var_constraints"] = range_constraints |
There was a problem hiding this comment.
Please use tir_var_upper_bound to annotate upper bound.
I grepped the tvm code base and I found that there's no lower bound annotation. So I don't think we need to keep it at the moment. If we have a real use case for it, it's fine to keep it in the Relax module.
https://github.com/apache/tvm/blob/main/src/relax/transform/static_plan_block_memory.cc#L62-L66
There was a problem hiding this comment.
Thanks for letting me know, I've updated to use tir_var_upper_bound to annotate upper bound.
| mod = from_exported_program(exported_program) | ||
|
|
||
| main_func = mod["main"] | ||
| assert hasattr(main_func, "attrs"), "Function should have attributes" |
There was a problem hiding this comment.
Please use structual equality instead of manually checking the attributes.
There was a problem hiding this comment.
Sure, I've updated the test to follow the convention.
![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 adds support for PyTorch range constraints in the Relax frontend. The changes correctly extract range constraints from the ExportedProgram and attach them as function attributes. A new test case is included to verify this functionality. My review includes a couple of suggestions to improve code conciseness and test robustness.
| dynamic_shapes = {"x1": {0: batch}, "x2": {0: batch}} | ||
| exported_program = export(DynamicModel(), args=example_args, dynamic_shapes=dynamic_shapes) | ||
|
|
||
| mod = from_exported_program(exported_program) |
There was a problem hiding this comment.
Should be from_exported_program(exported_program, run_ep_decomposition=True)
See #18399
There was a problem hiding this comment.
Thanks for providing useful info!
There was a problem hiding this comment.
Should we also need to update tests in https://github.com/apache/tvm/blob/f574031657faa035ffd1875050316b11821187cf/tests/python/relax/test_frontend_from_exported_program.py?
I could help with this if needed.
There was a problem hiding this comment.
Of course you can help update it!
There was a problem hiding this comment.
Sure! I would open another PR to handle this, thanks for your reply.
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
guan404ming
force-pushed
the
fix-relax-pytorch-constraints
branch
from
93b699f to
34476a4
Compare
|
Thanks for your detailed comments and reviews! |
guan404ming
changed the title
3 participants
Related Issue
related to #17818
Why
How
shape_var_constraintsfunction attributetest test_dynamic_shape_with_range_constraintsto verify constraint extraction works correctly