feat: Add _to_copy, operator.get and clone

- Add ATen converters for key operators in the pipeline of multiple
models
- Add robust testing and patch issues in interpreter
- Add evaluator and casting utilities to the converter utils

Author: gs-olive

py/torch_tensorrt/dynamo/conversion/_TRTInterpreter.py

@@ -27,6 +27,10 @@
 ] = Observer("TRT_INTERPRETER_CALL_PRE_OBSERVER")
 
 
+class UnsupportedOperatorException(RuntimeError):
+    pass
+
+
 class TRTInterpreterResult(NamedTuple):
     engine: Any
     input_names: Sequence[str]
@@ -301,7 +305,7 @@ def call_module(
         converter = CONVERTERS.get(self._cur_node)
 
         if not converter:
-            raise RuntimeError(
+            raise UnsupportedOperatorException(
                 f"Conversion of module of type {submod_type} not currently supported!"
             )
 
@@ -312,7 +316,7 @@ 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(
+            raise UnsupportedOperatorException(
                 f"Conversion of function {torch.typename(target)} not currently supported!"
             )
 
@@ -324,7 +328,7 @@ def call_method(self, target: str, args: Any, kwargs: Any) -> Any:
         converter = CONVERTERS.get(self._cur_node)
 
         if not converter:
-            raise RuntimeError(
+            raise UnsupportedOperatorException(
                 f"Conversion of method {target} not currently supported!"
             )
 

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -1,6 +1,8 @@
 import logging
+import operator
 from typing import Any, Dict, Optional, Sequence, Tuple, Union
 
+import tensorrt as trt
 import torch
 from torch.fx.node import Argument, Node, Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
@@ -12,8 +14,6 @@
 from torch_tensorrt.fx.converters import acc_ops_converters
 from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
 
-import tensorrt as trt
-
 from .converter_registry import dynamo_tensorrt_converter
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
@@ -76,13 +76,13 @@ def aten_ops_div(
         kwargs_new["input"].dtype == trt.int8 or kwargs_new["input"].dtype == trt.int32
     ):
         kwargs_new["input"] = cast_trt_tensor(
-            network, kwargs_new["input"], trt.float32, name
+            network, kwargs_new["input"], trt.float32, name, target
         )
     elif isinstance(args[1], TRTTensor) and (
         kwargs_new["other"].dtype == trt.int8 or kwargs_new["other"].dtype == trt.int32
     ):
         kwargs_new["other"] = cast_trt_tensor(
-            network, kwargs_new["other"], trt.float32, name
+            network, kwargs_new["other"], trt.float32, name, target
         )
     rounding_mode = kwargs.get("rounding_mode")
     if rounding_mode is None:
@@ -101,7 +101,7 @@ def aten_ops_div(
         )
 
 
-def embedding_param_validator(embedding_node: Node):
+def embedding_param_validator(embedding_node: Node) -> bool:
     scale_grad_by_freq = args_bounds_check(embedding_node.args, 3)
     sparse = args_bounds_check(embedding_node.args, 4)
 
@@ -365,3 +365,77 @@ def aten_ops_permute(
         args[0],
         args[1],
     )
+
+
+def to_copy_dtype_validator(to_copy_node: Node) -> bool:
+    allowed_casts = {torch.float, torch.int32, torch.bool, torch.int8, torch.float16}
+
+    # Validate input node has convertible kwargs
+    if "dtype" in to_copy_node.kwargs:
+        if to_copy_node.kwargs["dtype"] in allowed_casts:
+            return True
+        else:
+            _LOGGER.debug(
+                f"_to_copy converter rejected node {to_copy_node} with dtype {to_copy_node.kwargs['dtype']}"
+            )
+            return False
+    else:
+        _LOGGER.debug(
+            f"_to_copy converter rejected node {to_copy_node} with kwargs {to_copy_node.kwargs}"
+        )
+        return False
+
+
+@dynamo_tensorrt_converter(
+    torch.ops.aten._to_copy.default, capability_validator=to_copy_dtype_validator
+)
+def aten_ops_to_copy_dtype(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.cast.to_copy(
+        network,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        kwargs["dtype"],
+    )
+
+
+@dynamo_tensorrt_converter(operator.getitem)
+def operator_getitem(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.evaluators.getitem(
+        network,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+    )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.clone.default)
+def aten_ops_clone(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.evaluators.clone(
+        network,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+    )

py/torch_tensorrt/dynamo/conversion/converter_registry.py

@@ -66,7 +66,7 @@ def dynamo_tensorrt_converter(
     enabled: bool = True,
     capability_validator: Optional[Callable[[Node], bool]] = None,
     priority: ConverterPriority = ConverterPriority.STANDARD,
-) -> Callable[[Any], Any]:
+) -> Callable[[Any], Union[TRTTensor, Sequence[TRTTensor]]]:
     """Decorator for Dynamo TensorRT Converter
 
     Registers the decorated function in the DYNAMO_ATEN_CONVERTERS registry

py/torch_tensorrt/dynamo/conversion/converter_utils.py

@@ -1,15 +1,18 @@
 import logging
 import re
-from typing import List
+from typing import List, Optional
 
 import tensorrt as trt
 import torch
+from torch.fx.node import Target, _get_qualified_name
 from torch_tensorrt.fx.converters.converter_utils import (
     Frameworks,
     unified_dtype_converter,
 )
 from torch_tensorrt.fx.types import TRTDataType, TRTNetwork, TRTTensor
 
+from .._SourceIR import SourceIR
+
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
 
@@ -71,24 +74,35 @@ def cast_trt_tensor(
     input_val: TRTTensor,
     dtype: TRTDataType,
     name: str,
+    target: Target = "",
+    source_ir: Optional[SourceIR] = None,
 ) -> TRTTensor:
     """
     Given a TRT Tensor, convert that Tensor to the specified dtype
     Adds an Identity layer to the network which performs the conversion
     Args:
         network (TRTNetwork): A TensorRT network
         input_val (TRTTensor): A TRT Tensor to cast to a new data type
-        dtype (TRTDataType): The TRTDataType to cast the input Tensor to
+        dtype (TRTDataType, torch.dtype, np.dtype): The data type to cast the input Tensor to
         name (str): Name of the calling layer
+        target (Target): Target of calling node
+        source_ir (SourceIR): SourceIR of calling converter
     Returns:
         A TensorRT ITensor which has been casted to the specified dtype
     """
     trt_dtype = unified_dtype_converter(dtype, Frameworks.TRT)
 
     if input_val.dtype != trt_dtype:
+        source_ir = source_ir if source_ir is not None else SourceIR.UNKNOWN
+        target_name = (
+            f"{source_ir}_ops{'.' + target if target else ''}"
+            if (isinstance(target, str))
+            else f"{source_ir}_ops.{_get_qualified_name(target)}"
+        )
+
         identity_layer = network.add_identity(input_val)
         identity_layer.set_output_type(0, trt_dtype)
-        identity_layer.name = f"Cast ITensor {input_val.name} from {input_val.dtype} to {trt_dtype} - {name}"
+        identity_layer.name = f"Cast ITensor {input_val.name} from {input_val.dtype} to {trt_dtype} -{name}-[{target_name}]-[{name}]"
         return identity_layer.get_output(0)
     else:
         return input_val

py/torch_tensorrt/dynamo/conversion/impl/__init__.py

@@ -2,9 +2,11 @@
 
 from . import (
     activation,
+    cast,
     condition,
     elementwise,
     embedding,
+    evaluators,
     matmul,
     normalization,
     permutation,

py/torch_tensorrt/dynamo/conversion/impl/cast.py

@@ -0,0 +1,23 @@
+from typing import Optional
+
+from torch.fx.node import Target
+from torch_tensorrt.dynamo._SourceIR import SourceIR
+from torch_tensorrt.dynamo.conversion.converter_utils import cast_trt_tensor
+from torch_tensorrt.fx.types import TRTDataType, TRTNetwork, TRTTensor
+
+
+def to_copy(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+    dtype: TRTDataType,
+) -> TRTTensor:
+    if not isinstance(input, TRTTensor):
+        raise RuntimeError(
+            f"to_copy received input {input} that is not a TensorRT ITensor"
+        )
+
+    casted_tensor = cast_trt_tensor(network, input, dtype, name, target, source_ir)
+    return casted_tensor

py/torch_tensorrt/dynamo/conversion/impl/elementwise/base.py

@@ -2,6 +2,7 @@
 import warnings
 from typing import Any, Callable, Optional, Union
 
+import tensorrt as trt
 import torch
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
@@ -15,8 +16,6 @@
 from torch_tensorrt.fx.types import TRTElementWiseOp, TRTNetwork, TRTTensor
 from torch_tensorrt.fx.utils import Frameworks, unified_dtype_converter
 
-import tensorrt as trt
-
 
 def get_python_op_from_trt_elementwise_op(
     trt_op: TRTElementWiseOp,
@@ -132,9 +131,13 @@ def convert_binary_elementwise(
     trt_promoted_type = unified_dtype_converter(promoted_type, Frameworks.TRT)
 
     if trt_promoted_type != lhs_val.dtype:
-        lhs_val = cast_trt_tensor(network, lhs_val, trt_promoted_type, name)
+        lhs_val = cast_trt_tensor(
+            network, lhs_val, trt_promoted_type, name, target, source_ir
+        )
     if trt_promoted_type != rhs_val.dtype:
-        rhs_val = cast_trt_tensor(network, rhs_val, trt_promoted_type, name)
+        rhs_val = cast_trt_tensor(
+            network, rhs_val, trt_promoted_type, name, target, source_ir
+        )
 
     # Check the limitation in the doc string.
     if network.has_implicit_batch_dimension:

py/torch_tensorrt/dynamo/conversion/impl/evaluators.py

@@ -0,0 +1,40 @@
+import logging
+import operator
+from typing import Optional, Sequence
+
+from torch.fx.node import Target
+from torch_tensorrt.dynamo._SourceIR import SourceIR
+from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
+
+LOGGER: logging.Logger = logging.getLogger(__name__)
+
+
+def getitem(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: Sequence[TRTTensor],
+    index: int,
+) -> TRTTensor:
+    LOGGER.debug(f"Evaluating getitem on object with name: {name}")
+
+    # Directly index the input sequence and return the value
+    return operator.getitem(input, index)
+
+
+def clone(
+    network: TRTNetwork,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input: TRTTensor,
+) -> TRTTensor:
+    if not isinstance(input, TRTTensor):
+        raise RuntimeError(
+            f"clone received input {input} that is not a TensorRT ITensor"
+        )
+
+    LOGGER.debug(f"Evaluating clone on object with name: {name}")
+
+    return input