-
Notifications
You must be signed in to change notification settings - Fork 25.7k
[PyTorch][2/4] Support static dispatch with multiple backends #75605
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
Conversation
Usecase: Milan models have multiple backends and need to use static dispatch to save on static initialization time and to hit native functions directly from the unboxed APIs. This change passes in List[BackendIndex] and adds ability to generate code for multiple static backends with 1 or 0 kernels Differential Revision: [D35510644](https://our.internmc.facebook.com/intern/diff/D35510644/) [ghstack-poisoned]
🔗 Helpful links
💊 CI failures summary and remediationsAs of commit 5f2af30 (more details on the Dr. CI page): Expand to see more
🕵️ 1 new failure recognized by patternsThe following CI failures do not appear to be due to upstream breakages
|
…nds" Usecase: Milan models have multiple backends and need to use static dispatch to save on static initialization time and to hit native functions directly from the unboxed APIs. This change passes in List[BackendIndex] and adds ability to generate code for multiple static backends with 1 or 0 kernels Differential Revision: [D35510644](https://our.internmc.facebook.com/intern/diff/D35510644/) [ghstack-poisoned]
|
@priyaramani has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator. |
tools/codegen/gen.py
Outdated
| elif len(keys) == 0: | ||
| return generate_static_dispatch(f, cpp_sig, method=method, backend_index=backend_indices[0]) | ||
| else: | ||
| return f'TORCH_CHECK(false, "Static dispatch does not support backends with multiple kernels");' |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hmm... if you pass in two backends to the codegen (say CPU and QuantizedCPU), then won't the codegen emit this torch check for every f.structured_delegate NativeFunction?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The else part is implemented later in the stack #75607. Split them up to keep diff sizes manageable.
Also if we pass in multiple backends for a native function, and if they have a structured_delegate and one or more kernels for the same, then this torch check will be emitted as of this diff.
| if len(keys) == 1: | ||
| return generate_static_dispatch(f, cpp_sig, method=method, backend_index=keys[0]) | ||
| elif len(keys) == 0: | ||
| return generate_static_dispatch(f, cpp_sig, method=method, backend_index=backend_indices[0]) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
What's the idea behind passing backend_indices[0] here?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
@larryliu0820
If none of the static dispatch backends have kernels for the native function, then we just pass in the first one in the list, could be any one of them, as in generate_static_dispatch(f, ..) if the backend doesn't have a kernel or if the native function doesn't have a structured delegate we only see if the native function has compositeImplicitAugradKernel or compositeImplicitAugradKernel and use that if not emit a TORCH_CHECK https://github.com/pytorch/pytorch/pull/75605/files#diff-6382c847dac2d011f822c492e61f9fa0eceb47c2d6fc5fb4f0c42af726c1a3aaR283-R288
…nds" Usecase: Milan models have multiple backends and need to use static dispatch to save on static initialization time and to hit native functions directly from the unboxed APIs. This change passes in List[BackendIndex] and adds ability to generate code for multiple static backends with 1 or 0 kernels Differential Revision: [D35510644](https://our.internmc.facebook.com/intern/diff/D35510644/) [ghstack-poisoned]
…nds" Usecase: Milan models have multiple backends and need to use static dispatch to save on static initialization time and to hit native functions directly from the unboxed APIs. This change passes in List[BackendIndex] and adds ability to generate code for multiple static backends with 1 or 0 kernels Differential Revision: [D35510644](https://our.internmc.facebook.com/intern/diff/D35510644/) [ghstack-poisoned]
…nds" Usecase: Milan models have multiple backends and need to use static dispatch to save on static initialization time and to hit native functions directly from the unboxed APIs. This change passes in List[BackendIndex] and adds ability to generate code for multiple static backends with 1 or 0 kernels Differential Revision: [D35510644](https://our.internmc.facebook.com/intern/diff/D35510644/) [ghstack-poisoned]
…nds" Usecase: Milan models have multiple backends and need to use static dispatch to save on static initialization time and to hit native functions directly from the unboxed APIs. This change passes in List[BackendIndex] and adds ability to generate code for multiple static backends with 1 or 0 kernels Differential Revision: [D35510644](https://our.internmc.facebook.com/intern/diff/D35510644/) [ghstack-poisoned]
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Is it possible to add a unit test before commit? Just to demonstrate the ability we are adding to static dispatch. You can find two operators with different kernels and write a small model.
…nds" Usecase: Milan models have multiple backends and need to use static dispatch to save on static initialization time and to hit native functions directly from the unboxed APIs. This change passes in List[BackendIndex] and adds ability to generate code for multiple static backends with 1 or 0 kernels Differential Revision: [D35510644](https://our.internmc.facebook.com/intern/diff/D35510644/) [ghstack-poisoned]
…nds" Usecase: Milan models have multiple backends and need to use static dispatch to save on static initialization time and to hit native functions directly from the unboxed APIs. This change passes in List[BackendIndex] and adds ability to generate code for multiple static backends with 1 or 0 kernels Differential Revision: [D35510644](https://our.internmc.facebook.com/intern/diff/D35510644/) [ghstack-poisoned]
…nds" Usecase: Milan models have multiple backends and need to use static dispatch to save on static initialization time and to hit native functions directly from the unboxed APIs. This change passes in List[BackendIndex] and adds ability to generate code for multiple static backends with 1 or 0 kernels Differential Revision: [D35510644](https://our.internmc.facebook.com/intern/diff/D35510644/) [ghstack-poisoned]
Pull Request resolved: #75605 Usecase: Milan models have multiple backends and need to use static dispatch to save on static initialization time and to hit native functions directly from the unboxed APIs. This change passes in List[BackendIndex] and adds ability to generate code for multiple static backends with 1 or 0 kernels ghstack-source-id: 154525738 Differential Revision: [D35510644](https://our.internmc.facebook.com/intern/diff/D35510644/) **NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D35510644/)!
Summary: Pull Request resolved: #75605 Usecase: Milan models have multiple backends and need to use static dispatch to save on static initialization time and to hit native functions directly from the unboxed APIs. This change passes in List[BackendIndex] and adds ability to generate code for multiple static backends with 1 or 0 kernels ghstack-source-id: 154525738 (Note: this ignores all push blocking failures!) Test Plan: Builds lite_predictor_flatbuffer with multiple backends ``` buck build --config pt.enable_lightweight_dispatch=1 --config pt.static_dispatch_backend=CPU,QuantizedCPU,CompositeExplicitAutograd //xplat/caffe2/fb/lite_predictor:lite_predictor_flatbuffer ``` Reviewed By: larryliu0820 Differential Revision: D35510644 fbshipit-source-id: f985718ad066f8578b006b4759c4a3bd6caac176
|
Hey @priyaramani. |
Stack from ghstack (oldest at bottom):
Usecase: Milan models have multiple backends and need to use static dispatch to save on static initialization time and to hit native functions directly from the unboxed APIs.
This change passes in List[BackendIndex] and adds ability to generate code for multiple static backends with 1 or 0 kernels
Differential Revision: D35510644