fix: Reorganize Dynamo backends

- Rename key backends to establish default backend and optional
alternatives
- Update function headers and docstrings, as well as key imports
- Rename `torch_compile` folder to `backends` in accordance with
`torch.compile` terminology
- Update references throughout codebase
- Specify certain functions as private/helper via underscore

Author: gs-olive

py/torch_tensorrt/_compile.py

@@ -157,7 +157,7 @@ def compile(
             **kwargs,
         )
     elif target_ir == _IRType.torch_compile:
-        return torch_tensorrt.dynamo.torch_compile(
+        return torch_tensorrt.dynamo.compile(
             module, inputs=inputs, enabled_precisions=enabled_precisions, **kwargs
         )
     elif target_ir == _IRType.fx_ts_compat:

py/torch_tensorrt/dynamo/__init__.py

@@ -1,2 +1,2 @@
 from torch_tensorrt.dynamo import fx_ts_compat
-from .torch_compile import compile as torch_compile
+from .backend import compile

py/torch_tensorrt/dynamo/torch_compile/__init__.py renamed to py/torch_tensorrt/dynamo/backend/__init__.py

@@ -8,10 +8,10 @@
 from torch_tensorrt import EngineCapability, Device
 from torch_tensorrt.fx.utils import LowerPrecision
 
-from torch_tensorrt.dynamo.torch_compile._settings import CompilationSettings
-from torch_tensorrt.dynamo.torch_compile.utils import prepare_inputs, prepare_device
-from torch_tensorrt.dynamo.torch_compile.backends import tensorrt_backend
-from torch_tensorrt.dynamo.torch_compile._defaults import (
+from torch_tensorrt.dynamo.backend._settings import CompilationSettings
+from torch_tensorrt.dynamo.backend.utils import prepare_inputs, prepare_device
+from torch_tensorrt.dynamo.backend.backends import torch_tensorrt_backend
+from torch_tensorrt.dynamo.backend._defaults import (
     PRECISION,
     DEBUG,
     MAX_WORKSPACE_SIZE,
@@ -121,6 +121,6 @@ def create_backend(
     )
 
     return partial(
-        tensorrt_backend,
+        torch_tensorrt_backend,
         settings=settings,
     )

py/torch_tensorrt/dynamo/torch_compile/_defaults.py renamed to py/torch_tensorrt/dynamo/backend/_defaults.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 
 from torch_tensorrt.fx.utils import LowerPrecision
-from torch_tensorrt.dynamo.torch_compile._defaults import (
+from torch_tensorrt.dynamo.backend._defaults import (
     PRECISION,
     DEBUG,
     MAX_WORKSPACE_SIZE,

py/torch_tensorrt/dynamo/torch_compile/_settings.py renamed to py/torch_tensorrt/dynamo/backend/_settings.py

@@ -4,30 +4,42 @@
 from functools import partial
 import torch._dynamo as td
 
-from torch_tensorrt.dynamo.torch_compile._settings import CompilationSettings
-from torch_tensorrt.dynamo.torch_compile.lowering._decompositions import (
+from torch_tensorrt.dynamo.backend._settings import CompilationSettings
+from torch_tensorrt.dynamo.backend.lowering._decompositions import (
     get_decompositions,
 )
-from torch_tensorrt.dynamo.torch_compile.lowering._partition import (
+from torch_tensorrt.dynamo.backend.lowering._partition import (
     partition,
     get_submod_inputs,
 )
-from torch_tensorrt.dynamo.torch_compile.conversion import convert_module
+from torch_tensorrt.dynamo.backend.conversion import convert_module
 
 from torch._dynamo.backends.common import fake_tensor_unsupported
 
 from torch._functorch.aot_autograd import aot_module_simplified, make_boxed_compiler
 
 
-@td.register_backend(name="tensorrt")
+@td.register_backend(name="torch_tensorrt")
 @fake_tensor_unsupported
-def tensorrt_backend(
-    gm: torch.nn.Module,
+def torch_tensorrt_backend(
+    gm: torch.fx.GraphModule,
+    sample_inputs: Sequence[torch.Tensor],
+    settings: CompilationSettings = CompilationSettings(),
+):
+    DEFAULT_BACKEND = aot_torch_tensorrt_aten_backend
+
+    return DEFAULT_BACKEND(gm=gm, sample_inputs=sample_inputs, settings=settings)
+
+
+@td.register_backend(name="aot_torch_tensorrt_aten")
+@fake_tensor_unsupported
+def aot_torch_tensorrt_aten_backend(
+    gm: torch.fx.GraphModule,
     sample_inputs: Sequence[torch.Tensor],
     settings: CompilationSettings = CompilationSettings(),
 ):
     custom_backend = partial(
-        fx_dynamo_backend,
+        _pretraced_backend,
         settings=settings,
     )
 
@@ -40,14 +52,12 @@ def tensorrt_backend(
     )
 
 
-@td.register_backend(name="fx_tensorrt")
-@fake_tensor_unsupported
-def fx_dynamo_backend(
+def _pretraced_backend(
     gm: torch.fx.GraphModule,
-    example_inputs: Sequence[torch.Tensor],
+    sample_inputs: Sequence[torch.Tensor],
     settings: CompilationSettings = CompilationSettings(),
 ):
-    """Helper function to manage translation of FX module to TRT engines
+    """Helper function to manage translation of traced FX module to TRT engines
 
     Args:
         module: FX GraphModule to convert
@@ -57,9 +67,9 @@ def fx_dynamo_backend(
         Compiled FX GraphModule
     """
     try:
-        trt_compiled = compile_module(
+        trt_compiled = _compile_module(
             gm,
-            example_inputs,
+            sample_inputs,
             settings=settings,
         )
         return trt_compiled
@@ -72,12 +82,12 @@ def fx_dynamo_backend(
         return gm.forward
 
 
-def compile_module(
+def _compile_module(
     gm: torch.fx.GraphModule,
-    example_inputs: Sequence[torch.Tensor],
+    sample_inputs: Sequence[torch.Tensor],
     settings: CompilationSettings = CompilationSettings(),
 ) -> torch.fx.GraphModule:
-    """Compile an FX module
+    """Compile a traced FX module
 
     Includes: Partitioning + Conversion Phases
 
@@ -100,7 +110,7 @@ def compile_module(
 
         # Get submodule inputs
         submodule_inputs = get_submod_inputs(
-            partitioned_module, submodule, example_inputs
+            partitioned_module, submodule, sample_inputs
         )
 
         # Create TRT Module from submodule

py/torch_tensorrt/dynamo/torch_compile/backends.py renamed to py/torch_tensorrt/dynamo/backend/backends.py

@@ -0,0 +1,7 @@
+from torch_tensorrt.dynamo.backend.lowering._decompositions import (
+    get_decompositions,
+)
+from torch_tensorrt.dynamo.backend.lowering._partition import (
+    partition,
+    get_submod_inputs,
+)

py/torch_tensorrt/dynamo/torch_compile/conversion.py renamed to py/torch_tensorrt/dynamo/backend/conversion.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from torch_tensorrt.dynamo.torch_compile._defaults import MAX_NUM_TRT_ENGINES
+from torch_tensorrt.dynamo.backend._defaults import MAX_NUM_TRT_ENGINES
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
 from torch.fx.passes.operator_support import OperatorSupport