fix(dynamo): correct legacy exporter (retrace=False) submodule inlining for hybrid graphs

py/torch_tensorrt/dynamo/_exporter.py

@@ -4,6 +4,7 @@
 from typing import Any, Dict, Optional, Sequence, Tuple
 
 import torch
+import torch.utils._pytree as pytree
 from torch._export.non_strict_utils import make_constraints
 from torch._guards import detect_fake_mode
 from torch._library.fake_class_registry import FakeScriptObject
@@ -19,6 +20,7 @@
     OutputSpec,
     TensorArgument,
 )
+from torch.fx.graph import _PyTreeCodeGen
 from torch_tensorrt._features import ENABLED_FEATURES
 from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ENGINE_IDX, NAME_IDX
 
@@ -230,6 +232,12 @@ def lift(
                         kind=input_kind,
                         arg=input_spec_arg,
                         target=node.target,
+                        # torch>=2.3 requires an explicit persistent flag on BUFFER
+                        # specs. state_dict() excludes non-persistent buffers by
+                        # construction, so any buffer reaching this in-state_dict
+                        # branch is persistent (non-persistent buffers take the
+                        # not-in-state_dict path above and are lifted as constants).
+                        persistent=(True if input_kind == InputKind.BUFFER else None),
                     ),
                 )
                 non_user_input_idx += 1
@@ -245,29 +253,6 @@ def lift(
     return gm, graph_signature, state_dict, constants
 
 
-def get_duplicate_nodes(
-    gm: torch.fx.GraphModule, submodule: torch.fx.GraphModule
-) -> Tuple[Sequence[Any], Sequence[Any]]:
-    """
-    We check if there are duplicate nodes when we copy submodule graph into gm.
-    Handle the case where the subgraph input placeholders are same as
-    gm placeholders. This happens when the first submodule in the graph is
-    a pytorch submodule
-    """
-    submodule_placeholder_inputs = [
-        node for node in submodule.graph.nodes if node.op == "placeholder"
-    ]
-    submodule_input_node_names = [node.name for node in submodule_placeholder_inputs]
-    gm_node_names = [node.name for node in gm.graph.nodes]
-    submodule_duplicate_inputs = [
-        node for node in submodule_placeholder_inputs if node.name in gm_node_names
-    ]
-    gm_duplicate_inputs = [
-        node for node in gm.graph.nodes if node.name in submodule_input_node_names
-    ]
-    return submodule_duplicate_inputs, gm_duplicate_inputs
-
-
 def inline_torch_modules(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
     """
     Inline a submodule within the parent graph (gm). All `call_module` nodes
@@ -285,43 +270,31 @@ def inline_torch_modules(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
                 # or a placeholder of the main graph
                 submodule_inputs = gm_node.args
 
-                submodule_duplicate_inputs, gm_duplicate_inputs = get_duplicate_nodes(
-                    gm, submodule
-                )
-                assert len(submodule_duplicate_inputs) == len(gm_duplicate_inputs)
-                # Avoid creating new copies of duplicate inputs by creating a mapping
-                val_map = {}
-                for i in range(len(submodule_duplicate_inputs)):
-                    val_map[submodule_duplicate_inputs[i]] = gm_duplicate_inputs[i]
-
-                # Copy all nodes in the submodule into gm and
-                # store the output node of this submodule which is now present in gm
+                # Copy the submodule's nodes into gm, then wire its inputs POSITIONALLY.
+                #
+                # We deliberately do NOT pre-seed val_map by matching submodule input
+                # placeholders to gm nodes by NAME. Name matching silently binds an
+                # input to the WRONG node when names collide (e.g. a _run_on_gpu input
+                # placeholder whose name matches a different engine's getitem), which
+                # rewires a consumer to the wrong producer and orphans the real one --
+                # the orphan is then pruned by dead-code elimination, leaving a delegate
+                # short an output at runtime. It also leaks a mixed submodule's
+                # computed-intermediate inputs as spurious graph inputs. gm_node.args is
+                # the authoritative, ordered list of the real inputs, so we let
+                # graph_copy create a fresh (auto-renamed on collision) placeholder for
+                # every submodule input and rewire each to submodule_inputs[i] by
+                # position, then erase it.
+                val_map: Dict[Any, Any] = {}
                 submodule_output = gm.graph.graph_copy(submodule.graph, val_map)
 
-                # Get their references (since we copied) in the parent graph (gm)
-                if len(submodule_duplicate_inputs) == 0:
-                    submodule_placeholder_input_names = [
-                        node.name
-                        for node in submodule.graph.nodes
-                        if node.op == "placeholder"
-                    ]
-                    gm_added_placeholder_inputs = [
-                        node
-                        for node in gm.graph.nodes
-                        if node.name in submodule_placeholder_input_names
-                    ]
-
-                    assert len(submodule_inputs) == len(gm_added_placeholder_inputs)
-
-                    # Replace the added placeholder inputs with original inputs to this submodule node
-                    for idx in range(len(gm_added_placeholder_inputs)):
-                        gm_added_placeholder_inputs[idx].replace_all_uses_with(
-                            submodule_inputs[idx]
-                        )
-
-                    # Erase the placeholder input nodes in the gm
-                    for idx in range(len(gm_added_placeholder_inputs)):
-                        gm.graph.erase_node(gm_added_placeholder_inputs[idx])
+                submodule_placeholders = [
+                    node for node in submodule.graph.nodes if node.op == "placeholder"
+                ]
+                assert len(submodule_placeholders) == len(submodule_inputs)
+                for idx, submodule_placeholder in enumerate(submodule_placeholders):
+                    copied_placeholder = val_map[submodule_placeholder]
+                    copied_placeholder.replace_all_uses_with(submodule_inputs[idx])
+                    gm.graph.erase_node(copied_placeholder)
 
                 # Replace the pytorch submodule node (call_module) with the inlined subgraph output
                 # Special handling when submodule returns multiple outputs (tuple)
@@ -403,14 +376,29 @@ def create_trt_exp_program(
         input_specs=input_specs, output_specs=output_specs
     )
 
+    # A hybrid TRT+CUDA GraphModule from dynamo.compile carries a plain fx.CodeGen
+    # (no pytree_info): the module already returns a flat tuple, so rebuild
+    # in_spec/out_spec from the example inputs and the flat graph outputs. This
+    # matches retrace=True, which traces the same flat module and likewise cannot
+    # re-nest -- so the fallback never diverges from it.
+    codegen = gm.graph._codegen
+    if isinstance(codegen, _PyTreeCodeGen):
+        in_spec = codegen.pytree_info.in_spec
+        out_spec = codegen.pytree_info.out_spec
+    else:
+        example_args = tuple(arg_inputs) if arg_inputs is not None else ()
+        example_kwargs = kwarg_inputs or {}
+        in_spec = pytree.tree_flatten((example_args, example_kwargs))[1]
+        out_spec = pytree.tree_flatten(tuple(output_nodes))[1]
+
     module_call_graph = [
         ModuleCallEntry(
             "",
             ModuleCallSignature(
                 inputs=[],
                 outputs=[],
-                in_spec=gm.graph._codegen.pytree_info.in_spec,
-                out_spec=gm.graph._codegen.pytree_info.out_spec,
+                in_spec=in_spec,
+                out_spec=out_spec,
             ),
         )
     ]

tests/py/dynamo/models/test_exporter_inlining.py

@@ -0,0 +1,127 @@
+"""Unit tests for the legacy dynamo exporter's submodule inlining
+(torch_tensorrt.dynamo._exporter). These run on plain fx graphs and need neither a
+GPU nor a TensorRT build."""
+
+import operator
+
+import pytest
+import torch
+from torch_tensorrt.dynamo._exporter import inline_torch_modules
+
+
+@pytest.mark.unit
+def test_inline_torch_modules_wires_inputs_by_position():
+    """inline_torch_modules must wire a _run_on_gpu submodule's inputs from the
+    call_module args by POSITION, not by matching placeholder names to graph nodes.
+
+    Regression: the old name-matching path bound a submodule input to a same-named
+    but unrelated graph node, rewiring a consumer to the wrong producer (and, for a
+    submodule mixing graph-input and computed-intermediate inputs, leaking the
+    latter as spurious graph placeholders). Here the submodule's first input
+    placeholder is named "y", colliding with the parent's second input "y" even
+    though the first *argument* is the parent's "x"; positional wiring must ignore
+    the collision. Subtraction makes the input order observable.
+    """
+    # Submodule: out = first - second. First placeholder deliberately named "y".
+    sub_graph = torch.fx.Graph()
+    first = sub_graph.placeholder("y")
+    second = sub_graph.placeholder("z")
+    sub_graph.output(sub_graph.call_function(torch.sub, (first, second)))
+    submodule = torch.fx.GraphModule(torch.nn.Module(), sub_graph)
+
+    # Parent: inputs (x, y); call _run_on_gpu_0(x, y) -> expected x - y.
+    parent_graph = torch.fx.Graph()
+    x = parent_graph.placeholder("x")
+    y = parent_graph.placeholder("y")
+    root = torch.nn.Module()
+    root.add_module("_run_on_gpu_0", submodule)
+    call = parent_graph.call_module("_run_on_gpu_0", (x, y))
+    parent_graph.output(call)
+    parent = torch.fx.GraphModule(root, parent_graph)
+
+    n_placeholders_before = sum(1 for n in parent.graph.nodes if n.op == "placeholder")
+
+    inline_torch_modules(parent)
+    parent.recompile()
+
+    # No spurious placeholders leaked by the inlining.
+    assert (
+        sum(1 for n in parent.graph.nodes if n.op == "placeholder")
+        == n_placeholders_before
+    )
+    # No call_module node survives (the submodule was inlined).
+    assert not any(n.op == "call_module" for n in parent.graph.nodes)
+    # Positional wiring: first input <- x, second input <- y, so out == x - y.
+    out = parent(torch.tensor(5.0), torch.tensor(3.0))
+    assert torch.allclose(out, torch.tensor(2.0))
+
+
+@pytest.mark.unit
+def test_inline_torch_modules_preserves_all_submodule_outputs():
+    """A multi-output _run_on_gpu submodule must keep every output wired to its
+    consumer after inlining. Regression: a mis-wired input orphaned one submodule
+    output, which dead-code elimination then pruned, leaving a downstream consumer
+    (or, in the hybrid case, a TensorRT engine) short an output at runtime.
+    """
+    # Submodule returns (a + b, a - b); both outputs are consumed downstream.
+    sub_graph = torch.fx.Graph()
+    a = sub_graph.placeholder("a")
+    b = sub_graph.placeholder("b")
+    add = sub_graph.call_function(torch.add, (a, b))
+    sub = sub_graph.call_function(torch.sub, (a, b))
+    sub_graph.output((add, sub))
+    submodule = torch.fx.GraphModule(torch.nn.Module(), sub_graph)
+
+    parent_graph = torch.fx.Graph()
+    x = parent_graph.placeholder("x")
+    y = parent_graph.placeholder("y")
+    root = torch.nn.Module()
+    root.add_module("_run_on_gpu_0", submodule)
+    call = parent_graph.call_module("_run_on_gpu_0", (x, y))
+    o0 = parent_graph.call_function(operator.getitem, (call, 0))
+    o1 = parent_graph.call_function(operator.getitem, (call, 1))
+    # Consume both outputs: (a+b) * (a-b).
+    parent_graph.output(parent_graph.call_function(torch.mul, (o0, o1)))
+    parent = torch.fx.GraphModule(root, parent_graph)
+
+    inline_torch_modules(parent)
+    parent.recompile()
+
+    # (x+y)*(x-y) == x^2 - y^2 ; with x=5, y=3 -> 25 - 9 = 16.
+    out = parent(torch.tensor(5.0), torch.tensor(3.0))
+    assert torch.allclose(out, torch.tensor(16.0))
+
+
+@pytest.mark.unit
+def test_inline_torch_modules_computed_intermediate_inputs():
+    """A _run_on_gpu submodule whose inputs are computed intermediates (not top-level
+    graph placeholders, and not name-matching any graph node) must inline correctly.
+    This is the case the old zero-duplicate path handled; positional wiring preserves
+    it (and it is the shape that leaked spurious placeholders in the mixed case).
+    """
+    # Submodule: out = m + n. Names don't collide with anything in the parent.
+    sub_graph = torch.fx.Graph()
+    m = sub_graph.placeholder("m")
+    n = sub_graph.placeholder("n")
+    sub_graph.output(sub_graph.call_function(torch.add, (m, n)))
+    submodule = torch.fx.GraphModule(torch.nn.Module(), sub_graph)
+
+    # Parent: x -> c0 = x*2, c1 = x+1; call _run_on_gpu_0(c0, c1) -> c0 + c1.
+    parent_graph = torch.fx.Graph()
+    x = parent_graph.placeholder("x")
+    c0 = parent_graph.call_function(torch.mul, (x, 2))
+    c1 = parent_graph.call_function(torch.add, (x, 1))
+    root = torch.nn.Module()
+    root.add_module("_run_on_gpu_0", submodule)
+    call = parent_graph.call_module("_run_on_gpu_0", (c0, c1))
+    parent_graph.output(call)
+    parent = torch.fx.GraphModule(root, parent_graph)
+
+    inline_torch_modules(parent)
+    parent.recompile()
+
+    # No spurious placeholders leaked; the computed intermediates stay in-graph.
+    assert sum(1 for node in parent.graph.nodes if node.op == "placeholder") == 1
+    # out = (x*2) + (x+1); x=5 -> 10 + 6 = 16.
+    out = parent(torch.tensor(5.0))
+    assert torch.allclose(out, torch.tensor(16.0))