[Draft] feat: flag-gated upstream complex decomposition adapter

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,17 +93,27 @@
     _TENSORRT_RTX,
     _TRTLLM_AVAIL,
     _NATIVE_TRT_COLLECTIVES_AVAIL,
+    _COMPLEX_DECOMP_AVAIL,
 )
 
 T = TypeVar("T")
 
 
 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.

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

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)

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,
 ]

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

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)