From aa198024ef218f346aca193bb3e7cf62367103ee Mon Sep 17 00:00:00 2001 From: apbose Date: Wed, 29 Oct 2025 18:21:24 -0700 Subject: [PATCH 1/3] commit to add validator in nonzero for thor runs (currently fallback to torch) and make ATenAnyDimNegIndexConvertsCorrectly input deterministic on ORIN --- .../dynamo/conversion/aten_ops_converters.py | 9 +++++++++ tests/core/conversion/converters/test_reduce.cpp | 9 ++++++++- 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index 164f0c1065..df7d144254 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -24,6 +24,7 @@ get_positive_dim, is_only_operator_on_placeholder, ) +from torch_tensorrt.dynamo.utils import is_thor _LOGGER: logging.Logger = logging.getLogger(__name__) @@ -3601,10 +3602,18 @@ def aten_ops_full( ) +def nonzero_validator( + node: Node, settings: Optional[CompilationSettings] = None +) -> bool: + return not is_thor() + + # currently nonzero is not supported for tensorrt_rtx # TODO: lan to add the nonzero support once tensorrt_rtx team has added the support +# TODO: apbose to remove the capability validator once thor bug resolve in NGC @dynamo_tensorrt_converter( torch.ops.aten.nonzero.default, + capability_validator=nonzero_validator, supports_dynamic_shapes=True, requires_output_allocator=True, ) diff --git a/tests/core/conversion/converters/test_reduce.cpp b/tests/core/conversion/converters/test_reduce.cpp index 5e75723fbe..be9cf171ce 100644 --- a/tests/core/conversion/converters/test_reduce.cpp +++ b/tests/core/conversion/converters/test_reduce.cpp @@ -345,7 +345,14 @@ TEST(Converters, ATenAnyDimNegIndexConvertsCorrectly) { %3 : bool = prim::Constant[value=1]() %5 : Tensor = aten::any(%0, %1, %3) return (%5))IR"; - auto in = at::randint(-2, 2, {2, 32}, at::kCUDA); + std::vector data(64, 0); + for (int i = 0; i < 64; ++i) { + if (i % 7 == 0) + data[i] = 1; // some positives + if (i % 13 == 0) + data[i] = -1; // some negatives + } + auto in = at::tensor(data, at::TensorOptions().dtype(at::kInt).device(at::kCUDA)).reshape({2, 32}); // shape [2, 32] test_body(graph, in); } From cb75a49c9e5c9bbfe7e2f21240d3e24d65974fda Mon Sep 17 00:00:00 2001 From: apbose Date: Thu, 30 Oct 2025 11:03:12 -0700 Subject: [PATCH 2/3] validator for index non zero case- else error- Could not find any implementation for node [NON_ZERO]-[aten_ops.index.Tensor]-[index_tensor_bool_nonzero_1] --- .../dynamo/conversion/aten_ops_converters.py | 22 ++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index df7d144254..e5c937f6f5 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -9,7 +9,7 @@ from tensorrt import ITensor as TRTTensor from torch.fx.node import Argument, Node, Target from torch_tensorrt._features import needs_not_tensorrt_rtx -from torch_tensorrt._utils import is_tensorrt_version_supported +from torch_tensorrt._utils import is_tensorrt_version_supported, is_thor from torch_tensorrt.dynamo._settings import CompilationSettings from torch_tensorrt.dynamo._SourceIR import SourceIR from torch_tensorrt.dynamo.conversion import impl @@ -24,7 +24,8 @@ get_positive_dim, is_only_operator_on_placeholder, ) -from torch_tensorrt.dynamo.utils import is_thor + +from torch_tensorrt._utils import is_thor _LOGGER: logging.Logger = logging.getLogger(__name__) @@ -425,9 +426,24 @@ def index_dtype_validator( return True +def index_nonbool_validator( + node: Node, settings: Optional[CompilationSettings] = None +) -> bool: + # for thor, we don't support boolean indices + if is_thor(): + index = node.args[1] + for ind in index: + if ind is not None: + val = ind.meta.get("val") + if val is not None and val.dtype == torch.bool: + return False + return True + + @dynamo_tensorrt_converter( torch.ops.aten.index.Tensor, - capability_validator=index_dtype_validator, + capability_validator=lambda node, settings: index_dtype_validator(node, settings) + and index_nonbool_validator(node, settings), supports_dynamic_shapes=True, requires_output_allocator=True, ) From 8162718aa2909906446d06220fa761f6daa85fe9 Mon Sep 17 00:00:00 2001 From: apbose Date: Thu, 6 Nov 2025 14:02:47 -0800 Subject: [PATCH 3/3] changing is_thor() location based on function location in main branch --- py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py | 1 - 1 file changed, 1 deletion(-) diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index e5c937f6f5..c25f62d500 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -24,7 +24,6 @@ get_positive_dim, is_only_operator_on_placeholder, ) - from torch_tensorrt._utils import is_thor _LOGGER: logging.Logger = logging.getLogger(__name__)