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

fix: Add lowering pass to remove output repacking in convert_method_to_trt_engine calls #1945

Merged
merged 1 commit into from
May 31, 2023

Conversation

gs-olive
Copy link
Collaborator

@gs-olive gs-olive commented May 23, 2023

Description

  • Adds improved support for full-conversion and TRT engine export for a variety of models
  • Automatically remove output repacking for convert_method_to_trt_engine calls, to improve parity between models which can be converted directly to TRT engines, and models which can be fully compiled
  • Add new internal CompileSpec argument for lowering which indicates whether the lowering passes originate from a convert_method_to_trt_engine call or a regular compile call, which affects whether the lowering pass is applied
  • Regular TorchScript graphs cannot have this pass applied, as it can otherwise break the output graph. Newer versions of Torch disallow graph outputs with 0 or 2+ arguments which are not packed in a struct
  • Current lowering pass detects outputs which are flat Lists or Tuples of Tensors and returns the outputs as-is (direct from the TRT Engine), so the entire model can be converted to a single TRT engine

Fixes #1938
Fixes #1939

Type of change

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)

Checklist:

  • [ x ] My code follows the style guidelines of this project (You can use the linters)
  • [ x ] I have performed a self-review of my own code
  • [ x ] I have commented my code, particularly in hard-to-understand areas and hacks
  • [ x ] I have made corresponding changes to the documentation
  • [ x ] I have added tests to verify my fix or my feature
  • [ x ] New and existing unit tests pass locally with my changes
  • [ x ] I have added the relevant labels to my PR in so that relevant reviewers are notified

@gs-olive gs-olive requested a review from narendasan May 23, 2023 17:46
@gs-olive gs-olive self-assigned this May 23, 2023
@github-actions github-actions bot added component: api [Python] Issues re: Python API component: api [C++] Issues re: C++ API component: core Issues re: The core compiler component: lowering Issues re: The lowering / preprocessing passes component: tests Issues re: Tests labels May 23, 2023
@gs-olive gs-olive force-pushed the remove_output_collection_casts branch from 91a8d43 to e8b6eb1 Compare May 31, 2023 19:01
- Automatically remove output repacking for
`convert_method_to_trt_engine` calls, to improve parity between models
which can be converted directly to TRT engines, and models which can be
fully compiled
- Add new internal `CompileSpec` argument for lowering which indicates
whether the lowering passes originate from a
`convert_method_to_trt_engine` call or a regular `compile` call, which
affects whether the lowering pass is applied
- Regular TorchScript graphs cannot have this pass applied, as it can
otherwise break the output graph. Newer versions of Torch disallow graph
outputs with 0 or 2+ arguments which are not packed in a struct
- Current lowering pass detects outputs which are flat Lists or Tuples
of Tensors and returns the outputs as-is (direct from the TRT Engine),
so the entire model can be converted to a single TRT engine
@gs-olive gs-olive force-pushed the remove_output_collection_casts branch from e8b6eb1 to 62371d9 Compare May 31, 2023 19:04
Copy link
Collaborator

@narendasan narendasan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@gs-olive gs-olive merged commit dd31c9a into pytorch:main May 31, 2023
@gs-olive gs-olive deleted the remove_output_collection_casts branch May 31, 2023 20:52
narendasan pushed a commit that referenced this pull request Jun 2, 2023
narendasan pushed a commit that referenced this pull request Jun 2, 2023
narendasan pushed a commit that referenced this pull request Jun 3, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
cla signed component: api [C++] Issues re: C++ API component: api [Python] Issues re: Python API component: core Issues re: The core compiler component: lowering Issues re: The lowering / preprocessing passes component: tests Issues re: Tests
Projects
None yet
3 participants