Specialize sym node when used as device kwarg #131811
Conversation
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/131811
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 6667fbc with merge base 6c31e02 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
Fixes #131189. By specializing the sym node when used as kwargs device. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]
|
While it is possible that this is the best possible implementation, I am meh about it because it does not work with non strict export. Are we sure it's not possible to modify the argument parser to force the spec? |
Fixes #131189. By specializing the sym node when used as kwargs device. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]
Updated the pr. In python_arg_parser.cpp, we allow the signature check to pass when device receives a SymInt input and 2. specialize it to integer. Also add an non_strict export test. |
Fixes #131189. We specialize the symint in python_arg_parser when used as kwarg device. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]
ydwu4
changed the title
| // then cast it back to python. | ||
| auto guarded_int = py::cast<c10::SymInt>(py::handle(obj)) | ||
| .guard_int(__FILE__, __LINE__); | ||
| dst[i++] = py::cast(guarded_int).ptr(); |
There was a problem hiding this comment.
This sure looks like it will result in an excessive decref later (if dst is owning). Can you argue why this has the right lifetime? (cc @albanD)
There was a problem hiding this comment.
Yeah, the life time is wrong. I'm not very familar to pybind and digged a bit into it: dst doesn't own py::cast(guarded_int). The integer is a local variable so dst will have a dangling pointer to it after we exits the if-statement. The previous diff works probably because gc hasn't cleaned the symint up. Thanks for pointing it out!
Updated the diff to defer the specialization to when we're actually creating Device from argument.
Fixes #131189. We specialize the symint in python_arg_parser when used as kwarg device. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]
Fixes #131189. We specialize the symint in python_arg_parser when used as kwarg device. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]
albanD
left a comment
albanD
left a comment
There was a problem hiding this comment.
Updated code sounds good
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Stack from ghstack (oldest at bottom):
Fixes #131189.
We specialize the symint in python_arg_parser when used as kwarg device.
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames