chore/fix: Restructure Dynamo directory

gs-olive · gs-olive · commit 4ed3e3fe71e0 · 2023-06-05T10:21:23.000-07:00
- Add `common` directory which stores code common to both the compile
and export path, to reduce code duplication and better organize the
repository
- Update necessary imports, fix minor argument pass-through issues in
`fx_ts_compat`
diff --git a/py/torch_tensorrt/dynamo/backend/__init__.py b/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.backend._settings import CompilationSettings
+from torch_tensorrt.dynamo.common 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 (
+from torch_tensorrt.dynamo.common._defaults import (
     PRECISION,
     DEBUG,
     WORKSPACE_SIZE,
diff --git a/py/torch_tensorrt/dynamo/backend/backends.py b/py/torch_tensorrt/dynamo/backend/backends.py
@@ -5,7 +5,7 @@
 from dataclasses import replace, fields
 import torch._dynamo as td
 
-from torch_tensorrt.dynamo.backend._settings import CompilationSettings
+from torch_tensorrt.dynamo.common import CompilationSettings
 from torch_tensorrt.dynamo.backend.lowering._decompositions import (
     get_decompositions,
 )
diff --git a/py/torch_tensorrt/dynamo/backend/conversion.py b/py/torch_tensorrt/dynamo/backend/conversion.py
@@ -3,8 +3,8 @@
 import io
 from torch_tensorrt.fx.trt_module import TRTModule
 from torch_tensorrt import TRTModuleNext
-from torch_tensorrt.dynamo.backend._settings import CompilationSettings
-from torch_tensorrt.dynamo.fx_ts_compat.fx2trt import (
+from torch_tensorrt.dynamo.common import (
+    CompilationSettings,
     InputTensorSpec,
     TRTInterpreter,
 )
diff --git a/py/torch_tensorrt/dynamo/backend/lowering/_partition.py b/py/torch_tensorrt/dynamo/backend/lowering/_partition.py
@@ -3,7 +3,7 @@
 
 import torch
 
-from torch_tensorrt.dynamo.backend._defaults import MIN_BLOCK_SIZE
+from torch_tensorrt.dynamo.common._defaults import MIN_BLOCK_SIZE
 from torch_tensorrt.dynamo.backend.lowering import SUBSTITUTION_REGISTRY
 from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
 from torch.fx.graph_module import GraphModule
diff --git a/py/torch_tensorrt/dynamo/backend/test/test_backend_compiler.py b/py/torch_tensorrt/dynamo/backend/test/test_backend_compiler.py
@@ -4,7 +4,7 @@
 from copy import deepcopy
 from torch_tensorrt.dynamo import compile
 from utils import lower_graph_testing
-from torch_tensorrt.dynamo.common_utils.test_utils import DECIMALS_OF_AGREEMENT
+from torch_tensorrt.dynamo.common.test_utils import DECIMALS_OF_AGREEMENT
 
 
 class TestTRTModuleNextCompilation(TestCase):
diff --git a/py/torch_tensorrt/dynamo/backend/test/test_decompositions.py b/py/torch_tensorrt/dynamo/backend/test/test_decompositions.py
@@ -3,7 +3,7 @@
 from torch.testing._internal.common_utils import run_tests, TestCase
 import torch
 from torch_tensorrt.dynamo import compile
-from torch_tensorrt.dynamo.common_utils.test_utils import DECIMALS_OF_AGREEMENT
+from torch_tensorrt.dynamo.common.test_utils import DECIMALS_OF_AGREEMENT
 
 
 class TestLowering(TestCase):
diff --git a/py/torch_tensorrt/dynamo/common/__init__.py b/py/torch_tensorrt/dynamo/common/__init__.py
@@ -0,0 +1,4 @@
+from ._settings import CompilationSettings
+
+from .fx2trt import TRTInterpreter, TRTInterpreterResult
+from .input_tensor_spec import InputTensorSpec
diff --git a/py/torch_tensorrt/dynamo/common/_defaults.py b/py/torch_tensorrt/dynamo/common/_defaults.py
diff --git a/py/torch_tensorrt/dynamo/common/_settings.py b/py/torch_tensorrt/dynamo/common/_settings.py
@@ -2,7 +2,7 @@
 from typing import Optional, Sequence
 
 from torch_tensorrt.fx.utils import LowerPrecision
-from torch_tensorrt.dynamo.backend._defaults import (
+from ._defaults import (
     PRECISION,
     DEBUG,
     WORKSPACE_SIZE,
diff --git a/py/torch_tensorrt/dynamo/common/fx2trt.py b/py/torch_tensorrt/dynamo/common/fx2trt.py
diff --git a/py/torch_tensorrt/dynamo/common/input_tensor_spec.py b/py/torch_tensorrt/dynamo/common/input_tensor_spec.py
diff --git a/py/torch_tensorrt/dynamo/common/test_utils.py b/py/torch_tensorrt/dynamo/common/test_utils.py
diff --git a/py/torch_tensorrt/dynamo/common_utils/__init__.py b/py/torch_tensorrt/dynamo/common_utils/__init__.py
diff --git a/py/torch_tensorrt/dynamo/fx_ts_compat/__init__.py b/py/torch_tensorrt/dynamo/fx_ts_compat/__init__.py
@@ -6,8 +6,6 @@
     NO_IMPLICIT_BATCH_DIM_SUPPORT,
     tensorrt_converter,
 )
-from .fx2trt import TRTInterpreter, TRTInterpreterResult  # noqa
-from .input_tensor_spec import InputTensorSpec  # noqa
 from .lower_setting import LowerSetting  # noqa
 from .lower import compile  # usort: skip  #noqa
 
diff --git a/py/torch_tensorrt/dynamo/fx_ts_compat/lower.py b/py/torch_tensorrt/dynamo/fx_ts_compat/lower.py
@@ -10,7 +10,7 @@
 import torch_tensorrt.fx.tracer.dispatch_tracer.aten_tracer as aten_tracer
 from torch.fx.passes.splitter_base import SplitResult
 
-from .fx2trt import TRTInterpreter, TRTInterpreterResult
+from torch_tensorrt.dynamo.common import TRTInterpreter, TRTInterpreterResult
 from .lower_setting import LowerSetting
 from .passes.lower_pass_manager_builder import LowerPassManagerBuilder
 from .passes.pass_utils import PassFunc, validate_inference
@@ -21,6 +21,17 @@
 from torch_tensorrt.fx.trt_module import TRTModule
 from torch_tensorrt.fx.utils import LowerPrecision
 from torch_tensorrt._Device import Device
+from torch_tensorrt.dynamo.common._defaults import (
+    PRECISION,
+    DEBUG,
+    WORKSPACE_SIZE,
+    MIN_BLOCK_SIZE,
+    PASS_THROUGH_BUILD_FAILURES,
+    MAX_AUX_STREAMS,
+    VERSION_COMPATIBLE,
+    OPTIMIZATION_LEVEL,
+    USE_EXPERIMENTAL_RT,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -34,21 +45,25 @@ def compile(
     disable_tf32=False,
     sparse_weights=False,
     enabled_precisions=set(),
-    min_block_size: int = 3,
-    workspace_size=0,
+    min_block_size: int = MIN_BLOCK_SIZE,
+    workspace_size=WORKSPACE_SIZE,
     dla_sram_size=1048576,
     dla_local_dram_size=1073741824,
     dla_global_dram_size=536870912,
     calibrator=None,
     truncate_long_and_double=False,
     require_full_compilation=False,
-    debug=False,
+    explicit_batch_dimension=False,
+    debug=DEBUG,
     refit=False,
     timing_cache_prefix="",
     save_timing_cache=False,
     cuda_graph_batch_size=-1,
     is_aten=False,
-    use_experimental_fx_rt=False,
+    use_experimental_rt=USE_EXPERIMENTAL_RT,
+    max_aux_streams=MAX_AUX_STREAMS,
+    version_compatible=VERSION_COMPATIBLE,
+    optimization_level=OPTIMIZATION_LEVEL,
     num_avg_timing_iters=1,
     torch_executed_ops=[],
     torch_executed_modules=[],
@@ -67,11 +82,14 @@ def compile(
         timing_cache_prefix: Timing cache file name for timing cache used by fx2trt.
         save_timing_cache: Update timing cache with current timing cache data if set to True.
         cuda_graph_batch_size: Cuda graph batch size, default to be -1.
-        use_experimental_fx_rt: Uses the next generation TRTModule which supports both Python and TorchScript based execution (including in C++).
+        use_experimental_rt: Uses the next generation TRTModule which supports both Python and TorchScript based execution (including in C++).
+        max_aux_streams: max number of aux stream to use
+        version_compatible: enable version compatible feature
+        optimization_level: builder optimization level
     Returns:
         A torch.nn.Module lowered by TensorRT.
     """
-    if use_experimental_fx_rt and not explicit_batch_dimension:
+    if use_experimental_rt and not explicit_batch_dimension:
         raise ValueError(
             "The experimental unifed runtime only supports explicit batch. Please make sure to set explicit_batch_dimension=True when use_experimental_fx_rt=True"
         )
@@ -122,7 +140,11 @@ def compile(
         save_timing_cache=save_timing_cache,
         cuda_graph_batch_size=cuda_graph_batch_size,
         is_aten=is_aten,
-        use_experimental_rt=use_experimental_fx_rt,
+        use_experimental_rt=use_experimental_rt,
+        max_aux_streams=max_aux_streams,
+        version_compatible=version_compatible,
+        optimization_level=optimization_level,
+        use_experimental_rt=use_experimental_rt,
     )
     lowerer = Lowerer.create(lower_setting=lower_setting)
     return lowerer(module, inputs)
diff --git a/py/torch_tensorrt/dynamo/fx_ts_compat/lower_setting.py b/py/torch_tensorrt/dynamo/fx_ts_compat/lower_setting.py
@@ -4,7 +4,7 @@
 from torch import nn
 from torch.fx.passes.pass_manager import PassManager
 
-from .input_tensor_spec import InputTensorSpec
+from torch_tensorrt.dynamo.common import InputTensorSpec
 from torch_tensorrt.fx.passes.lower_basic_pass import (
     fuse_permute_linear,
     fuse_permute_matmul,
diff --git a/py/torch_tensorrt/dynamo/fx_ts_compat/passes/lower_pass_manager_builder.py b/py/torch_tensorrt/dynamo/fx_ts_compat/passes/lower_pass_manager_builder.py
@@ -10,7 +10,7 @@
 from torch.fx.passes.splitter_base import generate_inputs_for_submodules, SplitResult
 from torch_tensorrt.fx.utils import LowerPrecision
 from torch_tensorrt import _Input
-from ..input_tensor_spec import InputTensorSpec
+from torch_tensorrt.dynamo.common import InputTensorSpec
 
 from ..lower_setting import LowerSetting
 from torch_tensorrt.fx.observer import Observer
diff --git a/py/torch_tensorrt/dynamo/fx_ts_compat/test/core/test_input_tensor_spec.py b/py/torch_tensorrt/dynamo/fx_ts_compat/test/core/test_input_tensor_spec.py
@@ -5,7 +5,7 @@
 import torch
 import torch_tensorrt
 from torch.testing._internal.common_utils import run_tests, TestCase
-from torch_tensorrt.dynamo.fx_ts_compat import InputTensorSpec, LowerSetting
+from torch_tensorrt.dynamo.common import InputTensorSpec
 
 
 class TestTRTModule(TestCase):
diff --git a/py/torch_tensorrt/dynamo/fx_ts_compat/tools/common_fx2trt.py b/py/torch_tensorrt/dynamo/fx_ts_compat/tools/common_fx2trt.py
@@ -13,7 +13,7 @@
 from torch.fx.passes import shape_prop
 from torch.fx.passes.infra.pass_base import PassResult
 from torch.testing._internal.common_utils import TestCase
-from torch_tensorrt.dynamo.fx_ts_compat import InputTensorSpec, TRTInterpreter
+from torch_tensorrt.dynamo.common import InputTensorSpec, TRTInterpreter
 from torch_tensorrt.fx.passes.lower_basic_pass_aten import (
     compose_bmm,
     compose_chunk,

Original file line number	Diff line number	Diff line change
`@@ -5,7 +5,7 @@`
`5`	`5`	`from dataclasses import replace, fields`
`6`	`6`	`import torch._dynamo as td`
`7`	`7`
`8`		`-from torch_tensorrt.dynamo.backend._settings import CompilationSettings`
	`8`	`+from torch_tensorrt.dynamo.common import CompilationSettings`
`9`	`9`	`from torch_tensorrt.dynamo.backend.lowering._decompositions import (`
`10`	`10`	`get_decompositions,`
`11`	`11`	`)`