chore: ready to start review

Signed-off-by: Naren Dasan <naren@narendasan.com>
Signed-off-by: Naren Dasan <narens@nvidia.com>

Author: narendasan

.pre-commit-config.yaml

@@ -5,6 +5,7 @@ repos:
     hooks:
       - id: check-yaml
       - id: trailing-whitespace
+        exclude: ^docs
       - id: check-added-large-files
         args:
           - --maxkb=1000
@@ -13,11 +14,7 @@ repos:
       - id: mixed-line-ending
         args:
           - --fix=lf
-  - repo: https://github.com/psf/black
-    rev: 23.7.0
-    hooks:
-      - id: black
-        exclude: ^examples/custom_converters/elu_converter/setup.py
+        exclude: ^docs
   - repo: https://github.com/pre-commit/mirrors-clang-format
     rev: v16.0.6
     hooks:
@@ -30,21 +27,26 @@ repos:
         args:
           - --warnings=all
     -   id: buildifier-lint
-  - repo: https://github.com/astral-sh/ruff-pre-commit
-    # Ruff version.
-    rev: v0.0.278
-    hooks:
-      - id: ruff
   - repo: https://github.com/abravalheri/validate-pyproject
     rev: v0.13
     hooks:
       - id: validate-pyproject
+        python_version: "3.11"
   - repo: https://github.com/pre-commit/mirrors-mypy
     rev: 'v1.4.1'
     hooks:
     -   id: mypy
-        exclude: "^py/torch_tensorrt/fx|^examples|^tests|^tools"
-        python_version: "3.11"
+        exclude: "^py/torch_tensorrt/fx|^examples|^tests|^tools|^docs|noxfile.py|setup.py|versions.py"
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    # Ruff version.
+    rev: v0.0.278
+    hooks:
+      - id: ruff
+  - repo: https://github.com/psf/black
+    rev: 23.7.0
+    hooks:
+      - id: black
+        exclude: ^examples/custom_converters/elu_converter/setup.py|^docs
   - repo: local
     hooks:
     -   id: dont-commit-upstream

core/conversion/evaluators/aten.cpp

@@ -103,7 +103,12 @@ DEFINE_ARITHMATIC_TWO_INPUT_EVALUATOR(
         "aten::pow.float_int(float a, int b) -> (float)",
     }));
 
-DEFINE_TWO_INPUT_SIMPLE_EVALUATOR(and, "aten::__and__", a&& b, bool, std::set<std::string>({"aten::__and__(int a, int b) -> (bool)", "aten::__and__.bool(bool a, bool b) -> (bool)"}));
+DEFINE_TWO_INPUT_SIMPLE_EVALUATOR(
+    and,
+    "aten::__and__",
+    a&& b,
+    bool,
+    std::set<std::string>({"aten::__and__(int a, int b) -> (bool)", "aten::__and__.bool(bool a, bool b) -> (bool)"}));
 DEFINE_TWO_INPUT_SIMPLE_EVALUATOR(or, "aten::__or__", a || b, bool, {"aten::__or__(int a, int b) -> (bool)"});
 DEFINE_TWO_INPUT_SIMPLE_EVALUATOR(
     xor,

docs/_downloads/e1ef5a42560a98a132f56a79d0b66f79/dynamo_compile_advanced_usage.py

@@ -16,6 +16,7 @@
 
 # %%
 
+
 # We begin by defining a model
 class Model(torch.nn.Module):
     def __init__(self) -> None:

docs/v1.4.0/_downloads/e1ef5a42560a98a132f56a79d0b66f79/dynamo_compile_advanced_usage.py

@@ -16,6 +16,7 @@
 
 # %%
 
+
 # We begin by defining a model
 class Model(torch.nn.Module):
     def __init__(self) -> None:

py/torch_tensorrt/__init__.py

@@ -82,10 +82,10 @@ def _find_lib(name: str, paths: List[str]) -> str:
 
 import torch
 from torch_tensorrt._compile import *  # noqa: F403
+from torch_tensorrt._Device import Device  # noqa: F401
 from torch_tensorrt._enums import *  # noqa: F403
-from torch_tensorrt._util import *  # noqa: F403
-from torch_tensorrt._Input import Input # noqa: F401
-from torch_tensorrt._Device import Device # noqa: F401
+from torch_tensorrt._Input import Input  # noqa: F401
+from torch_tensorrt._utils import *  # noqa: F403
 
 if version.parse(sanitized_torch_version()) >= version.parse("2.1.dev"):
     from torch_tensorrt import dynamo  # noqa: F401

py/torch_tensorrt/_compile.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Any, Callable, List, Optional, Set, TypeGuard
+from typing import Any, Callable, List, Optional, Set, TypeGuard, Sequence
 
 import torch
 import torch.fx
@@ -15,13 +15,13 @@
 
 
 def _non_fx_input_interface(
-    inputs: List[Input | torch.Tensor | InputTensorSpec],
+    inputs: Sequence[Input | torch.Tensor | InputTensorSpec],
 ) -> TypeGuard[List[Input | torch.Tensor]]:
     return all(isinstance(i, torch.Tensor | Input) for i in inputs)
 
 
 def _fx_input_interface(
-    inputs: List[Input | torch.Tensor | InputTensorSpec],
+    inputs: Sequence[Input | torch.Tensor | InputTensorSpec],
 ) -> TypeGuard[List[InputTensorSpec | torch.Tensor]]:
     return all(isinstance(i, torch.Tensor | InputTensorSpec) for i in inputs)
 
@@ -97,7 +97,7 @@ def _get_target_ir(module_type: _ModuleType, ir: str) -> _IRType:
 def compile(
     module: Any,
     ir: str = "default",
-    inputs: Optional[List[Input | torch.Tensor | InputTensorSpec]] = None,
+    inputs: Optional[Sequence[Input | torch.Tensor | InputTensorSpec]] = None,
     enabled_precisions: Optional[Set[torch.dtype | dtype]] = None,
     **kwargs: Any,
 ) -> (
@@ -201,7 +201,7 @@ def compile(
         raise RuntimeError("Module is an unknown format or the ir requested is unknown")
 
 
-def torch_compile(module: torch.nn.Module, **kwargs: Any) -> Callable[..., Any]:
+def torch_compile(module: torch.nn.Module, **kwargs: Any) -> Any:
     """
     Returns a boxed model which is the output of torch.compile.
     This does not compile the model to TRT. Execute this model on
@@ -216,8 +216,8 @@ def torch_compile(module: torch.nn.Module, **kwargs: Any) -> Callable[..., Any]:
 
 def convert_method_to_trt_engine(
     module: Any,
-    inputs: List[Input | torch.Tensor],
-    method_name: str,
+    method_name: str = "forward",
+    inputs: Optional[Sequence[Input | torch.Tensor]] = None,
     ir: str = "default",
     enabled_precisions: Optional[Set[torch.dtype | dtype]] = None,
     **kwargs: Any,

py/torch_tensorrt/_enums.py

@@ -1,2 +1,3 @@
-from torch_tensorrt._C import dtype, EngineCapability, TensorFormat  # noqa: F401
+from torch_tensorrt._C import EngineCapability, TensorFormat, dtype  # noqa: F401
+
 from tensorrt import DeviceType  # noqa: F401

py/torch_tensorrt/_util.py renamed to py/torch_tensorrt/_utils.py

@@ -1,5 +1,8 @@
+from typing import Any
+
 import torch
-from torch_tensorrt import _C, __version__
+from torch_tensorrt._version import __version__
+from torch_tensorrt import _C
 
 
 def dump_build_info() -> None:
@@ -30,7 +33,7 @@ def set_device(gpu_id: int) -> None:
     _C.set_device(gpu_id)
 
 
-def sanitized_torch_version() -> str:
+def sanitized_torch_version() -> Any:
     return (
         torch.__version__
         if ".nv" not in torch.__version__

py/torch_tensorrt/dynamo/_SourceIR.py

@@ -9,7 +9,7 @@ class SourceIR(Enum):
     TORCHTRT_LOWERED = auto()
     UNKNOWN = auto()
 
-    def __str__(self):
+    def __str__(self) -> str:
         if self == SourceIR.NN:
             return "nn"
         elif self == SourceIR.ACC:

py/torch_tensorrt/dynamo/__init__.py

@@ -1,13 +1,11 @@
 from packaging import version
-from torch_tensorrt._util import sanitized_torch_version
+from torch_tensorrt._utils import sanitized_torch_version
 
 if version.parse(sanitized_torch_version()) >= version.parse("2.1.dev"):
     from ._settings import *  # noqa: F403
-    from .conversion import * # noqa: F403
-    from .aten_tracer import trace # noqa: F403
-    from .conversion.converter_registry import (
-        DYNAMO_CONVERTERS, # noqa: F403
-        dynamo_tensorrt_converter, # noqa: F403
-    )
-    from .compile import compile # noqa: F403
-    from ._SourceIR import SourceIR # noqa: F403
+    from ._SourceIR import SourceIR  # noqa: F403
+    from .aten_tracer import trace  # noqa: F403
+    from .compile import compile  # noqa: F403
+    from .conversion import *  # noqa: F403
+    from .conversion.converter_registry import DYNAMO_CONVERTERS  # noqa: F403
+    from .conversion.converter_registry import dynamo_tensorrt_converter  # noqa: F403

py/torch_tensorrt/dynamo/_settings.py

@@ -9,13 +9,10 @@
     OPTIMIZATION_LEVEL,
     PASS_THROUGH_BUILD_FAILURES,
     PRECISION,
-    USE_PYTHON_RUNTIME,
-<<<<<<< HEAD
     TRUNCATE_LONG_AND_DOUBLE,
-=======
+    USE_PYTHON_RUNTIME,
     VERSION_COMPATIBLE,
     WORKSPACE_SIZE,
->>>>>>> e39abb60d (chore: adding isort to pre-commit)
 )
 
 

py/torch_tensorrt/dynamo/aten_tracer.py

@@ -1,11 +1,10 @@
 import copy
 import sys
 from contextlib import contextmanager
-from typing import Any, Callable, Dict, Generator, List, Optional, Set, Tuple, Union
+from typing import Any, Callable, Dict, Generator, List, Optional, Tuple
 
 import torch
 import torch._dynamo as torchdynamo
-from torch import _guards
 from torch.fx.passes.infra.pass_base import PassResult
 from torch_tensorrt.dynamo.utils import req_torch_version
 from torch_tensorrt.fx.passes.lower_basic_pass_aten import (
@@ -23,11 +22,7 @@
 )
 from typing_extensions import TypeAlias
 
-Value: TypeAlias = Union[
-    Tuple["Value", ...],
-    List["Value"],
-    Dict[str, "Value"],
-]
+Value: TypeAlias = Tuple["Value", ...] | List["Value"] | Dict[str, "Value"]
 
 
 class DynamoConfig:
@@ -96,7 +91,7 @@ def dynamo_trace(
     aten_graph: bool,
     tracing_mode: str = "real",
     dynamo_config: Optional[DynamoConfig] = None,
-) -> Tuple[torch.fx.GraphModule, Set[_guards.Guard]]:
+) -> Any:  # Tuple[torch.fx.GraphModule, Set[_guards.Guard]]:
     """
     TODO: Once we fully migrate to torchdynamo frontend, we will remove
     this config option alltogether.  For now, it helps with quick
@@ -150,11 +145,11 @@ def trace(
     fx_module, __package__ = dynamo_trace(model, inputs, True, "symbolic")
     print(fx_module.graph)
     for passes in passes_list:
-        pr: PassResult = passes(fx_module)  # type: ignore[assignment] #The type hints in fx are wrong
+        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  # type: ignore[no-any-return]
+    return fx_module

py/torch_tensorrt/dynamo/backend/__init__.py

@@ -1 +1 @@
-from .backends import torch_tensorrt_backend # noqa: F401
+from .backends import torch_tensorrt_backend  # noqa: F401

py/torch_tensorrt/dynamo/backend/backends.py

@@ -4,30 +4,12 @@
 
 import torch
 import torch._dynamo as td
-<<<<<<< HEAD
-
+from torch._functorch.aot_autograd import aot_module_simplified, make_boxed_compiler
 from torch_tensorrt.dynamo import CompilationSettings
-from torch_tensorrt.dynamo.lowering._decompositions import (
-    get_decompositions,
-)
-from torch_tensorrt.dynamo.lowering._pre_aot_lowering import (
-    pre_aot_substitutions,
-)
-from torch_tensorrt.dynamo.lowering._partition import (
-    partition,
-    get_submod_inputs,
-)
-from torch_tensorrt.dynamo.utils import parse_dynamo_kwargs
 from torch_tensorrt.dynamo.conversion import (
     convert_module,
     repair_long_or_double_inputs,
 )
-
-=======
->>>>>>> e39abb60d (chore: adding isort to pre-commit)
-from torch._functorch.aot_autograd import aot_module_simplified, make_boxed_compiler
-from torch_tensorrt.dynamo import CompilationSettings
-from torch_tensorrt.dynamo.conversion import convert_module
 from torch_tensorrt.dynamo.lowering._decompositions import get_decompositions
 from torch_tensorrt.dynamo.lowering._partition import get_submod_inputs, partition
 from torch_tensorrt.dynamo.lowering._pre_aot_lowering import pre_aot_substitutions
@@ -64,7 +46,7 @@ def aot_torch_tensorrt_aten_backend(
     return aot_module_simplified(
         gm,
         sample_inputs,
-        fw_compiler=make_boxed_compiler(custom_backend),  # type: ignore[no-untyped-call]
+        fw_compiler=make_boxed_compiler(custom_backend),
         decompositions=get_decompositions(),
     )
 

py/torch_tensorrt/dynamo/compile.py

@@ -6,7 +6,11 @@
 import torch_tensorrt
 from torch.fx.passes.pass_manager import PassManager
 from torch.fx.passes.splitter_base import SplitResult
-from torch_tensorrt import Device, EngineCapability
+from torch_tensorrt.fx.tools.trt_splitter import TRTSplitter, TRTSplitterSetting
+from torch_tensorrt._Device import Device
+from torch_tensorrt._enums import (
+    EngineCapability,
+)  # TODO: Should probabably be the TRT EngineCapability Enum
 from torch_tensorrt.dynamo import CompilationSettings
 from torch_tensorrt.dynamo._defaults import (
     DEBUG,
@@ -15,32 +19,18 @@
     OPTIMIZATION_LEVEL,
     PASS_THROUGH_BUILD_FAILURES,
     PRECISION,
+    TRUNCATE_LONG_AND_DOUBLE,
     USE_PYTHON_RUNTIME,
     VERSION_COMPATIBLE,
     WORKSPACE_SIZE,
 )
 from torch_tensorrt.dynamo.backend.backends import _compile_module
 from torch_tensorrt.dynamo.conversion import convert_module
-<<<<<<< HEAD
-
-from torch_tensorrt.dynamo._defaults import (
-    PRECISION,
-    DEBUG,
-    WORKSPACE_SIZE,
-    MIN_BLOCK_SIZE,
-    PASS_THROUGH_BUILD_FAILURES,
-    MAX_AUX_STREAMS,
-    VERSION_COMPATIBLE,
-    OPTIMIZATION_LEVEL,
-    USE_PYTHON_RUNTIME,
-    TRUNCATE_LONG_AND_DOUBLE,
+from torch_tensorrt.dynamo.lowering._fusers import (
+    fuse_permute_linear,
+    fuse_permute_matmul,
 )
-
-=======
-from torch_tensorrt.dynamo.lowering import fuse_permute_linear, fuse_permute_matmul
 from torch_tensorrt.dynamo.utils import prepare_device, prepare_inputs
-from torch_tensorrt.fx.tools.trt_splitter import TRTSplitter, TRTSplitterSetting
->>>>>>> e39abb60d (chore: adding isort to pre-commit)
 
 logger = logging.getLogger(__name__)
 
@@ -89,7 +79,7 @@ def compile(
     if not isinstance(inputs, collections.abc.Sequence):
         inputs = [inputs]
 
-    torchtrt_inputs, torch_inputs = prepare_inputs(inputs, prepare_device(device))
+    _, torch_inputs = prepare_inputs(inputs, prepare_device(device))
 
     if (
         torch.float16 in enabled_precisions
@@ -125,7 +115,7 @@ def compile(
         "truncate_long_and_double": truncate_long_and_double,
     }
 
-    settings = CompilationSettings(**compilation_options)  # type: ignore[arg-type]
+    settings = CompilationSettings(**compilation_options)
     if kwargs.get("use_capability_partitioner", None):
         model = lower_model(gm, torch_inputs)
         return _compile_module(model, torch_inputs, settings)
@@ -163,7 +153,7 @@ def lower_model_using_trt_splitter(
 ) -> SplitResult:
     # Perform basic lowering
     model = lower_model(model, inputs)
-    splitter_setting = TRTSplitterSetting()  # type: ignore[no-untyped-call]
+    splitter_setting = TRTSplitterSetting()
     splitter_setting.use_implicit_batch_dim = False
     splitter_setting.min_acc_module_size = 1
     splitter_setting.use_experimental_rt = False
@@ -177,7 +167,7 @@ def lower_model_using_trt_splitter(
 def lower_model(
     model: torch.nn.Module, inputs: Any, **kwargs: Any
 ) -> torch.fx.GraphModule:
-    graph_optimization_pm = PassManager.build_from_passlist(  # type: ignore[no-untyped-call]
+    graph_optimization_pm = PassManager.build_from_passlist(
         [fuse_permute_matmul, fuse_permute_linear]
     )
     lowered_model: torch.fx.GraphModule = graph_optimization_pm(model)