fix: bind ExecuTorch TRT engine inputs in delegate arg order

## The problem

When a model with more than one input is compiled to a TensorRT engine and
run through ExecuTorch, the inputs can be delivered to the wrong engine slots.
The failure looks like this at runtime:

    input 'timesteps' rank 4 does not match profile rank 1

The model is fine -- the wrong tensor is simply arriving at that slot.

Why it happens: the ExecuTorch TensorRT runtime backend matches inputs to
engine bindings *by position*. Argument `i` handed to `execute()` is fed to
`input_binding_names[i]`. There are no names attached to the incoming tensors
at runtime; the backend just trusts that the i-th tensor belongs to the i-th
binding name.

But those two lists are built by two different stages that never coordinate:

1. `input_binding_names` is recorded during TensorRT conversion, in the order
   the interpreter visits the subgraph placeholders
   (`_TRTInterpreter.placeholder`).
2. The tensors passed to `executorch_call_delegate` at runtime are ordered by
   ExecuTorch's own lowering: its FX fuser sorts the delegate's placeholder
   inputs by top-level graph order, which can be a different permutation.

For a single-input engine there is nothing to reorder, so it always works.
For a multi-input engine whose inputs get laid out differently after lowering,
the two orders disagree and a tensor lands in the wrong binding -- which either
fails loudly (rank/shape mismatch, as above) or, if two inputs happen to share
a shape, silently produces wrong results.

## Alternatives considered

- Match tensors to bindings by name at runtime. Not possible: the ExecuTorch
  delegate ABI hands `execute()` only positional values, with no names. Adding
  names would require changing the serialized blob format and the runtime ABI.

- Match names to placeholders by name at serialization time. Does not work: the
  TensorRT binding names are the original semantic placeholder targets (e.g.
  `timesteps`, `position_ids`) while the lowered ExecuTorch placeholders are
  generic (`arg_0`, `arg_1`, ...). There is no reliable name correspondence.

- Fix the order earlier, during TensorRT conversion/export. Does not stick:
  ExecuTorch performs another lowering/fusion pass afterward that can re-permute
  the placeholders, so any alignment done before that pass is undone.

## The solution

Reorder `input_binding_names` inside the ExecuTorch backend `preprocess()` --
the first point where both the TensorRT binding order and the final delegate
argument order are visible. The permutation is recovered by *node identity*,
not by name: the engine node's first argument lists its input nodes in binding
order, and each of those nodes is a placeholder in the lowered graph whose
position is exactly the runtime argument slot it will occupy. So we sort the
binding names by each input node's placeholder position.

Single-input engines are naturally identity, so models that already worked are
unaffected. The runtime is unchanged, and the per-binding optimization-profile
bounds follow automatically because they are looked up by name at engine init
(`getProfileShape`).

Adds unit tests in `tests/py/dynamo/executorch/test_backend.py`: single-input
identity, node-identity reorder, generic-name multi-input permutation, and a
guard for a binding-count mismatch.

Author: shoumikhin

py/torch_tensorrt/executorch/backend.py

@@ -198,6 +198,43 @@ def _parse_device_id(value: Any) -> int:
         return 0
 
 
+def _reorder_input_names_for_executorch(
+    edge_program: ExportedProgram, input_names: List[str]
+) -> List[str]:
+    """Reorder TRT binding names into executorch_call_delegate argument order.
+
+    The runtime binds positionally (``execute()`` arg ``i`` -> input_binding_names
+    ``[i]``), but ExecuTorch fusion may permute the delegate placeholders relative
+    to the TRT-submodule order that produced ``input_binding_names``. The names
+    can't be matched (TRT names are semantic, lowered placeholders are generic
+    ``arg_N``), so recover the permutation by node identity: the engine node's
+    first arg lists its input nodes in binding order, so sort the names by each
+    node's slot among the graph placeholders (its runtime delegate-arg position).
+    """
+    input_nodes = list(_get_engine_nodes_from_edge_program(edge_program)[0].args[0])
+    if len(input_nodes) != len(input_names):
+        raise ValueError(
+            "TensorRT ExecuTorch backend: engine has "
+            f"{len(input_names)} input binding names but {len(input_nodes)} "
+            "engine input nodes; cannot establish a reliable binding order."
+        )
+    slot = {
+        node: i
+        for i, node in enumerate(
+            n for n in edge_program.graph_module.graph.nodes if n.op == "placeholder"
+        )
+    }
+    missing = [n for n in input_nodes if n not in slot]
+    if missing:
+        raise ValueError(
+            "TensorRT ExecuTorch backend: engine inputs "
+            f"{[n.name for n in missing]} are not delegate runtime placeholders; "
+            "cannot determine their argument position."
+        )
+    order = sorted(range(len(input_nodes)), key=lambda i: slot[input_nodes[i]])
+    return [input_names[i] for i in order]
+
+
 def _get_str(engine_info: List[Any], index: int, default: str = "") -> str:
     if index < 0 or index >= len(engine_info):
         return default
@@ -240,8 +277,9 @@ def preprocess(
             )
         elif not isinstance(serialized_engine, (bytes, bytearray)):
             engine_info[ENGINE_IDX] = bytes(serialized_engine)
-        input_names = _split_binding_names(
-            _get_str(engine_info, INPUT_BINDING_NAMES_IDX)
+        input_names = _reorder_input_names_for_executorch(
+            edge_program,
+            _split_binding_names(_get_str(engine_info, INPUT_BINDING_NAMES_IDX)),
         )
         output_names = _split_binding_names(
             _get_str(engine_info, OUTPUT_BINDING_NAMES_IDX)

tests/py/dynamo/executorch/test_backend.py

@@ -5,6 +5,7 @@
 executorch = pytest.importorskip("executorch.exir")
 
 import torch  # noqa: E402
+from torch.export.graph_signature import InputKind  # noqa: E402
 from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import (  # noqa: E402
     DEVICE_IDX,
     ENGINE_IDX,
@@ -24,22 +25,82 @@
 
 
 class _SchemaTarget:
+    """Callable stand-in for a tensorrt:: op overload carrying a `_schema`.
+
+    fx `call_function` requires a callable target; the backend only reads
+    `target._schema.name`, so a no-op callable with the schema attribute is
+    enough to build a realistic engine node in a real fx graph.
+    """
+
     def __init__(self, name):
         self._schema = SimpleNamespace(name=name)
+        self.__name__ = name.replace("::", "_")
+
+    def __call__(self, *args, **kwargs):  # pragma: no cover - never executed
+        raise AssertionError("engine op target is not meant to be called")
+
+
+_ENGINE_OP = _SchemaTarget("tensorrt::no_op_placeholder_for_execute_engine")
+
 
+def _build_edge_program(
+    engine_info,
+    input_names=("x",),
+    binding_order=None,
+):
+    """Construct a real fx-graph-backed edge program for the backend.
 
-def _make_placeholder_node(*engine_info):
+    The engine node's first arg is the ordered list of input placeholder NODES
+    (mirroring torch_tensorrt's inline_trt_modules), in `binding_order` — the
+    TRT binding order. The graph placeholders themselves are laid out in
+    `input_names` order — the fusion-arranged runtime arg order. The backend must
+    recover the permutation from node identity, not from names.
+
+    A bare `torch.fx.Graph` (not a GraphModule) is used so the nodes carry real
+    fx identity — which is what the reorder logic keys on — without triggering
+    GraphModule codegen of the fake engine op / raw engine tensor.
+    """
+    if binding_order is None:
+        binding_order = tuple(input_names)
+
+    graph = torch.fx.Graph()
+    placeholders = {name: graph.placeholder(name) for name in input_names}
+    engine_inputs = [placeholders[name] for name in binding_order]
+    engine_node = graph.call_function(
+        _ENGINE_OP,
+        (engine_inputs, *engine_info),
+    )
+    graph.output((engine_node,))
+
+    graph_signature = SimpleNamespace(
+        input_specs=[
+            SimpleNamespace(kind=InputKind.USER_INPUT, arg=SimpleNamespace(name=name))
+            for name in input_names
+        ]
+    )
     return SimpleNamespace(
-        op="call_function",
-        target=_SchemaTarget("tensorrt::no_op_placeholder_for_execute_engine"),
-        args=(["x"], *engine_info),
-        name="trt_node",
+        graph_module=SimpleNamespace(graph=graph),
+        graph_signature=graph_signature,
+        constants={},
     )
 
 
-def _make_edge_program(*nodes):
+def _make_multi_engine_program(engine_info):
+    """An edge program whose graph contains two engine nodes (invalid)."""
+    graph = torch.fx.Graph()
+    x = graph.placeholder("x")
+    n1 = graph.call_function(_ENGINE_OP, ([x], *engine_info))
+    n2 = graph.call_function(_ENGINE_OP, ([x], *engine_info))
+    graph.output((n1, n2))
     return SimpleNamespace(
-        graph_module=SimpleNamespace(graph=SimpleNamespace(nodes=list(nodes))),
+        graph_module=SimpleNamespace(graph=graph),
+        graph_signature=SimpleNamespace(
+            input_specs=[
+                SimpleNamespace(
+                    kind=InputKind.USER_INPUT, arg=SimpleNamespace(name="x")
+                )
+            ]
+        ),
         constants={},
     )
 
@@ -52,10 +113,7 @@ def _engine_tensor(payload: bytes) -> torch.Tensor:
 def test_get_engine_info_rejects_multiple_engine_nodes():
     engine_info = [""] * SERIALIZATION_LEN
     engine_info[ENGINE_IDX] = _engine_tensor(b"engine")
-    edge_program = _make_edge_program(
-        _make_placeholder_node(*engine_info),
-        _make_placeholder_node(*engine_info),
-    )
+    edge_program = _make_multi_engine_program(engine_info)
 
     with pytest.raises(RuntimeError, match="exactly 1 engine node"):
         _get_engine_info_from_edge_program(edge_program)
@@ -66,7 +124,7 @@ def test_preprocess_rejects_output_allocator():
     engine_info = [""] * SERIALIZATION_LEN
     engine_info[ENGINE_IDX] = _engine_tensor(b"engine")
     engine_info[REQUIRES_OUTPUT_ALLOCATOR_IDX] = "1"
-    edge_program = _make_edge_program(_make_placeholder_node(*engine_info))
+    edge_program = _build_edge_program(engine_info)
 
     with pytest.raises(RuntimeError, match="output allocator"):
         TensorRTBackend.preprocess(edge_program, [])
@@ -79,7 +137,7 @@ def test_preprocess_serializes_engine_blob():
     engine_info[DEVICE_IDX] = "2%8%0%0%GPU"
     engine_info[INPUT_BINDING_NAMES_IDX] = "x"
     engine_info[OUTPUT_BINDING_NAMES_IDX] = "y"
-    edge_program = _make_edge_program(_make_placeholder_node(*engine_info))
+    edge_program = _build_edge_program(engine_info)
 
     result = TensorRTBackend.preprocess(edge_program, [])
 
@@ -90,3 +148,102 @@ def test_preprocess_serializes_engine_blob():
     assert metadata.device_id == 2
     assert [binding.name for binding in metadata.io_bindings] == ["x", "y"]
     assert [binding.is_input for binding in metadata.io_bindings] == [True, False]
+
+
+@pytest.mark.unit
+def test_preprocess_single_input_is_identity():
+    # Single-input engines have zero ordering ambiguity: the one binding maps to
+    # the one placeholder regardless of name (TRT name may be semantic, the
+    # runtime placeholder generic). Must never reorder or reject.
+    engine_info = [""] * SERIALIZATION_LEN
+    engine_info[ENGINE_IDX] = _engine_tensor(b"engine-bytes")
+    engine_info[INPUT_BINDING_NAMES_IDX] = "pixel_values"
+    engine_info[OUTPUT_BINDING_NAMES_IDX] = "out"
+    # Graph placeholder is generic ("arg_0"); the binding name differs entirely.
+    edge_program = _build_edge_program(engine_info, input_names=("arg_0",))
+
+    _, metadata = deserialize_engine(
+        TensorRTBackend.preprocess(edge_program, []).processed_bytes
+    )
+    assert [binding.name for binding in metadata.io_bindings] == [
+        "pixel_values",
+        "out",
+    ]
+
+
+@pytest.mark.unit
+def test_preprocess_reorders_by_node_identity_not_name():
+    # TRT binding order is [second, first]; the engine node feeds those input
+    # placeholders in that order. The graph lays the placeholders out as
+    # [first, second] (the runtime delegate arg order). The backend must emit
+    # the binding names permuted to runtime order: [first, second].
+    engine_info = [""] * SERIALIZATION_LEN
+    engine_info[ENGINE_IDX] = _engine_tensor(b"engine-bytes")
+    engine_info[INPUT_BINDING_NAMES_IDX] = "second%first"
+    engine_info[OUTPUT_BINDING_NAMES_IDX] = "out"
+    edge_program = _build_edge_program(
+        engine_info,
+        input_names=("first", "second"),
+        binding_order=("second", "first"),
+    )
+
+    _, metadata = deserialize_engine(
+        TensorRTBackend.preprocess(edge_program, []).processed_bytes
+    )
+    assert [binding.name for binding in metadata.io_bindings] == [
+        "first",
+        "second",
+        "out",
+    ]
+    assert [binding.is_input for binding in metadata.io_bindings] == [
+        True,
+        True,
+        False,
+    ]
+
+
+@pytest.mark.unit
+def test_preprocess_generic_names_multi_input_permutation():
+    # Realistic lowered case: TRT names are semantic, runtime placeholders are
+    # generic arg_N and DO NOT match the TRT names. Reorder must still work via
+    # node identity. Binding order [timesteps, prefix, x]; graph placeholder
+    # (runtime) order is [arg_0=x, arg_1=timesteps, arg_2=prefix].
+    engine_info = [""] * SERIALIZATION_LEN
+    engine_info[ENGINE_IDX] = _engine_tensor(b"engine-bytes")
+    engine_info[INPUT_BINDING_NAMES_IDX] = "timesteps%prefix%x"
+    engine_info[OUTPUT_BINDING_NAMES_IDX] = "out"
+    # runtime placeholder order:      arg_0, arg_1,     arg_2
+    # semantic meaning of each:       x,     timesteps, prefix
+    # engine feeds bindings in order: timesteps(arg_1), prefix(arg_2), x(arg_0)
+    edge_program = _build_edge_program(
+        engine_info,
+        input_names=("arg_0", "arg_1", "arg_2"),
+        binding_order=("arg_1", "arg_2", "arg_0"),
+    )
+
+    _, metadata = deserialize_engine(
+        TensorRTBackend.preprocess(edge_program, []).processed_bytes
+    )
+    # Binding names are carried with their node; sorted by runtime slot the
+    # result is [name-fed-by-arg_0, name-fed-by-arg_1, name-fed-by-arg_2] =
+    # [x, timesteps, prefix].
+    assert [binding.name for binding in metadata.io_bindings] == [
+        "x",
+        "timesteps",
+        "prefix",
+        "out",
+    ]
+
+
+@pytest.mark.unit
+def test_preprocess_rejects_binding_count_mismatch():
+    # Defensive invariant: number of binding names must equal number of engine
+    # input nodes. Two names but one node -> refuse rather than corrupt.
+    engine_info = [""] * SERIALIZATION_LEN
+    engine_info[ENGINE_IDX] = _engine_tensor(b"engine-bytes")
+    engine_info[INPUT_BINDING_NAMES_IDX] = "a%b"
+    engine_info[OUTPUT_BINDING_NAMES_IDX] = "out"
+    edge_program = _build_edge_program(engine_info, input_names=("arg_0",))
+
+    with pytest.raises(ValueError, match="input binding names"):
+        TensorRTBackend.preprocess(edge_program, [])