Skip to content

feat: Add ExportedProgram as an IR #2191

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.

Sign up for GitHub

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

Merged
merged 4 commits into from
Aug 16, 2023
Merged

feat: Add ExportedProgram as an IR #2191

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c2c2e39
feat: Add ExportedProgram as an IR
peri044 Aug 10, 2023
f9c836c
chore: linter fixes
peri044 Aug 11, 2023
ead2aa6
chore: modify behavior of exported_program
peri044 Aug 11, 2023
140fefe
chore: rebase to main
peri044 Aug 16, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions py/torch_tensorrt/_compile.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import torch
import torch.fx
import torch_tensorrt.ts
from torch._export import ExportedProgram
from torch_tensorrt._enums import dtype
from torch_tensorrt._Input import Input
from torch_tensorrt.dynamo.compile import compile as dynamo_compile
Expand Down Expand Up @@ -43,6 +44,7 @@ class _IRType(Enum):
fx = 1
dynamo = 2
torch_compile = 3
exported_program = 4
Copy link
Collaborator

Choose a reason for hiding this comment

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

Unrelated to this change, but the docstring on this Enum is not relevant to the content



class _ModuleType(Enum):
Expand All @@ -51,6 +53,7 @@ class _ModuleType(Enum):
nn = 0
ts = 1
fx = 2
ep = 3


def _parse_module_type(module: Any) -> _ModuleType:
Expand All @@ -61,6 +64,8 @@ def _parse_module_type(module: Any) -> _ModuleType:
return _ModuleType.ts
elif isinstance(module, torch.fx.GraphModule):
return _ModuleType.fx
elif isinstance(module, ExportedProgram):
return _ModuleType.ep
elif isinstance(module, torch.nn.Module):
return _ModuleType.nn
else:
Expand All @@ -70,6 +75,7 @@ def _parse_module_type(module: Any) -> _ModuleType:
def _get_target_ir(module_type: _ModuleType, ir: str) -> _IRType:
module_is_tsable = any(module_type == t for t in [_ModuleType.nn, _ModuleType.ts])
module_is_fxable = any(module_type == t for t in [_ModuleType.nn, _ModuleType.fx])
module_is_exportable = module_type == _ModuleType.ep

ir_targets_torchscript = any(ir == opt for opt in ["torchscript", "ts"])
ir_targets_fx = ir == "fx"
Expand All @@ -95,8 +101,16 @@ def _get_target_ir(module_type: _ModuleType, ir: str) -> _IRType:
"Input graph is a Torchscript module but the ir provided is default (dynamo). Please set ir=torchscript to suppress the warning. Compiling the module with ir=torchscript"
)
return _IRType.ts
elif module_is_exportable:
raise ValueError(
"Input graph is an ExportedProgram which is not currently supported. Please provide torch.nn.Module or torch.fx.GraphModule as input."
)
else:
raise ValueError("Module was provided in an unsupported format")
elif ir == "exported_program":
raise ValueError(
"ir=exported_program is not currently supported. Supported ir options : ts|fx|dynamo"
)
Comment on lines +110 to +113
Copy link
Collaborator

Choose a reason for hiding this comment

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

Same as above for input type.

else:
raise ValueError("Unknown ir was requested")

Expand Down