diff --git a/py/torch_tensorrt/dynamo/conversion/__init__.py b/py/torch_tensorrt/dynamo/conversion/__init__.py index 4536ff0e7b..9cbfff950e 100644 --- a/py/torch_tensorrt/dynamo/conversion/__init__.py +++ b/py/torch_tensorrt/dynamo/conversion/__init__.py @@ -1,4 +1,5 @@ from ._TRTInterpreter import * # noqa: F403 from .aten_ops_converters import * # noqa: F403 from .conversion import * # noqa: F403 +from .op_evaluators import * # noqa: F403 from .truncate_long_and_double import repair_long_or_double_inputs diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index 8d1d9b7ecf..75a7782354 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -1,5 +1,4 @@ import logging -import operator from typing import Any, Dict, Optional, Sequence, Tuple, Union import tensorrt as trt @@ -406,24 +405,6 @@ def aten_ops_to_copy_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, @@ -432,7 +413,7 @@ def aten_ops_clone( kwargs: Dict[str, Argument], name: str, ) -> Union[TRTTensor, Sequence[TRTTensor]]: - return impl.evaluators.clone( + return impl.cast.clone( network, target, SourceIR.ATEN, diff --git a/py/torch_tensorrt/dynamo/conversion/converter_utils.py b/py/torch_tensorrt/dynamo/conversion/converter_utils.py index bbb57b9da7..d29393b18e 100644 --- a/py/torch_tensorrt/dynamo/conversion/converter_utils.py +++ b/py/torch_tensorrt/dynamo/conversion/converter_utils.py @@ -4,7 +4,7 @@ import tensorrt as trt import torch -from torch.fx.node import Target, _get_qualified_name +from torch.fx.node import Target from torch_tensorrt.fx.converters.converter_utils import ( Frameworks, unified_dtype_converter, @@ -12,6 +12,7 @@ from torch_tensorrt.fx.types import TRTDataType, TRTNetwork, TRTTensor from .._SourceIR import SourceIR +from .converter_registry import ConverterRegistry _LOGGER: logging.Logger = logging.getLogger(__name__) @@ -94,15 +95,12 @@ def cast_trt_tensor( 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)}" - ) + target_str = ConverterRegistry.qualified_name_or_str(target) + target_name = f"{source_ir}_ops{'.' + target_str if target_str else ''}" 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}-[{target_name}]-[{name}]" + identity_layer.name = f"Cast ITensor {input_val.name} from {input_val.dtype} to {trt_dtype} - [{target_name}]-[{name}]" return identity_layer.get_output(0) else: return input_val diff --git a/py/torch_tensorrt/dynamo/conversion/impl/__init__.py b/py/torch_tensorrt/dynamo/conversion/impl/__init__.py index 611dc630fa..8f7ab1badc 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/__init__.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/__init__.py @@ -6,7 +6,6 @@ condition, elementwise, embedding, - evaluators, matmul, normalization, permutation, diff --git a/py/torch_tensorrt/dynamo/conversion/impl/cast.py b/py/torch_tensorrt/dynamo/conversion/impl/cast.py index 68899de766..0c55731169 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/cast.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/cast.py @@ -1,3 +1,4 @@ +import logging from typing import Optional from torch.fx.node import Target @@ -5,6 +6,8 @@ from torch_tensorrt.dynamo.conversion.converter_utils import cast_trt_tensor from torch_tensorrt.fx.types import TRTDataType, TRTNetwork, TRTTensor +LOGGER: logging.Logger = logging.getLogger(__name__) + def to_copy( network: TRTNetwork, @@ -21,3 +24,20 @@ def to_copy( casted_tensor = cast_trt_tensor(network, input, dtype, name, target, source_ir) return casted_tensor + + +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 diff --git a/py/torch_tensorrt/dynamo/conversion/impl/evaluators.py b/py/torch_tensorrt/dynamo/conversion/impl/evaluators.py deleted file mode 100644 index cb61fb6158..0000000000 --- a/py/torch_tensorrt/dynamo/conversion/impl/evaluators.py +++ /dev/null @@ -1,40 +0,0 @@ -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 diff --git a/py/torch_tensorrt/dynamo/conversion/op_evaluators.py b/py/torch_tensorrt/dynamo/conversion/op_evaluators.py new file mode 100644 index 0000000000..a546e34305 --- /dev/null +++ b/py/torch_tensorrt/dynamo/conversion/op_evaluators.py @@ -0,0 +1,32 @@ +import logging +import operator +from typing import Dict, Sequence, Tuple, Union + +from torch.fx.node import Argument, Node, Target +from torch_tensorrt.fx.types import TRTNetwork, TRTTensor + +from .converter_registry import ConverterRegistry, dynamo_tensorrt_converter + +_LOGGER: logging.Logger = logging.getLogger(__name__) + + +def getitem_validator(getitem_node: Node) -> bool: + from torch_tensorrt.dynamo.conversion.converter_registry import DYNAMO_CONVERTERS + + # Getitem nodes can only be converted if their parent node also can + return getitem_node.args[0] in DYNAMO_CONVERTERS + + +# TODO: Subsequent evaluators should be registered here with their own validators +@dynamo_tensorrt_converter(operator.getitem, capability_validator=getitem_validator) +def generic_evaluator( + network: TRTNetwork, + target: Target, + args: Tuple[Argument, ...], + kwargs: Dict[str, Argument], + name: str, +) -> Union[TRTTensor, Sequence[TRTTensor]]: + _LOGGER.debug( + f"Evaluating {ConverterRegistry.qualified_name_or_str(target)} on object with name: {name}" + ) + return target(*args) diff --git a/tests/py/dynamo/converters/test_casts.py b/tests/py/dynamo/converters/test_casts.py index 4bb05ef463..3a4fd65610 100644 --- a/tests/py/dynamo/converters/test_casts.py +++ b/tests/py/dynamo/converters/test_casts.py @@ -5,6 +5,36 @@ from torch_tensorrt.dynamo.conversion import UnsupportedOperatorException +class TestCloneConverter(DispatchTestCase): + def test_clone_contiguous(self): + class Clone(nn.Module): + def forward(self, x): + y = torch.clone(x, memory_format=torch.contiguous_format) + return y + 1 + + inputs = [torch.randn((1, 3, 10))] + self.run_test( + Clone(), + inputs, + expected_ops={torch.ops.aten.clone.default}, + disable_passes=True, + ) + + def test_clone_regular(self): + class Clone(nn.Module): + def forward(self, x): + y = torch.clone(x) + return y + 1 + + inputs = [torch.randn((8, 2, 10))] + self.run_test( + Clone(), + inputs, + expected_ops={torch.ops.aten.clone.default}, + disable_passes=True, + ) + + class TestToCopyConverter(DispatchTestCase): def test_to_copy_half(self): class ToCopyHalf(nn.Module): diff --git a/tests/py/dynamo/converters/test_evaluators.py b/tests/py/dynamo/converters/test_evaluators.py index cf42009495..64dd303727 100644 --- a/tests/py/dynamo/converters/test_evaluators.py +++ b/tests/py/dynamo/converters/test_evaluators.py @@ -7,36 +7,6 @@ from torch.testing._internal.common_utils import run_tests -class TestCloneConverter(DispatchTestCase): - def test_clone_contiguous(self): - class Clone(nn.Module): - def forward(self, x): - y = torch.clone(x, memory_format=torch.contiguous_format) - return y + 1 - - inputs = [torch.randn((1, 3, 10))] - self.run_test( - Clone(), - inputs, - expected_ops={torch.ops.aten.clone.default}, - disable_passes=True, - ) - - def test_clone_regular(self): - class Clone(nn.Module): - def forward(self, x): - y = torch.clone(x) - return y + 1 - - inputs = [torch.randn((8, 2, 10))] - self.run_test( - Clone(), - inputs, - expected_ops={torch.ops.aten.clone.default}, - disable_passes=True, - ) - - # TODO: Switch this test back to self.run_test once an implementation exists # for a converter that returns a list, such as aten.split @unittest.skip("Pending aten.split converter. Currently tested by E2E") diff --git a/tests/py/dynamo/models/test_models.py b/tests/py/dynamo/models/test_models.py index c8f730e2e6..50d7fcbbd9 100644 --- a/tests/py/dynamo/models/test_models.py +++ b/tests/py/dynamo/models/test_models.py @@ -27,7 +27,6 @@ def test_resnet18(ir): "ir": ir, "pass_through_build_failures": True, "optimization_level": 1, - "min_block_size": 10, "ir": "torch_compile", } @@ -176,7 +175,6 @@ def test_resnet18_half(ir): "ir": ir, "pass_through_build_failures": True, "optimization_level": 1, - "min_block_size": 10, "ir": "torch_compile", }