fix: guard the 3D convolutions the pad-folding pass hides from the conv guard

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

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,

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__()

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()