feat: Add new convert_module function

- Improve overall documentation and commenting, improve code delineation
and separation of functionality

Author: gs-olive

py/torch_tensorrt/dynamo/_compiler.py

@@ -1,26 +1,17 @@
 import torch
 import logging
+from typing import Sequence, Any
 
 from torch_tensorrt import EngineCapability, Device
 
-from torch_tensorrt.dynamo.lowering._partition import partition
 from torch_tensorrt.dynamo import create_backend
 
-from torch_tensorrt.fx.fx2trt import (
-    InputTensorSpec,
-    TRTInterpreter,
-)
-import tensorrt as trt
-
-from torch_tensorrt.fx.trt_module import TRTModule
-from torch_tensorrt.fx.utils import LowerPrecision
-
 logger = logging.getLogger(__name__)
 
 
 def compile(
     gm: torch.Module,
-    example_inputs,
+    example_inputs: Sequence[Any],
     *,
     device=Device._current_device(),
     disable_tf32=False,
@@ -30,7 +21,7 @@ def compile(
     debug=False,
     capability=EngineCapability.default,
     num_avg_timing_iters=1,
-    workspace_size=0,
+    workspace_size=20 << 30,
     dla_sram_size=1048576,
     dla_local_dram_size=1073741824,
     dla_global_dram_size=536870912,
@@ -63,52 +54,8 @@ def compile(
     )
 
     model = torch.compile(gm, backend=custom_backend)
-    # Ensure compilation
-    model(example_inputs)
-
-    return model
-
-
-def compile_logic(gm: torch.fx.GraphModule, example_inputs):
-    partitioned = partition(gm)
-
-    precision = LowerPrecision.FP32
-
-    def get_submod_inputs(mod, submod, inputs):
-        """Helper function to get inputs to submodule"""
-        acc_inputs = None
 
-        def get_input(self, inputs):
-            nonlocal acc_inputs
-            acc_inputs = inputs
+    # Ensure compilation occurs by calling the function with provided inputs
+    model(*example_inputs)
 
-        handle = submod.register_forward_pre_hook(get_input)
-        mod(*inputs)
-        handle.remove()
-        return acc_inputs
-
-    for name, _ in partitioned.named_children():
-        submod = getattr(partitioned, name)
-
-        # Get submodule inputs
-        acc_inputs = get_submod_inputs(partitioned, submod, example_inputs)
-
-        # Create TRT Module from submodule
-        interp = TRTInterpreter(
-            submod,
-            InputTensorSpec.from_tensors(acc_inputs),
-            explicit_batch_dimension=True,
-            logger_level=trt.Logger.VERBOSE,
-        )
-
-        r = interp.run(
-            max_workspace_size=20 << 30,
-            lower_precision=precision,
-            profiling_verbosity=trt.ProfilingVerbosity.VERBOSE,
-        )
-        trt_mod = TRTModule(*r)
-
-        # Replace FX Module with TRT Module
-        setattr(partitioned, name, trt_mod)
-
-    return partitioned
+    return model

py/torch_tensorrt/dynamo/backends.py

@@ -6,11 +6,22 @@
 from torch_tensorrt import EngineCapability, Device
 from torch_tensorrt.dynamo import compile
 
+from torch_tensorrt.dynamo.lowering._decompositions import get_decompositions
+from torch_tensorrt.dynamo.lowering._partition import partition, get_submod_inputs
+from torch_tensorrt.dynamo.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
 
-from torch_tensorrt.dynamo.lowering._decompositions import get_decompositions
+from torch_tensorrt.fx.fx2trt import (
+    InputTensorSpec,
+    TRTInterpreter,
+)
+import tensorrt as trt
+
+from torch_tensorrt.fx.trt_module import TRTModule
+from torch_tensorrt.fx.utils import LowerPrecision
 
 logger = logging.getLogger(__name__)
 
@@ -97,7 +108,7 @@ def fx_dynamo_backend(
 ):
     """Helper function to manage translation of FX module to TRT engines"""
     try:
-        trt_compiled = compile(gm, example_inputs)
+        trt_compiled = compile_module(gm, example_inputs)
         return trt_compiled
     except:
         traceback.print_exc()
@@ -106,3 +117,48 @@ def fx_dynamo_backend(
             + "Returning GraphModule forward instead."
         )
         return gm.forward
+
+
+def compile_module(
+    gm: torch.fx.GraphModule,
+    example_inputs,
+    debug: bool = False,
+    workspace_size: int = 20 << 30,
+    precision: LowerPrecision = LowerPrecision.FP32,
+) -> torch.fx.GraphModule:
+    """Convert an FX module to a TRT module
+    Args:
+        module: FX GraphModule to convert
+        inputs: Inputs to the module
+        debug: Whether to print out verbose debugging information
+        workspace_size: Maximum workspace TRT is allowed to use for the module
+        precision: Model Layer precision
+    Returns:
+        TRTModule or TRTModuleNext
+    """
+    # Partition module into components that can be TRT-accelerated
+    partitioned_module = partition(gm)
+
+    # Iterate over all components that can be accelerated
+    # Generate the corresponding TRT Module for those
+    for name, _ in partitioned_module.named_children():
+        submodule = getattr(partitioned_module, name)
+
+        # Get submodule inputs
+        submodule_inputs = get_submod_inputs(
+            partitioned_module, submodule, example_inputs
+        )
+
+        # Create TRT Module from submodule
+        trt_mod = convert_module(
+            submodule,
+            submodule_inputs,
+            debug=debug,
+            workspace_size=workspace_size,
+            precision=precision,
+        )
+
+        # Replace FX Module with TRT Module
+        setattr(partitioned_module, name, trt_mod)
+
+    return partitioned_module

py/torch_tensorrt/dynamo/conversion.py

@@ -0,0 +1,48 @@
+from typing import Sequence, Union
+import torch
+from torch_tensorrt.fx.trt_module import TRTModule
+from torch_tensorrt import TRTModuleNext
+from torch_tensorrt.fx.fx2trt import (
+    InputTensorSpec,
+    TRTInterpreter,
+)
+from torch_tensorrt.fx.utils import LowerPrecision
+
+import tensorrt as trt
+
+
+def convert_module(
+    module: torch.fx.GraphModule,
+    inputs: Sequence[torch.Tensor],
+    debug: bool = False,
+    workspace_size: int = 20 << 30,
+    precision: LowerPrecision = LowerPrecision.FP32,
+) -> Union[TRTModuleNext, TRTModule]:
+    """Convert an FX module to a TRT module
+    Args:
+        module: FX GraphModule to convert
+        inputs: Sequence of Tensors representing inputs to the module
+        debug: Whether to print out verbose debugging information
+        workspace_size: Maximum workspace TRT is allowed to use for the module
+        precision: Model Layer precision
+    Returns:
+        TRTModule or TRTModuleNext
+    """
+    interp = TRTInterpreter(
+        module,
+        InputTensorSpec.from_tensors(inputs),
+        explicit_batch_dimension=True,
+        logger_level=(trt.Logger.VERBOSE if debug else trt.Logger.WARNING),
+    )
+
+    r = interp.run(
+        max_workspace_size=workspace_size,
+        lower_precision=precision,
+        profiling_verbosity=(
+            trt.ProfilingVerbosity.VERBOSE
+            if debug
+            else trt.ProfilingVerbosity.LAYER_NAMES_ONLY
+        ),
+    )
+
+    return TRTModule(*r)

py/torch_tensorrt/dynamo/lowering/__init__.py

@@ -1,2 +1,2 @@
 from torch_tensorrt.dynamo.lowering._decompositions import get_decompositions
-from torch_tensorrt.dynamo.lowering._partition import partition
+from torch_tensorrt.dynamo.lowering._partition import partition, get_submod_inputs

py/torch_tensorrt/dynamo/lowering/_partition.py

@@ -1,4 +1,4 @@
-from typing import Dict
+from typing import Dict, Optional, Sequence
 
 import torch
 
@@ -12,7 +12,7 @@
 
 
 class TorchTensorRTOperatorSupport(OperatorSupport):
-    """Class to determine whether the aten operators have converters"""
+    """Class to determine whether operators within a module are supported"""
 
     def __init__(self, support_dict=None):
         super().__init__(support_dict)
@@ -38,7 +38,7 @@ def is_node_supported(
 
             return False
 
-    def print_support_overview(self, num_trt_blocks=None):
+    def print_support_overview(self, num_trt_blocks: Optional[int] = None):
         if num_trt_blocks is not None:
             print(f"Number of TensorRT-Accelerated Subgraphs: {num_trt_blocks}\n")
 
@@ -51,9 +51,20 @@ def print_support_overview(self, num_trt_blocks=None):
             print(node_name)
 
 
-def partition(gm: torch.fx.GraphModule, verbose=True):
+def partition(
+    gm: torch.fx.GraphModule,
+    verbose: bool = True,
+    max_num_trt_engines: int = MAX_NUM_TRT_ENGINES,
+) -> torch.fx.GraphModule:
     """Partition an FX GraphModule with aten ops into TRT engines
-    Partitioning is based on operator support
+    Partitioning is based on converter operator support
+
+    Args:
+        gm: FX GraphModule to partition
+        verbose: Bool representing whether to print operator support
+        max_num_trt_engines: Maximum number of allowed TRT engines in partitioning
+    Returns:
+        torch.fx.GraphModule
     """
     supported_ops = TorchTensorRTOperatorSupport()
     partitioner = CapabilityBasedPartitioner(gm, supported_ops)
@@ -62,10 +73,10 @@ def partition(gm: torch.fx.GraphModule, verbose=True):
     # exceeds a specified threshold
     partitions = partitioner.propose_partitions()
     num_blocks = len(partitions)
-    if num_blocks > MAX_NUM_TRT_ENGINES:
+    if num_blocks > max_num_trt_engines:
         raise AssertionError(
             f"The graph module has {num_blocks} TRT Engines which is larger than the "
-            + f"threshold={MAX_NUM_TRT_ENGINES}. Falling back to non-TRT module."
+            + f"threshold={max_num_trt_engines}. Falling back to non-TRT module."
         )
 
     # Fuse partitions and display overview of supported/unsupported operators
@@ -76,3 +87,27 @@ def partition(gm: torch.fx.GraphModule, verbose=True):
         supported_ops.print_support_overview(num_blocks)
 
     return fused_graph
+
+
+def get_submod_inputs(
+    mod: torch.fx.GraphModule, submod: torch.fx.GraphModule, inputs
+) -> Sequence[torch.Tensor]:
+    """Helper function to get inputs to a Torch submodule
+
+    Args:
+        mod: Parent FX GraphModule
+        submod: Child FX GraphModule
+        inputs: Sample inputs to parent module
+    Returns:
+        Sequence of Tensors representing inputs to child module
+    """
+    acc_inputs = None
+
+    def get_input(self, inputs):
+        nonlocal acc_inputs
+        acc_inputs = inputs
+
+    handle = submod.register_forward_pre_hook(get_input)
+    mod(*inputs)
+    handle.remove()
+    return acc_inputs