Skip to content

Commit 1bc58ea

Browse files
committed
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.
1 parent 2aac435 commit 1bc58ea

2 files changed

Lines changed: 296 additions & 25 deletions

File tree

py/torch_tensorrt/executorch/backend.py

Lines changed: 62 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,17 @@ def _get_engine_nodes_from_edge_program(edge_program: ExportedProgram) -> List[A
5858
return _get_engine_nodes_in(edge_program.graph_module.graph.nodes)
5959

6060

61+
def _get_single_engine_node(edge_program: ExportedProgram) -> Any:
62+
"""Return the partition's one TRT engine node, or raise if not exactly one."""
63+
engine_nodes = _get_engine_nodes_from_edge_program(edge_program)
64+
if len(engine_nodes) != 1:
65+
raise RuntimeError(
66+
"TensorRT ExecuTorch backend expects exactly 1 engine node per "
67+
f"partition, found {len(engine_nodes)}."
68+
)
69+
return engine_nodes[0]
70+
71+
6172
def _get_engine_info_from_edge_program(edge_program: ExportedProgram) -> List[Any]:
6273
"""Extract engine info (list of strings/bytes) from the partition's TRT node.
6374
@@ -71,13 +82,9 @@ def _get_engine_info_from_edge_program(edge_program: ExportedProgram) -> List[An
7182
Uses schema name comparison (not object identity) so it works for both
7283
OpOverload and EdgeOpOverload targets.
7384
"""
74-
engine_nodes = _get_engine_nodes_from_edge_program(edge_program)
75-
if len(engine_nodes) != 1:
76-
raise RuntimeError(
77-
"TensorRT ExecuTorch backend expects exactly 1 engine node per "
78-
f"partition, found {len(engine_nodes)}."
79-
)
80-
return _get_engine_info_for_node(edge_program, engine_nodes[0])
85+
return _get_engine_info_for_node(
86+
edge_program, _get_single_engine_node(edge_program)
87+
)
8188

8289

8390
def _get_engine_info_for_node(
@@ -198,6 +205,48 @@ def _parse_device_id(value: Any) -> int:
198205
return 0
199206

200207

208+
def _reorder_input_names_for_executorch(
209+
edge_program: ExportedProgram, engine_node: Any, input_names: List[str]
210+
) -> List[str]:
211+
"""Reorder TRT binding names into executorch_call_delegate argument order.
212+
213+
The runtime binds positionally (``execute()`` arg ``i`` -> input_binding_names
214+
``[i]``), but ExecuTorch fusion may permute the delegate placeholders relative
215+
to the TRT-submodule order that produced ``input_binding_names``. The names
216+
can't be matched (TRT names are semantic, lowered placeholders are generic
217+
``arg_N``), so recover the permutation by node identity: the engine node's
218+
first arg lists its input nodes in binding order, so sort the names by each
219+
node's slot among the graph placeholders (its runtime delegate-arg position).
220+
221+
Only inputs need this. Outputs are also bound positionally by the runtime,
222+
but they are ``getitem(engine_node, idx)`` nodes whose index order is fixed
223+
at creation and equals the engine output-binding order, so ExecuTorch fusion
224+
(which only re-sorts external placeholder inputs) cannot permute them.
225+
"""
226+
input_nodes = list(engine_node.args[0])
227+
if len(input_nodes) != len(input_names):
228+
raise ValueError(
229+
"TensorRT ExecuTorch backend: engine has "
230+
f"{len(input_names)} input binding names but {len(input_nodes)} "
231+
"engine input nodes; cannot establish a reliable binding order."
232+
)
233+
slot = {
234+
node: i
235+
for i, node in enumerate(
236+
n for n in edge_program.graph_module.graph.nodes if n.op == "placeholder"
237+
)
238+
}
239+
missing = [n for n in input_nodes if n not in slot]
240+
if missing:
241+
raise ValueError(
242+
"TensorRT ExecuTorch backend: engine inputs "
243+
f"{[n.name for n in missing]} are not delegate runtime placeholders; "
244+
"cannot determine their argument position."
245+
)
246+
order = sorted(range(len(input_nodes)), key=lambda i: slot[input_nodes[i]])
247+
return [input_names[i] for i in order]
248+
249+
201250
def _get_str(engine_info: List[Any], index: int, default: str = "") -> str:
202251
if index < 0 or index >= len(engine_info):
203252
return default
@@ -223,7 +272,8 @@ def preprocess(
223272
edge_program: ExportedProgram,
224273
compile_specs: List[CompileSpec],
225274
) -> PreprocessResult:
226-
engine_info = _get_engine_info_from_edge_program(edge_program)
275+
engine_node = _get_single_engine_node(edge_program)
276+
engine_info = _get_engine_info_for_node(edge_program, engine_node)
227277
engine_info = list(engine_info)
228278
_validate_engine_info(engine_info)
229279
serialized_engine = engine_info[ENGINE_IDX]
@@ -240,8 +290,10 @@ def preprocess(
240290
)
241291
elif not isinstance(serialized_engine, (bytes, bytearray)):
242292
engine_info[ENGINE_IDX] = bytes(serialized_engine)
243-
input_names = _split_binding_names(
244-
_get_str(engine_info, INPUT_BINDING_NAMES_IDX)
293+
input_names = _reorder_input_names_for_executorch(
294+
edge_program,
295+
engine_node,
296+
_split_binding_names(_get_str(engine_info, INPUT_BINDING_NAMES_IDX)),
245297
)
246298
output_names = _split_binding_names(
247299
_get_str(engine_info, OUTPUT_BINDING_NAMES_IDX)

0 commit comments

Comments
 (0)