diff --git a/py/torch_tensorrt/_features.py b/py/torch_tensorrt/_features.py index 85a1595009..1c5e013fce 100644 --- a/py/torch_tensorrt/_features.py +++ b/py/torch_tensorrt/_features.py @@ -27,6 +27,7 @@ "tensorrt_rtx", "trtllm_for_nccl", "native_trt_collectives", + "complex_decomposition", ], ) @@ -49,6 +50,11 @@ _TS_FE_AVAIL = os.path.isfile(linked_file_full_path) _TORCHTRT_RT_AVAIL = _TS_FE_AVAIL or os.path.isfile(linked_file_runtime_full_path) _DYNAMO_FE_AVAIL = version.parse(sanitized_torch_version()) >= version.parse("2.1.dev") +# PyTorch's upstream complex-tensor decomposition (pytorch/pytorch#169832) ships +# in torch>=2.14.dev; gates the complex_decomposition_adapter lowering pass. +_COMPLEX_DECOMP_AVAIL = version.parse(sanitized_torch_version()) >= version.parse( + "2.14.dev" +) _FX_FE_AVAIL = False if _TENSORRT_RTX else True _REFIT_AVAIL = True _WINDOWS_CROSS_COMPILE = check_cross_compile_trt_win_lib() @@ -87,6 +93,7 @@ _TENSORRT_RTX, _TRTLLM_AVAIL, _NATIVE_TRT_COLLECTIVES_AVAIL, + _COMPLEX_DECOMP_AVAIL, ) T = TypeVar("T") @@ -94,10 +101,19 @@ def _enabled_features_str() -> str: enabled = lambda x: "ENABLED" if x else "DISABLED" - out_str: str = f"Enabled Features:\n - Dynamo Frontend: {enabled(_DYNAMO_FE_AVAIL)}\n - Torch-TensorRT Runtime: {enabled(_TORCHTRT_RT_AVAIL)}\n - FX Frontend: {enabled(_FX_FE_AVAIL)}\n - TorchScript Frontend: {enabled(_TS_FE_AVAIL)}\n - Refit: {enabled(_REFIT_AVAIL)}\n - QDP Plugin: {enabled(_QDP_PLUGIN_AVAIL)} \n - TensorRT-RTX: {enabled(_TENSORRT_RTX)}\n - TensorRT-LLM for NCCL: {enabled(_TRTLLM_AVAIL)}\n - Native TRT Collectives: {enabled(_NATIVE_TRT_COLLECTIVES_AVAIL)}\n" # type: ignore[no-untyped-call] + out_str: str = f"Enabled Features:\n - Dynamo Frontend: {enabled(_DYNAMO_FE_AVAIL)}\n - Torch-TensorRT Runtime: {enabled(_TORCHTRT_RT_AVAIL)}\n - FX Frontend: {enabled(_FX_FE_AVAIL)}\n - TorchScript Frontend: {enabled(_TS_FE_AVAIL)}\n - Refit: {enabled(_REFIT_AVAIL)}\n - QDP Plugin: {enabled(_QDP_PLUGIN_AVAIL)} \n - TensorRT-RTX: {enabled(_TENSORRT_RTX)}\n - TensorRT-LLM for NCCL: {enabled(_TRTLLM_AVAIL)}\n - Native TRT Collectives: {enabled(_NATIVE_TRT_COLLECTIVES_AVAIL)}\n - Complex Decomposition: {enabled(_COMPLEX_DECOMP_AVAIL)}\n" # type: ignore[no-untyped-call] return out_str +def has_complex_decomposition() -> bool: + """Check if PyTorch's upstream complex-tensor decomposition is available. + + Returns: + bool: True if torch>=2.14.dev (ships decompose_complex_in_graph, PR #169832) + """ + return bool(ENABLED_FEATURES.complex_decomposition) + + # Inline helper functions for checking feature availability def has_torch_tensorrt_runtime() -> bool: """Check if Torch-TensorRT C++ runtime is available. diff --git a/py/torch_tensorrt/dynamo/_defaults.py b/py/torch_tensorrt/dynamo/_defaults.py index 9c8a1f9f90..d4d5d1d3f3 100644 --- a/py/torch_tensorrt/dynamo/_defaults.py +++ b/py/torch_tensorrt/dynamo/_defaults.py @@ -22,6 +22,11 @@ OPTIMIZATION_LEVEL = None SPARSE_WEIGHTS = False TRUNCATE_DOUBLE = False +# When True, use PyTorch's upstream complex decomposition (pytorch/pytorch#169832) +# via complex_decomposition_adapter instead of the legacy hand-rolled +# complex_graph_detection pass. Default False until the new path is validated +# against the complex/RoPE test suites (issue #4390). +USE_COMPLEX_DECOMPOSITION = False USE_FAST_PARTITIONER = True ENABLE_EXPERIMENTAL_DECOMPOSITIONS = False REQUIRE_FULL_COMPILATION = False diff --git a/py/torch_tensorrt/dynamo/_settings.py b/py/torch_tensorrt/dynamo/_settings.py index 7ff5f6bdff..ca130f31a4 100644 --- a/py/torch_tensorrt/dynamo/_settings.py +++ b/py/torch_tensorrt/dynamo/_settings.py @@ -49,6 +49,7 @@ TILING_OPTIMIZATION_LEVEL, TIMING_CACHE_PATH, TRUNCATE_DOUBLE, + USE_COMPLEX_DECOMPOSITION, USE_DISTRIBUTED_MODE_TRACE, USE_FAST_PARTITIONER, USE_FP32_ACC, @@ -131,6 +132,7 @@ class CompilationSettings: version_compatible: bool = VERSION_COMPATIBLE optimization_level: Optional[int] = OPTIMIZATION_LEVEL truncate_double: bool = TRUNCATE_DOUBLE + use_complex_decomposition: bool = USE_COMPLEX_DECOMPOSITION use_fast_partitioner: bool = USE_FAST_PARTITIONER enable_experimental_decompositions: bool = ENABLE_EXPERIMENTAL_DECOMPOSITIONS device: Device = field(default_factory=default_device) diff --git a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py index 271f7c98b7..0bb003dc82 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py @@ -11,6 +11,7 @@ ) from .annotate_fp8_sdpa import annotate_fp8_sdpa +from .complex_decomposition_adapter import complex_decomposition_adapter from .complex_graph_rewrite import complex_graph_detection from .constant_folding import constant_fold from .force_causal_efficient_attention import force_causal_efficient_attention @@ -31,6 +32,19 @@ rule_based_autocast, ] + +def complex_lowering_pass( + gm: torch.fx.GraphModule, settings: CompilationSettings +) -> torch.fx.GraphModule: + """Select the complex-numerics lowering by settings flag (issue #4390). + + ``use_complex_decomposition`` routes to PyTorch's upstream decomposition; + otherwise the legacy hand-rolled rewriter runs. + """ + if getattr(settings, "use_complex_decomposition", False): + return complex_decomposition_adapter(gm, settings) + return complex_graph_detection(gm, settings) + post_lowering_pass_list = [ replace_fused_rms_norm, remove_input_alias_fixing_clones, @@ -40,7 +54,7 @@ replace_max_pool_with_indices, remove_assert_nodes, remove_num_users_is_0_nodes, - complex_graph_detection, + complex_lowering_pass, force_causal_efficient_attention, annotate_fp8_sdpa, ] diff --git a/py/torch_tensorrt/dynamo/lowering/passes/complex_decomposition_adapter.py b/py/torch_tensorrt/dynamo/lowering/passes/complex_decomposition_adapter.py new file mode 100644 index 0000000000..7599222c0a --- /dev/null +++ b/py/torch_tensorrt/dynamo/lowering/passes/complex_decomposition_adapter.py @@ -0,0 +1,169 @@ +"""Adopt PyTorch's upstream complex decomposition (pytorch/pytorch#169832). + +This is the "augment, not replace" path proposed in issue #4390. Instead of the +hand-maintained per-op rewriter in ``complex_graph_rewrite.py``, we delegate the +interior complex -> real expansion to PyTorch's ``decompose_complex_in_graph`` +(a tensor-subclass / SoA retrace) and keep only the TRT-specific glue: + + 1. capture the complex I/O signature before rewriting (drives the boundary + adapters in ``_compiler._insert_complex_io_adapters``); + 2. call upstream to expand every complex op into real ops on separate re/im; + 3. normalize the SoA seams (``aten.complex`` / ``aten.real`` / ``aten.imag``) + into the interleaved ``[..., 2]`` layout the rest of the TRT flow expects, + so the engine only ever sees real tensors (Option A in the RFC). + +Gated by ``settings.use_complex_decomposition``; falls back to the legacy pass on +older torch (the upstream API only exists in torch>=2.14.0.dev). + +NOTE: the exact shape of what ``decompose_complex_in_graph`` emits at the +boundary must be confirmed by a spike (torch is required to run it). The steps +marked ``TODO(spike)`` encode assumptions to validate against real output. +""" + +import logging + +import torch +from torch.fx import GraphModule +from torch_tensorrt._features import has_complex_decomposition +from torch_tensorrt.dynamo._settings import CompilationSettings +from torch_tensorrt.dynamo.lowering.passes.pass_utils import ( + clean_up_graph_after_modifications, +) + +# Reuse the I/O-signature capture helpers already written for the legacy pass. +from .complex_graph_rewrite import ( + _get_complex_input_dtypes, + _get_complex_input_names, + _get_complex_output_indices, + complex_graph_detection, +) + +logger = logging.getLogger(__name__) + + +def complex_decomposition_adapter( + gm: GraphModule, settings: CompilationSettings +) -> GraphModule: + """Delegate complex decomposition to upstream, keep TRT boundary glue ours.""" + if not has_complex_decomposition(): + logger.warning( + "complex decomposition requires torch>=2.14.dev; falling back to the " + "legacy complex_graph_detection pass." + ) + return complex_graph_detection(gm, settings) + + # (1) Capture the complex I/O signature BEFORE any rewrite mutates the graph. + # Identical contract to the legacy pass so _insert_complex_io_adapters + # keeps working unchanged. + gm.meta["complex_output_indices"] = _get_complex_output_indices(gm) + gm.meta["complex_input_names"] = _get_complex_input_names(gm) + gm.meta["complex_input_dtypes"] = _get_complex_input_dtypes(gm) + if not gm.meta["complex_input_names"] and not gm.meta["complex_output_indices"]: + # No complex I/O and no interior complex work -> nothing to do. (Interior- + # only complex still shows up via inputs/outputs of the complex region, so + # this early-out matches the legacy pass's behavior.) + if not _graph_has_complex(gm): + return gm + + # (2) Hand the interior expansion to upstream. + from torch._functorch._aot_autograd.complex_decomposition import ( + decompose_complex_in_graph, + ) + + flat_args = _fake_flat_args(gm) + logger.debug("complex_decomposition_adapter: retracing under ComplexTensor") + gm = decompose_complex_in_graph(gm, flat_args) + + # (3) Normalize SoA seams into the interleaved [..., 2] layout used by TRT. + gm = _normalize_complex_boundary_for_trt(gm) + gm = clean_up_graph_after_modifications(gm) + return gm + + +def _graph_has_complex(gm: GraphModule) -> bool: + from torch_tensorrt.dynamo.utils import COMPLEX_DTYPES + + for node in gm.graph.nodes: + val = node.meta.get("val", None) + if val is not None and getattr(val, "dtype", None) in COMPLEX_DTYPES: + return True + return False + + +def _fake_flat_args(gm: GraphModule) -> list: + """Build the example inputs the upstream retrace (make_fx) needs. + + We reuse the placeholders' existing fake tensors so the export ShapeEnv / + SymInt ranges are preserved through the retrace. Whether dynamic shapes + actually survive make_fx is RFC open-question #1 -- assert-worthy in the spike. + """ + args = [] + for node in gm.graph.nodes: + if node.op != "placeholder": + continue + val = node.meta.get("val", None) + if val is None: + tm = node.meta["tensor_meta"] + val = torch.empty(tm.shape, dtype=tm.dtype, device=tm.device) + args.append(val) + return args + + +def _normalize_complex_boundary_for_trt(gm: GraphModule) -> GraphModule: + """Fold upstream's SoA seams into the interleaved ``[..., 2]`` real layout. + + Upstream leaves the graph in terms of separate re/im joined by + ``aten.complex(re, im)`` and unpacked by ``aten.real`` / ``aten.imag``. TRT + has no converter for any of those. We: + + * ``real(z)`` / ``imag(z)`` where ``z = complex(re, im)`` -> re / im + (cancel the pack/unpack round-trip) + * remaining ``aten.complex(re, im)`` -> ``stack([re, im], -1)`` tagged as + complex-layout, so the [..., 2] tensor flows to the boundary adapters. + + TODO(spike): confirm against real decompose_complex_in_graph output -- + the op set at the boundary (aten.complex vs view_as_complex, real vs + select) and whether inputs are unpacked with real/imag or view_as_real. + """ + g = gm.graph + aten = torch.ops.aten + + # Pass A: cancel real(complex(re,im)) / imag(complex(re,im)) round-trips. + for node in list(g.nodes): + if node.op != "call_function" or node.target not in ( + aten.real.default, + aten.imag.default, + ): + continue + # node is real(z) / imag(z); node.args[0] is that single arg z -- the + # tensor being unpacked (the "source" producer feeding this unpack). + src = node.args[0] + # Only fold when z was itself produced by aten.complex(re, im): then + # real(complex(re,im)) == re and imag(complex(re,im)) == im, so the + # pack/unpack round-trip is pure noise. (isinstance guards against a + # non-Node arg, e.g. a constant, which has no .target.) + if isinstance(src, torch.fx.Node) and src.target == aten.complex.default: + # aten.complex stores re at args[0], im at args[1] -> pick the half + # this unpack wanted: real -> 0, imag -> 1. + idx = 0 if node.target == aten.real.default else 1 + # Rewire every consumer of real(z)/imag(z) straight to re/im, then + # delete the now-orphaned real()/imag() node. + node.replace_all_uses_with(src.args[idx]) + g.erase_node(node) + + # Pass B: turn surviving aten.complex(re, im) into stack([re, im], -1). + for node in list(g.nodes): + if node.op != "call_function" or node.target != aten.complex.default: + continue + re, im = node.args[0], node.args[1] + with g.inserting_before(node): + re_u = g.call_function(aten.unsqueeze.default, (re, -1)) + im_u = g.call_function(aten.unsqueeze.default, (im, -1)) + packed = g.call_function(aten.cat.default, ([re_u, im_u], -1)) + packed.meta["is_complex_layout"] = True + node.replace_all_uses_with(packed) + g.erase_node(node) + + g.lint() + gm.recompile() + return gm diff --git a/tests/py/dynamo/hlo/test_complex_ops.py b/tests/py/dynamo/hlo/test_complex_ops.py index 27d135c4f1..2fdcd12d7e 100644 --- a/tests/py/dynamo/hlo/test_complex_ops.py +++ b/tests/py/dynamo/hlo/test_complex_ops.py @@ -22,6 +22,7 @@ import torch.nn as nn import torch_tensorrt as torchtrt from torch.export import Dim +from torch_tensorrt._features import has_complex_decomposition from torch_tensorrt.dynamo.utils import COSINE_THRESHOLD, cosine_similarity # --------------------------------------------------------------------------- @@ -62,6 +63,34 @@ def _cossim_real(py_out: torch.Tensor, trt_out: torch.Tensor, tag: str) -> None: _COMPILE = dict(ir="dynamo", min_block_size=1, pass_through_build_failures=True) +# Extra compile kwargs for the "explicit compile" tests (dynamic-shape and +# truncate cases) that do NOT spread **_COMPILE. The autouse fixture below +# fills this with {"use_complex_decomposition": True} on the decomp param. +_DECOMP: dict = {} + + +@pytest.fixture(autouse=True, params=[False, True], ids=["legacy", "decomp"]) +def _complex_path(request): + """Run every complex test twice: once on the legacy hand-rolled rewriter and + once on PyTorch's upstream complex decomposition (issue #4390). + + The flag is toggled on the shared _COMPILE dict (picked up by the ~28 tests + that spread **_COMPILE) and mirrored into _DECOMP (spread by the handful of + explicit torchtrt.dynamo.compile / truncate tests). + """ + use_decomp = request.param + if use_decomp and not has_complex_decomposition(): + pytest.skip("decompose_complex_in_graph requires torch>=2.14.dev") + + _COMPILE.pop("use_complex_decomposition", None) + _DECOMP.clear() + if use_decomp: + _COMPILE["use_complex_decomposition"] = True + _DECOMP["use_complex_decomposition"] = True + yield + _COMPILE.pop("use_complex_decomposition", None) + _DECOMP.clear() + # =========================================================================== # 1. I/O boundary tests @@ -789,7 +818,11 @@ def test_complex_mul_dynamic_seqlen(): dynamic_shapes = ({1: seq}, {0: seq}) ep = torch.export.export(model, inputs, dynamic_shapes=dynamic_shapes) trt_model = torchtrt.dynamo.compile( - ep, inputs=inputs, min_block_size=1, pass_through_build_failures=True + ep, + inputs=inputs, + min_block_size=1, + pass_through_build_failures=True, + **_DECOMP, ) py_out = model(*inputs) trt_out = trt_model(*inputs) @@ -816,7 +849,11 @@ def test_complex_output_dynamic_batch(): dynamic_shapes = ({0: batch}, {}) ep = torch.export.export(model, inputs, dynamic_shapes=dynamic_shapes) trt_model = torchtrt.dynamo.compile( - ep, inputs=inputs, min_block_size=1, pass_through_build_failures=True + ep, + inputs=inputs, + min_block_size=1, + pass_through_build_failures=True, + **_DECOMP, ) py_out = model(*inputs) trt_out = trt_model(*inputs) @@ -856,6 +893,7 @@ def test_complex128_truncated_to_float32(): min_block_size=1, pass_through_build_failures=True, truncate_double=True, + **_DECOMP, ) py_out = model(*inputs).float() # cast reference to float32 for comparison trt_out = trt_model(*inputs) diff --git a/tests/py/dynamo/lowering/test_complex_decomposition_adapter.py b/tests/py/dynamo/lowering/test_complex_decomposition_adapter.py new file mode 100644 index 0000000000..e0b0c082dc --- /dev/null +++ b/tests/py/dynamo/lowering/test_complex_decomposition_adapter.py @@ -0,0 +1,119 @@ +"""Unit tests for the complex_decomposition_adapter lowering pass (issue #4390). + +These tests exercise the TRT-specific glue around PyTorch's upstream complex +decomposition -- they do NOT require a GPU or a TRT build: + + * _normalize_complex_boundary_for_trt: folds the aten.complex / real / imag + seams that decompose_complex_in_graph leaves behind into the interleaved + [..., 2] real layout the rest of the TRT flow expects. + * the torch-version feature gate: when the upstream API is unavailable the + adapter must fall back to the legacy complex_graph_detection pass. +""" + +import pytest +import torch +from torch_tensorrt.dynamo.lowering.passes import complex_decomposition_adapter as cda + +aten = torch.ops.aten + + +def _targets(gm): + return [n.target for n in gm.graph.nodes if n.op == "call_function"] + + +# --------------------------------------------------------------------------- +# _normalize_complex_boundary_for_trt +# --------------------------------------------------------------------------- + + +def test_normalize_folds_real_imag_of_complex(): + """real(complex(re,im)) -> re ; imag(complex(re,im)) -> im, and the + aten.complex / real / imag nodes are erased.""" + + class M(torch.nn.Module): + def forward(self, re, im): + z = torch.ops.aten.complex.default(re, im) + return torch.ops.aten.real.default(z), torch.ops.aten.imag.default(z) + + re = torch.randn(4) + im = torch.randn(4) + gm = torch.fx.symbolic_trace(M()) + + gm = cda._normalize_complex_boundary_for_trt(gm) + targets = _targets(gm) + + assert aten.real.default not in targets + assert aten.imag.default not in targets + assert aten.complex.default not in targets # unused after folding -> DCE'd + # outputs should be re, im directly + out_re, out_im = gm(re, im) + assert torch.equal(out_re, re) + assert torch.equal(out_im, im) + + +def test_normalize_packs_surviving_complex_to_interleaved(): + """A complex(re,im) whose result is actually used becomes + cat([unsqueeze(re,-1), unsqueeze(im,-1)], -1) tagged is_complex_layout.""" + + class M(torch.nn.Module): + def forward(self, re, im): + return torch.ops.aten.complex.default(re, im) + + re = torch.randn(3) + im = torch.randn(3) + gm = torch.fx.symbolic_trace(M()) + + gm = cda._normalize_complex_boundary_for_trt(gm) + targets = _targets(gm) + + assert aten.complex.default not in targets + assert aten.cat.default in targets + assert aten.unsqueeze.default in targets + + # the packed cat node must carry the complex-layout tag + cat_nodes = [ + n for n in gm.graph.nodes if n.target == aten.cat.default + ] + assert cat_nodes and cat_nodes[0].meta.get("is_complex_layout") is True + + # numerically: output is the [...,2] interleaved layout of re/im + out = gm(re, im) + assert out.shape == (3, 2) + assert torch.equal(out[..., 0], re) + assert torch.equal(out[..., 1], im) + + +# --------------------------------------------------------------------------- +# feature gate / fallback +# --------------------------------------------------------------------------- + + +def test_gate_falls_back_to_legacy_on_old_torch(monkeypatch): + """When has_complex_decomposition() is False, the adapter must delegate to + the legacy complex_graph_detection pass (and NOT import the upstream API).""" + calls = {} + + def fake_legacy(gm, settings): + calls["legacy"] = True + return gm + + monkeypatch.setattr(cda, "has_complex_decomposition", lambda: False) + monkeypatch.setattr(cda, "complex_graph_detection", fake_legacy) + + gm = torch.fx.symbolic_trace(torch.nn.Identity()) + out = cda.complex_decomposition_adapter(gm, settings=object()) + + assert calls.get("legacy") is True + assert out is gm + + +@pytest.mark.skipif( + not cda.has_complex_decomposition(), + reason="decompose_complex_in_graph requires torch>=2.14.dev", +) +def test_upstream_api_importable(): + """Smoke check: on a supported torch, the upstream entry point is importable + at the path the adapter uses.""" + from torch._functorch._aot_autograd.complex_decomposition import ( # noqa: F401 + decompose_complex_in_graph, + )