diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index 4b7385438a..eb43a8b55b 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -3369,7 +3369,55 @@ def aten_ops_convolution( ) -@dynamo_tensorrt_converter(torch.ops.aten._cdist_forward.default) +# Above this many rows in either operand, the p == 2 path of +# impl.normalization.cdist_forward switches from a broadcast-subtract to a matrix +# multiply. Kept deliberately in sync with the threshold in that converter. +# aten._cdist_forward(x1, x2, p, compute_mode) +_CDIST_ARG_P = 2 +_CDIST_ARG_COMPUTE_MODE = 3 + + +def cdist_forward_capability_validator( + node: Node, settings: Optional[CompilationSettings] = None +) -> bool: + """Reject the cdist variants whose converter emits a GEMM, on Turing (SM 7.5). + + The GEMM is emitted *inside* the converter, so the graph holds one + ``_cdist_forward`` node and no ``mm``/``bmm`` for ``gemm_capability_validator`` to + reject, and TensorRT-RTX then fails. Which arguments emit one is decided by + ``cdist_emits_matmul``, the converter's own predicate, so the two cannot drift. + """ + if not trt_rtx_targets_turing(settings): + return True + + def operand_rows(operand: Argument) -> Optional[int]: + val = operand.meta.get("val") if hasattr(operand, "meta") else None + shape = getattr(val, "shape", None) + if shape is None or len(shape) < 2: + return None + rows = shape[-2] + return rows if isinstance(rows, int) else None + + operands = node.args[:2] # x1, x2 + if not impl.normalization.ops.cdist_emits_matmul( + args_bounds_check(node.args, _CDIST_ARG_P, replacement=2.0), + args_bounds_check(node.args, _CDIST_ARG_COMPUTE_MODE, replacement=None), + [operand_rows(operand) for operand in operands], + ): + return True + + _LOGGER.debug( + "cdist '%s' computes p=2 as a matrix multiply, which is not supported on " + "TensorRT-RTX for Turing (SM 7.5). Falling back to PyTorch.", + node.name, + ) + return False + + +@dynamo_tensorrt_converter( + torch.ops.aten._cdist_forward.default, + capability_validator=cdist_forward_capability_validator, +) def aten_ops_cdist_forward( ctx: ConversionContext, target: Target, diff --git a/tests/py/dynamo/models/test_turing_capability_guards.py b/tests/py/dynamo/models/test_turing_capability_guards.py index 023bbe32d8..e982a22722 100644 --- a/tests/py/dynamo/models/test_turing_capability_guards.py +++ b/tests/py/dynamo/models/test_turing_capability_guards.py @@ -97,6 +97,30 @@ def forward(self, x): return self.conv(F.pad(x, (0, 1, 0, 1))) +class Cdist(nn.Module): + """A GEMM the explicit GEMM guards never see, emitted inside the cdist converter. + + The op has to be called directly: torch.cdist does not reach this converter on the + GEMM path, because ATen's cdist calls _euclidean_dist there, which lowers to + aten.matmul and is guarded already. Only _cdist_forward's own p == 2 branch emits a + matmul layer that no mm/bmm node stands for. + """ + + def __init__(self, p=2.0, compute_mode=None): + super().__init__() + self.p, self.compute_mode = p, compute_mode + + def forward(self, x1, x2): + return torch.ops.aten._cdist_forward.default(x1, x2, self.p, self.compute_mode) + + +def _cdist_inputs(rows1, rows2, dtype=torch.float32): + return ( + torch.randn(4, rows1, 5, dtype=dtype).cuda(), + torch.randn(4, rows2, 5, dtype=dtype).cuda(), + ) + + @unittest.skipIf( not ENABLED_FEATURES.tensorrt_rtx, "Turing capability guards only apply to TensorRT-RTX", @@ -205,6 +229,59 @@ def test_declared_turing_target_keeps_transposed_conv3d_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_cdist_p2_compute_mode_1(self): + # compute_mode=1 is "always use the matrix multiply", whatever the row count. + mod = Cdist(compute_mode=1) + inputs = _cdist_inputs(6, 8) + 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_falls_back_cdist_p2_default_compute_mode(self): + # An absent compute_mode means "use the matrix multiply if either operand has + # more than 25 rows". The validator has to normalise None to 0 as the converter does. + mod = Cdist() + inputs = _cdist_inputs(35, 45) + 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_falls_back_fp16_cdist_p2(self): + # Unlike every other guard here, this one is not FP32-only: on a T4 the fused + # Matmul_MUL_SUB_SQRT_ pattern fails for FP16 operands too, under every + # enabled_precisions setting, even though a bare FP16 GEMM runs there fine. + mod = Cdist(compute_mode=1) + inputs = _cdist_inputs(6, 8, dtype=torch.half) + 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_cdist_p2_compute_mode_2_on_trt(self): + # compute_mode=2 is "never use the matrix multiply", so there is no GEMM to reject. + mod = Cdist(compute_mode=2) + inputs = _cdist_inputs(35, 45) + 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_cdist_p2_small_rows_on_trt(self): + # At or below 25 rows the default compute_mode emits no GEMM. Rejecting these + # anyway would be worse than useless: PyTorch's cdist_cuda kernel has no Half + # implementation, so for FP16 the fallback raises where TensorRT-RTX succeeds. + mod = Cdist(compute_mode=0) + inputs = _cdist_inputs(6, 8) + 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_cdist_p1_on_trt(self): + # Only p == 2 has a matrix-multiply path at all. + mod = Cdist(p=1.0, compute_mode=0) + inputs = _cdist_inputs(35, 45) + compiled = self._compile(mod, inputs, target_compute_capabilities=[TURING]) + self.assertGreater(_trt_submodule_count(compiled), 0) + self._assert_matches_eager(mod, compiled, inputs) + @unittest.skipIf( _is_turing(), "on Turing the native path is already guarded; see the SM75 tests" ) @@ -387,6 +464,19 @@ def test_conv3d_falls_back(self): compiled = self._compile(mod, inputs) self.assertEqual(_trt_submodule_count(compiled), 0) + def test_cdist_p2_falls_back(self): + # Unguarded this built an engine whose createExecutionContext() then returned + # nullptr: "cuDNN graph compilation failed: No valid engine configs for + # Matmul_MUL_SUB_SQRT_". + mod = Cdist(compute_mode=1) + inputs = _cdist_inputs(6, 8) + compiled = self._compile(mod, inputs) + self.assertEqual(_trt_submodule_count(compiled), 0) + with torch.no_grad(): + ref = mod.eval().cuda()(*inputs) + out = compiled(*inputs) + self.assertEqual((ref - out).abs().max().item(), 0.0) + def test_bfloat16_falls_back_without_crashing(self): # Unguarded, compiling bfloat16 for SM 7.5 segfaulted the process. class Add(nn.Module):