Eager mode #635
Merged
Eager mode #635
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Closed
|
Can you add a test config for the e2e test framework? https://github.com/llvm/torch-mlir/tree/main/python/torch_mlir_e2e_test/torchscript/configs It should be possible to pass all the tests if you implement the fallback correctly. |
makslevental
commented
Mar 2, 2022
makslevental
commented
Mar 2, 2022
makslevental
commented
Mar 2, 2022
silvasean
reviewed
Mar 2, 2022
makslevental
commented
Mar 2, 2022
makslevental
commented
Mar 2, 2022
makslevental
force-pushed
the
eager_clean
branch
11 times, most recently
from
433330b to
47a6a3d
Compare
silvasean
requested changes
Mar 14, 2022
makslevental
force-pushed
the
eager_clean
branch
4 times, most recently
from
5aecdb6 to
8a0be06
Compare
silvasean
reviewed
Mar 16, 2022
silvasean
reviewed
Mar 16, 2022
makslevental
force-pushed
the
eager_clean
branch
2 times, most recently
from
7e23d5e to
248a3f9
Compare
makslevental
force-pushed
the
eager_clean
branch
from
248a3f9 to
8e27b5c
Compare
silvasean
requested changes
Mar 21, 2022
makslevental
force-pushed
the
eager_clean
branch
5 times, most recently
from
8bb719a to
dadb46f
Compare
makslevental
force-pushed
the
eager_clean
branch
3 times, most recently
from
bbc8e27 to
cf86e01
Compare
silvasean
approved these changes
Mar 22, 2022
| from torch_mlir_e2e_test.linalg_on_tensors_backends import refbackend | ||
|
|
||
|
|
||
| class TorchMLIRTensor(torch.Tensor): |
There was a problem hiding this comment.
Is there any upstream documentation describing the extension point you are using here? (_make_wrapper_subclass/`_torch_dispatch``, etc.) It would be good to link it in if it exists.
There was a problem hiding this comment.
I spoke with brian hirsch and he said the best documentation for how to use wrapper_subclass is https://github.com/albanD/subclass_zoo. I can add as a comment.
…h-mlir framework. This is accomplished by overriding the `__torch_dispatch__` class method on wrapper subclass `TorchMLIRTensor(torch.Tensor)`. Effectively, this mode works by compiling op by op as the NN is eagerly executed by PyTorch. Entailed in that compilation is building a representation of the op that can be `torch.jit.script`ed, importing using `ModuleBuilder`, and then executing (e.g., with `RefBackendLinalgOnTensorsBackend`). This mode includes a fallback to conventional PyTorch if anything in the torch-mlir compilation process fails (e.g., unsupported op). Currently, all e2e tests pass execpt for two that involve an upstream PyTorch bug (pytorch/pytorch#74400). High priority next steps: 1. A compile cache in order to speed up reruns of the same NN. 2. Integration with IREE (though not in this repo). 3. Integration with `torch.distributed`.
makslevental
force-pushed
the
eager_clean
branch
from
cf86e01 to
e22e2b6
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR implements an eager mode backend for PyTorch through the torch-mlir framework. This is accomplished by overriding the
__torch_dispatch__class method on wrapper subclassTorchMLIRTensor(torch.Tensor).Effectively, this mode works by compiling op by op as the NN is eagerly executed by PyTorch. Entailed in that compilation is building a representation of the op that can be
torch.jit.scripted, importing usingModuleBuilder, and then executing (e.g., withRefBackendLinalgOnTensorsBackend). This mode includes a fallback to conventional PyTorch if anything in the torch-mlir compilation process fails (e.g., unsupported op).Currently, all e2e tests pass execpt for two that involve an upstream PyTorch bug (pytorch/pytorch#74400).
High priority next steps:
torch.distributed.