chore(//py/torch_tensorrt/dynamo/conversion): mypy conforming

Signed-off-by: Naren Dasan <[email protected]>
Signed-off-by: Naren Dasan <[email protected]>

Author: narendasan

py/torch_tensorrt/dynamo/conversion/trt_interpreter.py renamed to py/torch_tensorrt/dynamo/conversion/TRTInterpreter.py

@@ -2,7 +2,7 @@
 import warnings
 from datetime import datetime
 from packaging import version
-from typing import Any, Callable, Dict, List, NamedTuple, Optional, Sequence
+from typing import Any, Callable, Dict, List, NamedTuple, Optional, Sequence, Set
 
 import numpy
 
@@ -41,12 +41,13 @@ def __init__(
         self,
         module: torch.fx.GraphModule,
         input_specs: List[Input],
-        logger_level=None,
-        output_dtypes=None,
+        logger_level: trt.ILogger.Severity = trt.ILogger.Severity.WARNING,
+        output_dtypes: Optional[List[torch.dtype]] = None,
     ):
         super().__init__(module)
 
-        self.logger = trt.Logger(logger_level or trt.Logger.WARNING)
+        # TODO: @narendasan replace with Torch-TensorRT Logger
+        self.logger = trt.Logger(logger_level)
         self.builder = trt.Builder(self.logger)
 
         flag = 0
@@ -59,12 +60,13 @@ def __init__(
 
         missing_ops = self.validate_conversion()
         if missing_ops:
+            # TODO: @narendasan make sure to set logging.captureWarnings(True)
             warnings.warn(
                 "Interpretation will fail due to missing operations \n"
                 + "\n".join(f"{i}" for i in missing_ops)
             )
 
-        self.optimization_profiles = (
+        self.optimization_profiles: Optional[List[trt.IOptimizationProfile]] = (
             [self.builder.create_optimization_profile()]
             if any(
                 input_spec.shape_mode == Input._ShapeMode.DYNAMIC
@@ -86,8 +88,8 @@ def __init__(
         # Data types for TRT Module output Tensors
         self.output_dtypes = output_dtypes
 
-    def validate_conversion(self):
-        missing_converter = set()
+    def validate_conversion(self) -> Set[str]:
+        missing_converters = set()
 
         for node in self.module.graph.nodes:
             if node.op == "call_function" and not CONVERTERS.get(node):
@@ -98,25 +100,25 @@ def validate_conversion(self):
                 submod = self.fetch_attr(node.target)
                 submod_type = getattr(submod, "_base_class_origin", type(submod))
                 if not CONVERTERS.get(node):
-                    missing_converter.add(f"{node.op} {torch.typename(submod_type)}")
+                    missing_converter.add(f"{node.op} {torch.typename(submod_type)}") # type: ignore[no-untyped-call]
 
-        return missing_converter
+        return missing_converters
 
     def run(
         self,
-        workspace_size=0,
-        precision=torch.float32,
-        sparse_weights=False,
-        disable_tf32=False,
-        force_fp32_output=False,
-        strict_type_constraints=False,
-        algorithm_selector=None,
-        timing_cache=None,
-        profiling_verbosity=None,
-        tactic_sources=None,
-        max_aux_streams=None,
-        version_compatible=False,
-        optimization_level=None,
+        workspace_size: int = 0,
+        precision: torch.dtype = torch.float32, # TODO: @peri044 Needs to be expanded to set
+        sparse_weights: bool = False,
+        disable_tf32: bool = False,
+        force_fp32_output: bool = False,
+        strict_type_constraints: bool = False,
+        algorithm_selector: Optional[trt.IAlgorithmSelector] = None,
+        timing_cache: Optional[trt.ITimingCache] =None,
+        profiling_verbosity: Optional[trt.ProfilingVerbosity] = None,
+        tactic_sources: Optional[int] = None,
+        max_aux_streams: Optional[int] = None,
+        version_compatible: bool = False,
+        optimization_level: Optional[int] = None,
     ) -> TRTInterpreterResult:
         """
         Build TensorRT engine with some configs.
@@ -204,7 +206,7 @@ def run(
         if strict_type_constraints:
             builder_config.set_flag(trt.BuilderFlag.STRICT_TYPES)
 
-        if self.optimization_profiles:
+        if len(self.optimization_profiles) > 0:
             for optimization_profile in self.optimization_profiles:
                 builder_config.add_optimization_profile(optimization_profile)
 
@@ -232,7 +234,7 @@ def run(
             engine, self._input_names, self._output_names, serialized_cache
         )
 
-    def run_node(self, n):
+    def run_node(self, n: torch.fx.Node) -> torch.fx.Node:
         self._cur_node_name = str(n)
         self._cur_node = n
         # add "_itensor_to_tensor_meta"
@@ -241,29 +243,31 @@ def run_node(self, n):
         n.kwargs = kwargs
 
         # run the node
-        trt_node = super().run_node(n)
+        trt_node: torch.fx.Node = super().run_node(n)
 
         # remove "_itensor_to_tensor_meta"
         kwargs = dict(n.kwargs)
         del kwargs["_itensor_to_tensor_meta"]
         n.kwargs = kwargs
 
         if isinstance(trt_node, trt.tensorrt.ITensor):
-            self._itensor_to_tensor_meta[trt_node] = n.meta.get("tensor_meta")
+            self._itensor_to_tensor_meta[trt_node] = n.meta.get("tensor_meta") #type: ignore[assignment]
 
         return trt_node
 
-    def placeholder(self, target, args, kwargs):
+    def placeholder(self, target: str, args: Any, kwargs: Any) -> trt.ITensor:
         self._input_names.append(target)
         current_input = self.input_specs[self.input_specs_iter]
         self.input_specs_iter += 1
         # Set optimization profile for dynamic input shape
-        shape = current_input.shape
+        shape = None
         if current_input.shape_mode == Input._ShapeMode.DYNAMIC:
+            assert isinstance(current_input.shape, dict)
             shape = []
             min_shape = current_input.shape["min_shape"]
             opt_shape = current_input.shape["opt_shape"]
             max_shape = current_input.shape["max_shape"]
+            # TODO: Does not support disjoint optimization profiles?
             self.optimization_profiles[0].set_shape(
                 target, min_shape, opt_shape, max_shape
             )
@@ -274,14 +278,20 @@ def placeholder(self, target, args, kwargs):
                 else:
                     # -1 to represent the dynamic dimension
                     shape.append(-1)
+        elif current_input.shape_mode == Input._ShapeMode.STATIC:
+            assert isinstance(current_input.shape, tuple)
+            shape = list(current_input.shape)
+        else:
+            raise RuntimeError(f"Unable to access shape spec for input: {target} (got: {current_input})")
+
 
         return self.network.add_input(
             name=target,
             shape=tuple(shape),
             dtype=unified_dtype_converter(current_input.torch_dtype, Frameworks.TRT),
         )
 
-    def call_module(self, target, args, kwargs):
+    def call_module(self, target: str, args: Any, kwargs: Any) -> Any: #Probably should be Tuple[trt.ITensor]? Case for Any?
         assert isinstance(target, str)
         submod = self.fetch_attr(target)
         submod_type = getattr(submod, "_base_class_origin", type(submod))
@@ -295,17 +305,18 @@ def call_module(self, target, args, kwargs):
         assert self._cur_node_name is not None
         return converter(self.network, submod, args, kwargs, self._cur_node_name)
 
-    def call_function(self, target, args, kwargs):
+    def call_function(self, target: str, args: Any, kwargs: Any) -> Any:
+        #TODO: Why is this stateful? We should be able to take in the inputs
         converter = CONVERTERS.get(self._cur_node)
         if not converter:
             raise RuntimeError(
-                f"Conversion of function {torch.typename(target)} not currently supported!"
+                f"Conversion of function {torch.typename(target)} not currently supported!" # type: ignore[no-untyped-call]
             )
 
         assert self._cur_node_name is not None
         return converter(self.network, target, args, kwargs, self._cur_node_name)
 
-    def call_method(self, target, args, kwargs):
+    def call_method(self, target: str, args: Any, kwargs: Any) -> Any:
         assert isinstance(target, str)
         converter = CONVERTERS.get(self._cur_node)
 
@@ -317,7 +328,7 @@ def call_method(self, target, args, kwargs):
         assert self._cur_node_name is not None
         return converter(self.network, target, args, kwargs, self._cur_node_name)
 
-    def output(self, target, args, kwargs):
+    def output(self, target: str, args: Any, kwargs: Any) -> None:
         assert len(args) == 1
         if isinstance(args[0], tuple):
             outputs = args[0]

py/torch_tensorrt/dynamo/conversion/__init__.py

@@ -1,4 +1,4 @@
 from .aten_ops_converters import *
-from .trt_interpreter import *
+from .TRTInterpreter import *
 from .conversion import *
 from .truncate_long_and_double import repair_long_or_double_inputs

py/torch_tensorrt/dynamo/conversion/conversion.py

@@ -1,7 +1,7 @@
 from typing import Sequence, Union
 import torch
 import io
-from torch_tensorrt.dynamo.runtime import _PythonTorchTRTModule
+from torch_tensorrt.dynamo.runtime import PythonTorchTensorRTModule, TorchTensorRTModule
 from torch_tensorrt.dynamo import CompilationSettings
 from torch_tensorrt import Input
 from torch_tensorrt.dynamo.conversion import TRTInterpreter
@@ -15,7 +15,7 @@ def convert_module(
     inputs: Sequence[torch.Tensor],
     settings: CompilationSettings = CompilationSettings(),
     name: str = "",
-):
+) -> Union[PythonTorchTensorRTModule, TorchTensorRTModule]:
     """Convert an FX module to a TRT module
     Args:
         module: FX GraphModule to convert