diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index cc679ec6c2..4b7385438a 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -3252,6 +3252,30 @@ def aten_ops_le( _CONV_ARG_TRANSPOSED = 6 +def turing_rejects_forward_convolution( + node: Node, + spatial_rank: Optional[int], + settings: Optional[CompilationSettings] = None, +) -> bool: + """Whether a forward convolution over ``spatial_rank`` spatial dims must fall back. + + No valid kernel config for 3D ConvFwd on SM 7.5 for TensorRT-RTX, a known gap. + Transposed 3D is a distinct layer and is unaffected, so callers must only ask about + non-transposed convolutions. + """ + unsupported_spatial_rank = 3 + if spatial_rank != unsupported_spatial_rank: + return False + if not trt_rtx_targets_turing(settings): + return False + _LOGGER.debug( + "3D convolution '%s' is not supported on TensorRT-RTX for Turing " + "(SM 7.5). Falling back to PyTorch.", + node.name, + ) + return True + + def convolution_capability_validator( node: Node, settings: Optional[CompilationSettings] = None ) -> bool: @@ -3278,20 +3302,14 @@ def convolution_capability_validator( ) return False - # No valid kernel config for 3D ConvFwd on SM 7.5: the engine builds, but - # createExecutionContext() then returns nullptr. Transposed 3D is a distinct layer - # and is unaffected. aten.convolution input is (N, C, *spatial), so ndim 5 is 3D. - is_forward_conv = not args_bounds_check(node.args, _CONV_ARG_TRANSPOSED) - if trt_rtx_targets_turing(settings) and is_forward_conv: + # aten.convolution input is (N, C, *spatial), so ndim - 2 is the spatial rank. + # Like every validator in this module this relies on meta["val"] and fails open + # when it is absent. + if not args_bounds_check(node.args, _CONV_ARG_TRANSPOSED): input_node = node.args[_CONV_ARG_INPUT] val = input_node.meta.get("val") if hasattr(input_node, "meta") else None - if (ndim := getattr(val, "ndim", None)) == 5: - _LOGGER.debug( - "3D convolution '%s' (ndim %s) is not supported on TensorRT-RTX for " - "Turing (SM 7.5). Falling back to PyTorch.", - node.name, - ndim, - ) + spatial_rank = val.ndim - 2 if val is not None else None + if turing_rejects_forward_convolution(node, spatial_rank, settings): return False return True diff --git a/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py index 3f15c25cad..9cc8532929 100644 --- a/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py @@ -1,17 +1,22 @@ # mypy: disallow-untyped-decorators=False import logging -from typing import Dict, Sequence, Tuple, Union +from typing import Dict, Optional, Sequence, Tuple, Union import tensorrt as trt -from torch.fx.node import Argument, Target +from torch.fx.node import Argument, Node, Target from torch_tensorrt._features import ENABLED_FEATURES +from torch_tensorrt.dynamo._settings import CompilationSettings from torch_tensorrt.dynamo._SourceIR import SourceIR from torch_tensorrt.dynamo.conversion import impl from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext from torch_tensorrt.dynamo.conversion._ConverterRegistry import ( dynamo_tensorrt_converter, ) +from torch_tensorrt.dynamo.conversion.aten_ops_converters import ( + turing_rejects_forward_convolution, +) +from torch_tensorrt.dynamo.conversion.converter_utils import args_bounds_check from torch_tensorrt.dynamo.lowering.passes.fuse_distributed_ops import ( tensorrt_fused_nccl_all_gather_op, tensorrt_fused_nccl_all_reduce_op, @@ -27,7 +32,34 @@ _LOGGER: logging.Logger = logging.getLogger(__name__) -@dynamo_tensorrt_converter(tensorrt_conv_asym_pad_op, supports_dynamic_shapes=True) +def conv_asym_pad_capability_validator( + node: Node, settings: Optional[CompilationSettings] = None +) -> bool: + """Apply the aten.convolution capability guards to the fused pad+conv op. + + ``fuse_pad_into_convolution`` rewrites constant_pad_nd -> aten.convolution into + this op during lowering, which erases the aten.convolution node that + ``convolution_capability_validator`` would otherwise have rejected. Without this, + the pass smuggles a 3D convolution past that guard onto Turing (SM 7.5), where it + surfaces at runtime as a null execution context. The pass runs in the normal + compile path, so this is not confined to tests. + + Spatial rank comes from the argument list rather than node meta, so the guard holds + whether or not the caller traced with dynamo. + """ + # conv_asym_pad(source, weight, bias, stride, pre_padding, post_padding, dilation, + # groups) -- no transposed argument; the pass skips those outright. + _ARG_STRIDE = 3 + stride = args_bounds_check(node.args, _ARG_STRIDE) + spatial_rank = len(stride) if stride is not None else None + return not turing_rejects_forward_convolution(node, spatial_rank, settings) + + +@dynamo_tensorrt_converter( + tensorrt_conv_asym_pad_op, + capability_validator=conv_asym_pad_capability_validator, + supports_dynamic_shapes=True, +) def conv_asym_pad( ctx: ConversionContext, target: Target, diff --git a/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py b/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py index f023aad984..a909cb4ffa 100644 --- a/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py +++ b/tests/py/dynamo/lowering/test_fuse_pad_into_convolution.py @@ -5,7 +5,7 @@ import torch.nn.functional as F from parameterized import parameterized -from ..conversion.harness import DispatchTestCase +from ..conversion.harness import DispatchTestCase, skip_if_trt_rtx_turing def _node_targets(gm: torch.fx.GraphModule) -> list: @@ -230,6 +230,11 @@ class TestFusePadIntoConvolutionConverter(DispatchTestCase): ] ) def test_padded_conv3d(self, _, pad, stride, dilation, groups, bias): + # The fusion turns this into tensorrt::conv_asym_pad, whose validator rejects + # 3D on Turing. DispatchTestCase has no PyTorch-fallback path, so without this + # the rejection surfaces as UnsupportedOperatorException. + skip_if_trt_rtx_turing(self, "3D convolution") + class PaddedConv3d(nn.Module): def __init__(self) -> None: super().__init__() diff --git a/tests/py/dynamo/models/test_turing_capability_guards.py b/tests/py/dynamo/models/test_turing_capability_guards.py index e7f5044f40..023bbe32d8 100644 --- a/tests/py/dynamo/models/test_turing_capability_guards.py +++ b/tests/py/dynamo/models/test_turing_capability_guards.py @@ -70,6 +70,33 @@ def forward(self, x): return self.conv(x) +class PaddedConv3d(nn.Module): + """A conv3d the 3D-conv guard never sees as a convolution. + + fuse_pad_into_convolution folds the pad into the conv and rewrites the pair as + tensorrt::conv_asym_pad, erasing the aten.convolution node. It fires for any + zero-fill, non-negative pad -- asymmetry is not required. + """ + + def __init__(self): + super().__init__() + self.conv = nn.Conv3d(4, 8, 3, padding=0) + + def forward(self, x): + return self.conv(F.pad(x, (1, 1, 1, 1, 2, 0))) + + +class PaddedConv2d(nn.Module): + """Same fusion, 2D -- which Turing does support.""" + + def __init__(self): + super().__init__() + self.conv = nn.Conv2d(4, 8, 3, padding=0) + + def forward(self, x): + return self.conv(F.pad(x, (0, 1, 0, 1))) + + @unittest.skipIf( not ENABLED_FEATURES.tensorrt_rtx, "Turing capability guards only apply to TensorRT-RTX", @@ -154,6 +181,23 @@ def test_declared_turing_target_keeps_conv2d_on_trt(self): compiled = self._compile(mod, inputs, target_compute_capabilities=[TURING]) self.assertGreater(_trt_submodule_count(compiled), 0) + def test_declared_turing_target_falls_back_padded_conv3d(self): + # The fused op needs its own guard: by partitioning time the aten.convolution + # node the conv guard keys on no longer exists. + mod = PaddedConv3d() + inputs = (torch.randn(1, 4, 8, 8, 8).cuda(),) + compiled = self._compile(mod, inputs, target_compute_capabilities=[TURING]) + self.assertEqual(_trt_submodule_count(compiled), 0) + self._assert_matches_eager(mod, compiled, inputs) + + def test_declared_turing_target_keeps_padded_conv2d_on_trt(self): + # Only rank 3 is rejected; the fused op must keep serving 2D. + mod = PaddedConv2d() + inputs = (torch.randn(1, 4, 16, 16).cuda(),) + compiled = self._compile(mod, inputs, target_compute_capabilities=[TURING]) + self.assertGreater(_trt_submodule_count(compiled), 0) + self._assert_matches_eager(mod, compiled, inputs) + def test_declared_turing_target_keeps_transposed_conv3d_on_trt(self): # Transposed 3D convolution is a distinct layer and does work on Turing. mod = ConvTranspose3d()