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,14 +2,15 @@
 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
 
 # @manual=//deeplearning/trt/python:py_tensorrt
 import tensorrt as trt
 import torch
 import torch.fx
+from torch.fx.node import Target
 from torch._ops import OpOverload
 from torch.fx.node import _get_qualified_name
 from torch.fx.passes.shape_prop import TensorMetadata
@@ -42,12 +43,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
@@ -60,12 +62,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: Optional[List] = None
+        self.optimization_profiles: List[trt.IOptimizationProfile] = []
         self.input_specs = input_specs
         self.input_specs_iter = 0
         self._cur_node_name: Optional[str] = None
@@ -78,37 +81,37 @@ 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.target):
-                missing_converter.add(f"{node.op} {_get_qualified_name(node.target)}")
+                missing_converters.add(f"{node.op} {_get_qualified_name(node.target)}")
             elif node.op == "call_method" and not CONVERTERS.get(node.target):
-                missing_converter.add(f"{node.op} torch.Tensor.{node.target}")
+                missing_converters.add(f"{node.op} torch.Tensor.{node.target}")
             elif node.op == "call_module":
                 submod = self.fetch_attr(node.target)
                 submod_type = getattr(submod, "_base_class_origin", type(submod))
                 if not CONVERTERS.get(submod_type):
-                    missing_converter.add(f"{node.op} {torch.typename(submod_type)}")
+                    missing_converters.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.
@@ -196,7 +199,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)
 
@@ -224,55 +227,63 @@ 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)
         # add "_itensor_to_tensor_meta"
         kwargs = dict(n.kwargs)
         kwargs["_itensor_to_tensor_meta"] = self._itensor_to_tensor_meta
         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"]
-            self.optimization_profiles[0].set_shape(
+            # TODO: Does not support disjoint optimization profiles?
+            self.optimization_profiles.append(trt.IOptimizationProfile().set_shape(
                 target, [min_shape, opt_shape, max_shape]
-            )
+            ))
             assert len(min_shape) == len(opt_shape) == len(max_shape)
             for i in range(len(min_shape)):
                 if min_shape[i] == opt_shape[i] == max_shape[i]:
                     shape.append(min_shape[i])
                 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))
@@ -286,17 +297,17 @@ 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:
         converter = CONVERTERS.get(target)
         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(target)
 
@@ -308,7 +319,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,2 +1,2 @@
-from .trt_interpreter import *
+from .TRTInterpreter import *
 from .conversion import *

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