perf: skip the base64 round trip when reading TensorRT engine info

[Conarnar]

Description

TRTEngine::serialize() base64-encodes the serialized engine so the whole record can travel as a std::vector<std::string>. That encoding is the only reason the engine is a string, and both ExecuTorch export consumers pay for it without wanting it:

• validate_engine_program reads only flags and binding names, but reaching them goes through __getstate__, which serializes the entire ICudaEngine and base64-encodes it first.
• replace_execute_engine wants the engine as a uint8 tensor, so it pays the encode in C++ and then an immediate matching decode in Python — an encode plus a decode of a multi-megabyte blob, purely to move bytes through a string.

This adds the two halves as accessors and points the consumers at them.

C++ (core/runtime)

accessor	returns
serialize_metadata_only()	every serialized_info field except ENGINE_IDX. All plain members, so no engine serialization happens at all.
serialized_engine_tensor()	the engine as a uint8at::Tensor.

serialized_engine_tensor() returns at::Tensor rather than std::string because TorchScript maps std::string to a Python str and engine bytes are not valid UTF-8 (UnicodeDecodeError); a tensor is also what the consumer builds anyway.

serialize() now fills ENGINE_IDX into serialize_metadata_only()'s record rather than carrying a second copy of the field list, so the two cannot drift apart when a field is added.

Python (py/torch_tensorrt/executorch/_export_utils.py)

A metadata_only=True kwarg threaded through get_engine_info_from_state / _resolve_engine_info routes metadata readers to the cheap accessor, and _resolve_engine_tensor takes the engine as a tensor directly (which also drops the torch.frombuffer rebuild).

The other engine reader, _get_engine_info_for_node in backend.py, is left alone. Its one caller (partitioner.py) runs after replace_execute_engine has rewritten those nodes to no_op_placeholder, whose engine argument is already a tensor, so it takes the node.args[1:] branch and serializes nothing. The execute_engine branch below it would still serialize, but no caller reaches it in the current pass order; threading metadata_only through it would change no work today.

Results

Exporting a model with a 67 MB engine; six runs per configuration, interleaved, one fresh process per run, medians with min–max. A is this PR's base #4440; B is that base plus these two commits.

Measured against the base (#4440), where the change lands — equivalent to a main-relative measurement here, since both resolve each engine's info once and pay the single base64 round trip this PR removes. #4440 already includes #4473 (fbb10c9565), whose bulk engine copy sits in both columns and does not move the delta.

	A (baseline)	B (this PR)
ExecuTorch lowering phase	1.81 s	1.03 s (−43%)
passes that read engine info	0.82 s	0.09 s
engine-state reads	0.53 s	0.05 s
Python-side base64 decode	0.24 s	~0

Validation drops from 0.53 s to about 0.1 ms: it reads only metadata now, and the metadata accessor touches no engine bytes. The rewrite still reads the engine, but as a tensor without the base64 round trip, so its engine read is 0.05 s rather than a 0.53 s __getstate__ plus a 0.24 s decode. The engine's raw serialization is about 0.08 s of a 0.57 s __getstate__, so most of what disappears is the base64 encode and the string copies feeding it. A metadata-only read costs about 10 µs.

The separation is clean — every B run is faster than every A run on each metric — so a permutation test on the lowering phase gives p ≈ 0.002, the floor for six-versus-six fully separated samples.

Neither #4440 nor #4473 overlaps in scope with this change.

Compatibility

• serialize() output is byte-identical, so existing serialized engines and .pte artifacts are unaffected.
• serialized_engine_tensor() wraps TensorRT's buffer with at::from_blob rather than copying it, so peak host memory during the call is one engine rather than two — which matters for the multi-GB engines the untyped_storage() comment in backend.py's preprocess already accounts for. The consequence is that the returned tensor's storage is not resizable; nothing on the export path resizes it.
• The accessors are purely additive; nothing existing changes signature or behavior.
• The Python side does not require a matching runtime. Both accessors fall back to the old path when absent, which happens only when the Torch-TensorRT C++ library is older than the Python package — a source or editable build where only the Python half was rebuilt. A wheel ships both halves together. Since the fallback is correct and merely slower, the symptom would otherwise be silence, so it logs a warning once per accessor naming the cause.

Dependency

Stacked on #4440 (composable Edge export API), which introduces py/torch_tensorrt/executorch/_export_utils.py. The C++ commit stands alone; the Python commit needs that file to exist.

Related work

#4489 is orthogonal and complementary. It stops fakification from serializing the engine through TRTEngine::__obj_flatten__; this PR removes the base64 round trip in the engine-info reads. Both cut export-time engine serialization, but at different points, and they compose — neither depends on the other.

Type of change

• New feature (non-breaking change which adds functionality)

Checklist:

• My code follows the style guidelines of this project (You can use the linters)
• I have performed a self-review of my own code
• I have commented my code, particularly in hard-to-understand areas and hacks
• I have made corresponding changes to the documentation
• I have added tests to verify my fix or my feature
• New and existing unit tests pass locally with my changes
• I have added the relevant labels to my PR in so that relevant reviewers are notified

[Gasoonjia]

thx for the update! What's the major component of ExecuTorch lowering phase now? aoti_compile?

[shoumikhin]

Reviewed this and found no defects. The tests are unusually good, they assert the performance property itself rather than only correctness. Notes, then two questions.

The memory handling in serialized_engine_tensor() is correct, which was my main worry.at::from_blob borrows memory instead of copying, so the TensorRT buffer has to outlive the tensor. I compiled a small probe to confirm the captured deleter does that:

source shared_ptr goes out of scope -> buffer NOT freed (tensor still holds it), data readable
tensor released -> buffer freed exactly once

So skipping at::empty plus memcpy really does avoid holding two copies of a multi-megabyte blob. One thing worth stating in a comment: this storage is not resizable, so any consumer calling resize_() on it will throw. The current consumer just registers it as a buffer, which is fine.

Having serialize() call serialize_metadata_only() is the right structure, because one field list means a newly added index cannot land in one function and be silently missing from the other.

I ran the Python logic directly (loading _export_utils.py by path, since my environment cannot import the matching native runtime):

metadata_only, new runtime : uses the cheap accessor, __getstate__ not called PASS
metadata_only, old runtime : falls back to __getstate__, returns real record PASS
accessor vs fallback : differ in exactly one slot, ENGINE_IDX PASS
without metadata_only : always uses __getstate__ PASS
missing-accessor warning : fires once across three calls PASS

The third one is the property that makes the fallback safe, so it is good that you have it as a test.

One latent footgun. With metadata_only=True, ENGINE_IDX is "" on a new runtime and a full base64 engine on an old one, and nothing about the record's shape says which you got. You documented it and guarded the one place it matters, so this is not a bug today, it just relies on every future caller reading the docstring. Returning None instead of "" would turn a misuse into a TypeError, though that needs a nullable value from the C++ side, so it may not be worth it.

On CI, the three failures are not yours.mergeable_state says unstable, which looks alarming, so for the record, comparing the latest run per check:

your head: 25 success, 7 skipped, 3 failure
main (82c033fcb): 63 success, 4 skipped, 21 failure

Main is already red in every executorch-runtime-build variant and all py-core variants, which is where all three of your failures sit.

Two questions:

1. Any measurement? The mechanism is obviously right, but the description does not quantify the win, and a before and after (export time, or peak host memory) would also answer @Gasoonjia's question above about what dominates the lowering phase now.
2. _get_engine_info_for_node in backend.py is left alone because its caller runs after replace_execute_engine has rewritten those nodes. That reasoning looks right, but it is an ordering invariant with nothing asserting it. Worth a comment at the call site, or a cheap test that the expensive branch is not reached?