Implemented executorch's multi-optimization profile

[cehongwang]

Description

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

Fixes # (issue)

Type of change

Please delete options that are not relevant and/or add your own.

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

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

[shoumikhin]

I reviewed the multi-profile runtime, policy tests, and reference documentation. I left one runtime correctness concern and two nonblocking test/documentation suggestions inline.

[shoumikhin]

Once setOptimizationProfileAsync(profile, stream) succeeds, several later failures can return before enqueueV3() and before the existing completion event is recorded. The profile switch may therefore remain in flight while the next execute() or destroy() reconfigures or destroys the same IExecutionContext. Could you add cleanup for every post-switch error path, either by synchronizing the stream or recording completion that the next execution and teardown will wait for? Delaying the profiles.active assignment alone would not address the context-lifetime race.

[shoumikhin]

Nonblocking test suggestion: could we add (1) a two-input table where profile 0 fits input 0 but not input 1, verifying auto-selection skips it, and (2) a three-profile table where the active profile no longer fits and two lower profiles do, verifying the rescan selects the lowest matching index? The implementation handles both today, but the current tests cover only one-input tables and a rescan with one fitting alternative.

[shoumikhin]

Nit: kPrefillProfile and kDecodeProfile are example-local constants and are not defined or exported by the public API, so this copied snippet does not compile as written. Could you define them in the snippet, or use profile indices 1 and 0 with a note that the indices follow the export-time profile order?

[shoumikhin]

Re-reviewed after the fixup commit. The three earlier comments are all addressed, thanks.

The core design holds up well. I checked the stream ordering against the contract in NvInferRuntime.h (the switch, the H2D copies, and the enqueue all use the one stream, so the required happens-before comes for free with no host sync), confirmed a fresh IExecutionContext really does start on profile 0, confirmed the default kSTATIC allocation means a switch never has to allocate, and confirmed every read and write of profiles.active is under the handle mutex. The selection policy matches _TRTEngine._auto_select_profile and TRTEngine::auto_select_profile exactly.

Two things I would like resolved before merge.

1. executorch-static-build is currently failing on this PR because the two new example .cpp files are not packaged. The job log shows Cannot find source file: multi_profile_main.cpp. Same job passes on the base commit, so it is this change. Because cmake aborts the whole configure, this also breaks the pre-existing runner for anyone unpacking the release tarball.

2. The mark_inflight refactor dropped an error return the previous code had, so a faulted inference can now report success on the skip-sync path.

The rest is smaller: a bare OptimizationProfileSelection.h on every consumer's include path, the static-vs-dynamic pin inconsistency, and the -1 sentinel in the public API, which is cheap to change now and expensive after release.

One request on scope. This is about 1760 added lines mixing the runtime feature, the in-flight-event refactor, two example programs, an export script, and a rewrite of the Python dynamo example. That last one touches the Python runtime example and is a separate concern from the C++ delegate. Splitting would make each piece much easier to review.

Validation note: static review only, no GPU run. The TensorRT semantics above come from the bundled header docs, and I did not reproduce the benchmark numbers.

[shoumikhin]

These two new runners are not packaged into the release tarball, and this is already failing CI.

examples/executorch_reference_runner/BUILD has a source_files filegroup that decides what ships:

filegroup(
    name = "source_files",
    srcs = ["CMakeLists.txt", "README.md", "main.cpp"],
)

That feeds executorch_reference_runner_pkg_files -> executorch_source_package -> libtorchtrt_tar. The PR does not touch it, so the tarball gets the new CMakeLists.txt but neither multi_profile_main.cpp nor multi_profile_benchmark.cpp.

The executorch-static-build job on this commit shows exactly that:

CMake Error at CMakeLists.txt:65 (add_executable):
  Cannot find source file:
    multi_profile_main.cpp
CMake Error at CMakeLists.txt:75 (add_executable):
  Cannot find source file:
    multi_profile_benchmark.cpp
CMake Generate step failed.  Build files cannot be regenerated correctly.

The same job passes on the base commit, so this is from this change. Worth noting the blast radius is wider than the two new targets: cmake aborts the whole configure step and emits no build files, so example_executorch_runner cannot be built either.

Could you add both files to the filegroup?

    srcs = [
        "CMakeLists.txt",
        "README.md",
        "main.cpp",
        "multi_profile_benchmark.cpp",
        "multi_profile_main.cpp",
    ],

The two new backend headers were correctly added to //cpp:executorch_backend_source_files, so this is the one spot that was missed. A require_tar_entry line per new file in verify-executorch-reference-runner.sh would also catch this at packaging time rather than at configure time.

[shoumikhin]

This refactor drops an error that the previous version reported.

Before, on the skip-sync path, a failed cudaEventRecord logged, drained, and returned an error:

  } else {
    cuda_err = cudaEventRecord(engine->inflight_event, stream);
    if (cuda_err != cudaSuccess) {
      ...
      (void)cudaStreamSynchronize(stream);
      engine->inflight_pending = false;
      return Error::InvalidProgram;   // <-- gone now
    }

Now mark_inflight returns void and also discards the return code of its own fallback cudaStreamSynchronize, so with must_sync == false execute() goes on to return Error::Ok.

Why it matters: the likely reason cudaEventRecord fails right after enqueueV3 is a sticky asynchronous CUDA error from the enqueue itself. In that case we now report success for a call whose inference actually faulted, and the error resurfaces later attributed to some unrelated operator. Rare path, but it used to be reported and now is not.

Could you have it return an Error and propagate at both call sites?

Error mark_inflight(EngineHandle& engine, cudaStream_t stream) {
  const cudaError_t rec = cudaEventRecord(engine.inflight_event, stream);
  engine.inflight_pending = (rec == cudaSuccess);
  if (rec == cudaSuccess) {
    return Error::Ok;
  }
  ET_LOG(Error, "TensorRTBackend::execute: cudaEventRecord failed: %s", cudaGetErrorString(rec));
  return cudaStreamSynchronize(stream) == cudaSuccess ? Error::Ok : Error::InvalidProgram;
}

Minor related note: on the must_sync path the event recorded at the tail is synchronized away immediately after, so that record is wasted work. Harmless, just noting it since the helper is now unconditional.

[shoumikhin]

This puts a bare OptimizationProfileSelection.h on the include path of every app that links the backend.

strip_include_prefix here is the header's own directory, so the header ends up with no directory prefix at all. Every other header target in this repo strips to a directory:

strip_include_prefix = "include"   # -> torch_tensorrt/executorch/TensorRTBlobHeader.h
strip_include_prefix = "include"   # -> torch_tensorrt/executorch/TensorRTBindingNames.h

Since this target is in deps of tensorrt_executorch_backend, the virtual include dir propagates transitively, so a downstream app with its own file of that name can shadow ours (or vice versa). Bazel-only hygiene, no behavior change, so low priority, but the rest of the project deliberately avoids this.

If you do change it, two things to watch. The bare spelling is what makes one #include work in both build systems, because the header lives in src/ and the CMake build only puts cpp/include on the path (cpp/src/torch_tensorrt/executorch/CMakeLists.txt:37), so switching to strip_include_prefix = "src" alone would fix Bazel and break CMake. And the two existing include sites would need updating too:

• cpp/src/torch_tensorrt/executorch/EngineHandle.h
• tests/cpp/executorch/test_optimization_profile_selection.cpp

Simplest version is probably to move the header to cpp/include/torch_tensorrt/executorch/, include it as "torch_tensorrt/executorch/OptimizationProfileSelection.h" in both places, and strip to include like the siblings. It is already effectively public since the test depends on it.

[shoumikhin]

This tolerance treats two engines that are equally unable to honor the pin differently.

Pinning index 1:

single-profile STATIC  engine -> silently runs profile 0, returns kOk
single-profile DYNAMIC engine -> kRequestedProfileUnavailable

Both have exactly one profile, so neither has an index 1. The comment says the point is not to fail an innocent single-profile sibling when the guard was aimed at a multi-profile one, but that applies just as much to the dynamic sibling, which still fails. It also does not help two multi-profile engines with different counts (say 3 and 2, pin index 2), since neither is size() == 1.

The two existing runtimes each pick one rule and stick to it:

	out-of-range pin, single-profile engine
Python _TorchTensorRTModule.set_optimization_profile	raises ValueError, always
C++ TRTEngine::set_active_profile_with_stream	silently no-ops for all single-profile engines

Could we match one of them? If you keep the tolerance, applying it to all single-profile engines regardless of static or dynamic, plus a warning log, would at least make an ineffective pin visible rather than silent.

Narrow case in practice (one .pte mixing a static engine with a multi-profile one, and a nonzero pin), so not blocking, but the asymmetry will be hard to explain later.

[shoumikhin]

Could you document the multi-delegate hazard here?

The guard sets one thread-local that every TensorRT delegate in the method reads, so if a .pte has two engines whose profile lists differ, index 1 can mean prefill in one and decode in the other. The comment below says the delegates see one consistent request, which is true of the integer but not of its meaning.

For contrast, the other two runtimes both target something specific:

# Python: targets a module object, can be scoped to one submodule
with optimization_profile(trt_gm, 1): ...

and the C++ runtime keeps active_profile_index per engine instance.

I think the thread-local is the right call here given the ExecuTorch BackendInterface. Its official set_option channel is process-global, which would be worse under concurrency. So this is not a redesign request, just a docs one so a user with two engines is not surprised.

One idea worth a thought, not for this PR: BackendExecutionContext::get_method_name() is available inside execute(), so if prefill and decode were exported as two methods the profile could be chosen from the method name with no ambient state at all.

[shoumikhin]

Any dtype that is not Half reads through a float* here, including BFloat16.

const_data_ptr<T>() is an unchecked static_cast in the portable tensor type, so on a bf16 tensor this walks 4 bytes per element through a 2-bytes-per-element buffer and reads roughly twice past the end. No assert, just wrong numbers and an out-of-bounds read.

Latent today since export_multi_profile.py does .to(torch.float16), so the paired model never hits it. Still worth guarding, especially as IndexTensor just above carries a comment about dtype mismatch being silent corruption. A BFloat16 branch plus a log-and-skip for anything unexpected would close it.

[shoumikhin]

Two small flag edge cases in the benchmark.

--block_rounds=1 means the r != 0 guard drops every prefill sample, so prefill stats always come out n=0. The default is 3 so this only bites someone who passes 1, but the empty result is silent rather than explained.

--blocks=0, which is also what atoi returns for garbage input, leaves wall_ms == 0, so the wall-time percentage at the end computes 0.0 / 0.0 and prints nan. percentile() would also underflow size() - 1 to SIZE_MAX on an empty vector; it is currently shielded only by the empty check in summarize.

Rejecting non-positive parsed values up front would handle all three.

[shoumikhin]

These two files quote the same measurement at different rates.

Here: 1.9 GB and roughly 3.7 s/MB.
export_multi_profile.py:38: roughly 3.6 seconds per megabyte and ~2 GB.

Worth making them agree, or dropping the per-MB rate from one of the two. The nearby ~3.6 ms switch cost also reads confusingly next to 3.6 s/MB, since they are unrelated quantities that happen to share a number.

[shoumikhin]

EngineHandle.h in srcs rather than hdrs deviates from the rest of the repo, where .cpp goes in srcs and .h in hdrs without exception.

It works, and putting a deliberately private header in srcs is a legitimate Bazel idiom that matches the "not installed" note in the file. Just unexpected for a reader, so a one-line comment saying it is intentionally private would help.

[shoumikhin]

Typo: "seperately" -> "separately".