diff --git a/py/torch_tensorrt/dynamo/_exporter.py b/py/torch_tensorrt/dynamo/_exporter.py index c6c992d720..34c5b5b838 100644 --- a/py/torch_tensorrt/dynamo/_exporter.py +++ b/py/torch_tensorrt/dynamo/_exporter.py @@ -386,6 +386,41 @@ def copy_submodule_attributes( _assign_attr(value, gm, key, _AttrKind.BUFFER) +def _order_placeholders_first(gm: torch.fx.GraphModule) -> None: + """Hoist every placeholder to the front of the node list, keeping their + relative order. + + ``make_constraints()`` indexes ``flat_dynamic_shapes`` by *node* position while + walking ``enumerate(gm.graph.nodes)``, so it silently requires placeholders to lead + the graph. For ``nn.Linear(10, 5)`` with one dynamic input:: + + converted (placeholders lead) fully fallen back to PyTorch + [0] placeholder x [0] get_attr linear_weight + [1] get_attr engine_0 [1] get_attr linear_bias + [2] call_function ... [2] placeholder x + [3] output [3] call_function linear + [4] output + + Unlifting puts the parameters ahead of the user input, so on the right the lone + input sits at node index 2 and the lookup becomes ``flat_dynamic_shapes[2]`` on a + one-element list: ``save()`` dies with a bare ``IndexError`` while its own length + check passes. Placeholders take no arguments, so hoisting cannot break topological + order, and preserving their relative order leaves the forward signature, + ``input_nodes``, ``in_spec`` and InputSpec ordering unchanged. + """ + nodes = list(gm.graph.nodes) + placeholders = [node for node in nodes if node.op == "placeholder"] + if not placeholders or nodes[: len(placeholders)] == placeholders: + return + + for node in reversed(placeholders): + first = next(iter(gm.graph.nodes)) + if node is not first: + first.prepend(node) + + gm.graph.lint() + + def create_trt_exp_program( gm: torch.fx.GraphModule, *, @@ -397,6 +432,8 @@ def create_trt_exp_program( and constructs an Exported Program object with the new IO node names and state_dict """ + _order_placeholders_first(gm) + input_nodes = [node for node in gm.graph.nodes if node.op == "placeholder"] output_nodes = [node for node in gm.graph.nodes if node.op == "output"] assert output_nodes diff --git a/tests/py/dynamo/models/test_exporter_inlining.py b/tests/py/dynamo/models/test_exporter_inlining.py index 2fa2f0470e..9dad5de16b 100644 --- a/tests/py/dynamo/models/test_exporter_inlining.py +++ b/tests/py/dynamo/models/test_exporter_inlining.py @@ -71,7 +71,7 @@ def test_inline_torch_modules_wires_inputs_by_position(): @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 + consumer after inlining. Regression: a miswired 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. @@ -231,6 +231,64 @@ def forward(self, a, b): assert torch.allclose(out, torch.tensor(7.0)) +@pytest.mark.unit +def test_create_trt_exp_program_handles_get_attrs_before_placeholders(): + """create_trt_exp_program must build range_constraints on a graph whose + get_attr nodes precede its user input. + + Regression: when the partitioner sends ops back to PyTorch the parameters + come back as get_attr nodes AHEAD of the user inputs. torch's + make_constraints() indexes flat_dynamic_shapes by NODE position, not by + placeholder ordinal, so the lone input sitting at node index 2 made it read + flat_dynamic_shapes[2] on a one-element list and torch_tensorrt.save() died + with a bare "IndexError: list index out of range". Its own length check + passes, so nothing pointed at the real problem. An exported module reproduces + the node order exactly. + """ + + class Linear(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(10, 5) + + def forward(self, x): + return self.linear(x) + + batch = torch.export.Dim("batch", min=2, max=8) + gm = torch.export.export( + Linear().eval(), (torch.randn(4, 10),), dynamic_shapes={"x": {0: batch}} + ).module() + # A compiled TRT GraphModule carries no _guards_fn; drop it so the graph + # matches what create_trt_exp_program is handed in the normal flow. + for node in list(gm.graph.nodes): + if node.op == "call_module" and node.target == "_guards_fn": + gm.graph.erase_node(node) + break + gm.graph.lint() + gm.recompile() + + nodes = list(gm.graph.nodes) + assert nodes[0].op == "get_attr" + assert [n.op for n in nodes].index("placeholder") > 0 + + ep = create_trt_exp_program( + gm, arg_inputs=(torch.randn(4, 10),), dynamic_shapes={"x": {0: batch}} + ) + + # The user-specified Dim bound survived, so path 1 (make_constraints) ran. + assert [str(vr) for vr in ep.range_constraints.values()] == ["VR[2, 8]"] + assert [s.kind.name for s in ep.graph_signature.input_specs] == [ + "PARAMETER", + "PARAMETER", + "USER_INPUT", + ] + + # And the program still runs at another batch size within the range. + out = ep.module()(torch.randn(3, 10)) + out = out[0] if isinstance(out, (tuple, list)) else out + assert out.shape == (3, 5) + + @pytest.mark.unit def test_lift_sets_persistent_true_on_buffer_spec(): """lift() must mark a lifted BUFFER InputSpec as persistent.