refactor: Promote lowering to top level, remove input_tensor_spec, rename fx2trt

Signed-off-by: Dheeraj Peri <[email protected]>

chore: apply linting

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

py/torch_tensorrt/dynamo/aten_tracer.py

@@ -0,0 +1,157 @@
+import copy
+import sys
+from contextlib import contextmanager
+from typing import Any, Callable, Dict, Generator, List, Optional, Set, Tuple, Union
+from packaging import version
+
+import torch
+import torch._dynamo as torchdynamo
+
+from torch_tensorrt.fx.utils import req_torch_version
+from torch_tensorrt.fx.passes.lower_basic_pass_aten import (
+    compose_bmm,
+    compose_chunk,
+    compose_getitem_slice,
+    remove_ops,
+    replace_aten_op_with_indices,
+    replace_aten_reshape_alias_with_replace,
+    replace_builtin_ops,
+    replace_inplace_ops,
+    replace_native_layernorm_with_layernorm,
+    replace_transpose_mm_op_with_linear,
+    run_const_fold,
+)
+from typing_extensions import TypeAlias
+
+Value: TypeAlias = Union[
+    Tuple["Value", ...],
+    List["Value"],
+    Dict[str, "Value"],
+]
+
+
+class DynamoConfig:
+    """
+    Manage Exir-specific configurations of Dynamo.
+    """
+
+    def __init__(
+        self,
+        capture_scalar_outputs: bool = True,
+        guard_nn_modules: bool = True,
+        dynamic_shapes: bool = True,
+        specialize_int: bool = True,
+        verbose: bool = True,
+    ) -> None:
+
+        self.capture_scalar_outputs = capture_scalar_outputs
+        self.guard_nn_modules = guard_nn_modules
+        self.dynamic_shapes = dynamic_shapes
+        self.specialize_int = specialize_int
+        self.verbose = verbose
+
+    def activate(self) -> None:
+        torchdynamo.config.capture_scalar_outputs = self.capture_scalar_outputs
+        torchdynamo.config.guard_nn_modules = self.guard_nn_modules
+        torchdynamo.config.dynamic_shapes = self.dynamic_shapes
+        torchdynamo.config.specialize_int = self.specialize_int
+        torchdynamo.config.verbose = self.verbose
+
+    def deactivate(self) -> None:
+        torchdynamo.config.capture_scalar_outputs = True
+        torchdynamo.config.guard_nn_modules = True
+        torchdynamo.config.dynamic_shapes = True
+        torchdynamo.config.specialize_int = True
+        torchdynamo.config.verbose = True
+
+
+@contextmanager
+def using_config(config: DynamoConfig) -> Generator[DynamoConfig, None, None]:
+    config.activate()
+    try:
+        yield config
+    finally:
+        config.deactivate()
+
+
+@contextmanager
+def setting_python_recursive_limit(limit: int = 10000) -> Generator[None, None, None]:
+    """
+    Temporarily increase the python interpreter stack recursion limit.
+    This is mostly used for pickling large scale modules.
+    """
+    default = sys.getrecursionlimit()
+    if limit > default:
+        sys.setrecursionlimit(limit)
+    try:
+        yield
+    finally:
+        sys.setrecursionlimit(default)
+
+
+@req_torch_version("2.dev")
+def dynamo_trace(
+    f: Callable[..., Value],
+    # pyre-ignore
+    args: Tuple[Any, ...],
+    aten_graph: bool,
+    tracing_mode: str = "real",
+    dynamo_config: Optional[DynamoConfig] = None,
+) -> Tuple[torch.fx.GraphModule, Set]:
+    """
+    TODO: Once we fully migrate to torchdynamo frontend, we will remove
+    this config option alltogether.  For now, it helps with quick
+    experiments with playing around with TorchDynamo
+    """
+    if dynamo_config is None:
+        dynamo_config = DynamoConfig()
+    with using_config(dynamo_config), setting_python_recursive_limit(2000):
+        torchdynamo.reset()
+        try:
+            return torchdynamo.export(
+                f,
+                *copy.deepcopy(args),
+                aten_graph=aten_graph,
+                tracing_mode=tracing_mode,
+            )
+        except torchdynamo.exc.Unsupported as exc:
+            raise RuntimeError(
+                "The user code is using a feature we don't support. "
+                "Please try torchdynamo.explain() to get possible the reasons",
+            ) from exc
+        except Exception as exc:
+            raise RuntimeError(
+                "torchdynamo internal error occured. Please see above stacktrace"
+            ) from exc
+
+
+@req_torch_version("2.dev")
+def trace(model, inputs, **kwargs):
+    """
+    Optimized trace with necessary passes which re-compose some ops or replace some ops
+    These passes should be general and functional purpose
+    """
+    passes_list = [
+        compose_bmm,
+        compose_chunk,
+        compose_getitem_slice,
+        replace_aten_reshape_alias_with_replace,
+        replace_aten_op_with_indices,
+        replace_transpose_mm_op_with_linear,  # after compose_bmm
+        replace_native_layernorm_with_layernorm,
+        remove_ops,
+        replace_builtin_ops,  # after replace_native_layernorm_with_layernorm
+        replace_inplace_ops,  # remove it once functionalization is enabled
+    ]
+
+    fx_module, __package__ = dynamo_trace(model, inputs, True, "symbolic")
+    print(fx_module.graph)
+    for passes in passes_list:
+        pr: PassResult = passes(fx_module)
+        fx_module = pr.graph_module
+
+    fx_module(*inputs)
+
+    fx_module = run_const_fold(fx_module)
+    print(fx_module.graph)
+    return fx_module

py/torch_tensorrt/dynamo/backend/backends.py

@@ -5,13 +5,13 @@
 import torch._dynamo as td
 
 from torch_tensorrt.dynamo import CompilationSettings
-from torch_tensorrt.dynamo.backend.lowering._decompositions import (
+from torch_tensorrt.dynamo.lowering._decompositions import (
     get_decompositions,
 )
-from torch_tensorrt.dynamo.backend.lowering._pre_aot_lowering import (
+from torch_tensorrt.dynamo.lowering._pre_aot_lowering import (
     pre_aot_substitutions,
 )
-from torch_tensorrt.dynamo.backend.lowering._partition import (
+from torch_tensorrt.dynamo.lowering._partition import (
     partition,
     get_submod_inputs,
 )

py/torch_tensorrt/dynamo/compile.py

@@ -9,9 +9,9 @@
 from torch_tensorrt.fx.utils import LowerPrecision
 from torch.fx.passes.pass_manager import PassManager
 from torch.fx.passes.shape_prop import ShapeProp
-import torch_tensorrt.fx.tracer.dispatch_tracer.aten_tracer as aten_tracer
+from torch_tensorrt.dynamo.aten_tracer import trace
 from torch_tensorrt.fx.tools.trt_splitter import TRTSplitter, TRTSplitterSetting
-from torch_tensorrt.dynamo.backend.lowering import (
+from torch_tensorrt.dynamo.lowering import (
     fuse_permute_linear,
     fuse_permute_matmul,
 )
@@ -113,8 +113,7 @@ def compile(
     model = trace(gm, inputs, **kwargs)
 
     if kwargs.get("use_capability_partitioner", None):
-        traced_model = trace(model)
-        model = lower_model(traced_model, inputs)
+        model = lower_model(model, inputs)
         return _compile_module(model, inputs, settings)
     else:
         split_result = lower_model_using_trt_splitter(model, inputs)
@@ -146,33 +145,6 @@ def _compile_graph(
     return split_result.split_module
 
 
-def trace(
-    model: torch.nn.Module,
-    inputs: Any,
-    **kwargs,
-):
-    """Create torch.compile backend given specified arguments
-
-    Args:
-        precision: Model Layer precision
-        debug: Whether to print out verbose debugging information
-        workspace_size: Workspace TRT is allowed to use for the module (0 is default)
-        min_block_size: Minimum number of operators per TRT-Engine Block
-        torch_executed_ops: Sequence of operations to run in Torch, regardless of converter coverage
-        pass_through_build_failures: Whether to fail on TRT engine build errors (True) or not (False)
-        max_aux_streams: Maximum number of allowed auxiliary TRT streams for each engine
-        version_compatible: Provide version forward-compatibility for engine plan files
-        optimization_level: Builder optimization 0-5, higher levels imply longer build time,
-            searching for more optimization options. TRT defaults to 3
-        use_experimental_rt: Whether to use the new experimental TRTModuleNext for TRT engines
-    Returns:
-        Backend for torch.compile
-    """
-    model = aten_tracer.opt_trace(model, inputs)
-
-    return model
-
-
 def lower_model_using_trt_splitter(model: torch.nn.Module, inputs: Any, **kwargs):
     # Perform basic lowering
     model = lower_model(model, inputs)

py/torch_tensorrt/dynamo/conversion/__init__.py

@@ -1,3 +1,2 @@
-from .input_tensor_spec import *
-from .fx2trt import *
+from .trt_interpreter import *
 from .conversion import *

py/torch_tensorrt/dynamo/conversion/conversion.py

@@ -3,10 +3,9 @@
 import io
 from torch_tensorrt.fx.trt_module import TRTModule
 from torch_tensorrt.dynamo import CompilationSettings
-from torch_tensorrt.dynamo.conversion import (
-    InputTensorSpec,
-    TRTInterpreter,
-)
+from torch_tensorrt import Input
+from torch_tensorrt.dynamo.conversion import TRTInterpreter
+
 
 import tensorrt as trt
 
@@ -34,14 +33,12 @@ def convert_module(
         module_outputs = [module_outputs]
 
     output_dtypes = list(output.dtype for output in module_outputs)
-
     interpreter = TRTInterpreter(
         module,
-        InputTensorSpec.from_tensors(inputs),
+        Input.from_tensors(inputs),
         logger_level=(trt.Logger.VERBOSE if settings.debug else trt.Logger.WARNING),
         output_dtypes=output_dtypes,
     )
-
     interpreter_result = interpreter.run(
         workspace_size=settings.workspace_size,
         lower_precision=settings.precision,