Skip to content
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

[Pipeline] Use dedicated simplification pipeline for TorchDynamo frontend #3376

Merged
merged 1 commit into from
May 22, 2024

Conversation

sjain-stanford
Copy link
Member

@sjain-stanford sjain-stanford commented May 22, 2024

Discord Thread: https://discord.com/channels/636084430946959380/1238330633328005243

Context:

This was updated to support e2e tests for the TorchDynamo frontend in Torch-MLIR, where we run FX decompositions and import the FX IR to generate Torch dialect, followed by torch-function-to-torch-backend-pipeline, skipping only the shape/type refinement for now. However, we should be able to skip many of the torch simplification passes, as depicted in the frontend roadmap.

Based on IREE's TorchDynamo pipeline, the only two passes we seem to require are: ReduceOpVariantsPass and DecomposeComplexOpsPass. This is inline with our findings as well based on initial exploration.

This PR creates a dedicated frontend simplification pipeline for TorchDynamo / FX Importer which calls only ReduceOpVariantsPass and DecomposeComplexOpsPass. We rely on the e2e fx_importer tests to ensure we're not regressing by removing many of the passes that were historically needed for TorchScript.

One notable change here is that we do not call the LowerToBackendContractPass anymore, which used to call TorchSimplificationPipeline iteratively until VerifyBackendContract was clean. Some of this was required for the shape/type refinement to converge, which seems a non-issue for Dynamo frontend. Do we anticipate this (the iterative invocation of TorchSimplificationPipeline followed by VerifyBackendContract) to be worth retaining in the Dynamo frontend pipeline? If so, I can make those changes, PLMK.

@sjain-stanford sjain-stanford changed the title [Pipeline] Use dedicated simplification pipeline for TorchDynamo/fx_importer flows [Pipeline] Use dedicated simplification pipeline for TorchDynamo frontend May 22, 2024
@sjain-stanford sjain-stanford force-pushed the sambhav/fx_pipeline_fixes branch from 7e89083 to 2cb35fc Compare May 22, 2024 10:41
@sjain-stanford sjain-stanford merged commit 6e48557 into llvm:main May 22, 2024
3 checks passed
@sjain-stanford sjain-stanford deleted the sambhav/fx_pipeline_fixes branch May 22, 2024 12:23
@stellaraccident
Copy link
Collaborator

Can we get a post review from @rsuderman ?

While probably ok, I wish we had left time for a design change like this to get reviewed by all parties vs an author and merge in 2 hours.

@sjain-stanford
Copy link
Member Author

Appreciate the note @stellaraccident , I will be on standby for @rsuderman 's review and send a follow-on to address/revise/revert as appropriate.

BaneTrifa pushed a commit to BaneTrifa/torch-mlir that referenced this pull request May 24, 2024
…tend (llvm#3376)

Discord Thread:
https://discord.com/channels/636084430946959380/1238330633328005243

## Context: 

[This](https://github.com/llvm/torch-mlir/blob/main/python/torch_mlir/fx.py#L61)
was updated to support e2e tests for the TorchDynamo frontend in
Torch-MLIR, where we run FX decompositions and import the FX IR to
generate Torch dialect, followed by
`torch-function-to-torch-backend-pipeline`, skipping only the shape/type
refinement for now. However, we should be able to skip many of the torch
simplification passes, as depicted in the [frontend
roadmap](https://github.com/llvm/torch-mlir/blob/main/docs/images/roadmap_frontend.png).

Based on IREE's TorchDynamo
[pipeline](https://github.com/iree-org/iree/blob/main/compiler/plugins/input/Torch/InputConversion/Passes.cpp#L29),
the only two passes we seem to require are: `ReduceOpVariantsPass` and
`DecomposeComplexOpsPass`. This is inline with our findings as well
based on initial exploration.

This PR creates a dedicated frontend simplification pipeline for
TorchDynamo / FX Importer which calls only `ReduceOpVariantsPass` and
`DecomposeComplexOpsPass`. We rely on the e2e fx_importer tests to
ensure we're not regressing by removing many of the passes that were
historically needed for TorchScript.

One notable change here is that we do not call the
`LowerToBackendContractPass` anymore, which used to call
`TorchSimplificationPipeline` iteratively until VerifyBackendContract
was clean. Some of this was required for the shape/type refinement to
converge, which seems a non-issue for Dynamo frontend. Do we anticipate
this (the iterative invocation of TorchSimplificationPipeline followed
by VerifyBackendContract) to be worth retaining in the Dynamo frontend
pipeline? If so, I can make those changes, PLMK.
sjarus pushed a commit to sjarus/torch-mlir that referenced this pull request Jun 6, 2024
…tend (llvm#3376)

Discord Thread:
https://discord.com/channels/636084430946959380/1238330633328005243

## Context: 

[This](https://github.com/llvm/torch-mlir/blob/main/python/torch_mlir/fx.py#L61)
was updated to support e2e tests for the TorchDynamo frontend in
Torch-MLIR, where we run FX decompositions and import the FX IR to
generate Torch dialect, followed by
`torch-function-to-torch-backend-pipeline`, skipping only the shape/type
refinement for now. However, we should be able to skip many of the torch
simplification passes, as depicted in the [frontend
roadmap](https://github.com/llvm/torch-mlir/blob/main/docs/images/roadmap_frontend.png).

Based on IREE's TorchDynamo
[pipeline](https://github.com/iree-org/iree/blob/main/compiler/plugins/input/Torch/InputConversion/Passes.cpp#L29),
the only two passes we seem to require are: `ReduceOpVariantsPass` and
`DecomposeComplexOpsPass`. This is inline with our findings as well
based on initial exploration.

This PR creates a dedicated frontend simplification pipeline for
TorchDynamo / FX Importer which calls only `ReduceOpVariantsPass` and
`DecomposeComplexOpsPass`. We rely on the e2e fx_importer tests to
ensure we're not regressing by removing many of the passes that were
historically needed for TorchScript.

One notable change here is that we do not call the
`LowerToBackendContractPass` anymore, which used to call
`TorchSimplificationPipeline` iteratively until VerifyBackendContract
was clean. Some of this was required for the shape/type refinement to
converge, which seems a non-issue for Dynamo frontend. Do we anticipate
this (the iterative invocation of TorchSimplificationPipeline followed
by VerifyBackendContract) to be worth retaining in the Dynamo frontend
pipeline? If so, I can make those changes, PLMK.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants