Implemented executorch's multi-optimization profile #4441
+1,760
−147
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,59 @@ | ||
| /* | ||
| * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. | ||
| * | ||
| * This source code is licensed under the BSD-style license found in the | ||
| * LICENSE file in the root directory of this source tree. | ||
| * | ||
| * Private state of a TensorRT ExecuTorch delegate. | ||
| * | ||
| * This header is deliberately not installed. EngineHandle grows fields as the | ||
| * backend gains features, so keeping it out of the public API means a new | ||
| * header can never disagree about its layout with an already-built backend | ||
| * archive. | ||
| */ | ||
| #pragma once | ||
|
|
||
| #include "OptimizationProfileSelection.h" | ||
| #include "torch_tensorrt/executorch/TensorRTBackend.h" | ||
|
|
||
| #include <NvInfer.h> | ||
| #include <cuda_runtime.h> | ||
|
|
||
| #include <cstdint> | ||
| #include <mutex> | ||
| #include <string> | ||
| #include <vector> | ||
|
|
||
| namespace torch_tensorrt { | ||
| namespace executorch_backend { | ||
|
|
||
| struct EngineHandle { | ||
| TRTLogger logger; | ||
| TRTUniquePtr<nvinfer1::IRuntime> runtime; | ||
| TRTUniquePtr<nvinfer1::ICudaEngine> engine; | ||
| TRTUniquePtr<nvinfer1::IExecutionContext> exec_ctx; | ||
| std::vector<std::string> input_binding_names; | ||
| std::vector<std::string> output_binding_names; | ||
| ProfileTable profiles; | ||
| std::vector<void*> cached_input_ptrs; | ||
| std::vector<size_t> cached_input_sizes; | ||
| std::vector<void*> cached_output_ptrs; | ||
| std::vector<size_t> cached_output_sizes; | ||
| size_t num_inputs = 0; | ||
| size_t num_outputs = 0; | ||
| int device_id = 0; | ||
| bool unified_memory = false; | ||
| std::mutex mu; | ||
| // Makes the skip-sync fast path safe to reuse: TensorRT forbids reconfiguring or | ||
| // destroying an execution context while one of its enqueues is in flight, so when | ||
| // execute() returns without an end sync it records this event; the next execute() | ||
| // and the destructor wait on it before touching exec_ctx. One event/flag pair | ||
| // suffices because a handle runs on a single thread at a time. | ||
| cudaEvent_t inflight_event = nullptr; | ||
| bool inflight_pending = false; | ||
|
|
||
| ~EngineHandle(); | ||
| }; | ||
|
|
||
| } // namespace executorch_backend | ||
| } // namespace torch_tensorrt |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change | ||||||
|---|---|---|---|---|---|---|---|---|
| @@ -0,0 +1,135 @@ | ||||||||
| /* | ||||||||
| * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. | ||||||||
| * | ||||||||
| * This source code is licensed under the BSD-style license found in the | ||||||||
| * LICENSE file in the root directory of this source tree. | ||||||||
| * | ||||||||
| * Which TensorRT optimization profile an execution runs under. | ||||||||
| * | ||||||||
| * Kept free of ExecuTorch, CUDA, and the engine itself so the policy can be | ||||||||
| * exercised without a GPU; reporting the outcome is left to the caller. | ||||||||
| */ | ||||||||
| #pragma once | ||||||||
|
|
||||||||
| #include <NvInfer.h> | ||||||||
|
|
||||||||
| #include <cstdint> | ||||||||
| #include <vector> | ||||||||
|
|
||||||||
| namespace torch_tensorrt { | ||||||||
| namespace executorch_backend { | ||||||||
|
|
||||||||
| // The [min, max] dim envelope one optimization profile allows for one input. | ||||||||
| struct InputProfileBounds { | ||||||||
| nvinfer1::Dims min{}; | ||||||||
| nvinfer1::Dims max{}; | ||||||||
| }; | ||||||||
|
|
||||||||
| // Everything a profile decision depends on, read from the engine once at init(). | ||||||||
| struct ProfileTable { | ||||||||
| // Indexed [profile][input]. The outer size is the engine's optimization | ||||||||
| // profile count, which is at least 1; a single-profile engine keeps exactly | ||||||||
| // one row and never switches. | ||||||||
| std::vector<std::vector<InputProfileBounds>> bounds; | ||||||||
| // True when every input dim is pinned to one extent in every profile. Such an | ||||||||
|
Contributor
There was a problem hiding this comment. This comment is not true for a multi-profile engine whose profiles are each individually static. Take profile 0 pinned at Harmless today because the only reader also requires |
||||||||
| // engine accepts one shape only, so a request for a profile it does not have | ||||||||
| // changes nothing it can do. | ||||||||
| bool all_inputs_static = true; | ||||||||
| // The profile currently loaded into the execution context. | ||||||||
| int32_t active = 0; | ||||||||
|
|
||||||||
| int32_t size() const { | ||||||||
| return static_cast<int32_t>(bounds.size()); | ||||||||
| } | ||||||||
| }; | ||||||||
|
|
||||||||
| // What the calling thread asked for, as resolved from OptimizationProfileGuard. | ||||||||
| enum class ProfileRequest { | ||||||||
| kUnset, // no guard in scope | ||||||||
| kPinned, // an exact index | ||||||||
| kAuto, // choose from the input shapes | ||||||||
| }; | ||||||||
|
|
||||||||
| enum class ProfileSelection { | ||||||||
| // Created this enum to decouple the profile header from executorch so that we can test it seperately | ||||||||
|
Contributor
There was a problem hiding this comment. Typo: "seperately" -> "separately". |
||||||||
| kOk, | ||||||||
| // A pinned index this engine does not have, and it is not static enough for | ||||||||
| // that to be harmless. | ||||||||
| kRequestedProfileUnavailable, | ||||||||
| // Auto-selection ran out of profiles. | ||||||||
| kNoProfileMatchesInputs, | ||||||||
| }; | ||||||||
|
|
||||||||
| inline bool dims_fit(const nvinfer1::Dims& dims, const InputProfileBounds& bounds) { | ||||||||
| if (dims.nbDims != bounds.min.nbDims) { | ||||||||
| return false; | ||||||||
| } | ||||||||
| for (int d = 0; d < dims.nbDims; ++d) { | ||||||||
| if (dims.d[d] < bounds.min.d[d] || dims.d[d] > bounds.max.d[d]) { | ||||||||
| return false; | ||||||||
| } | ||||||||
| } | ||||||||
| return true; | ||||||||
| } | ||||||||
|
|
||||||||
| inline bool profile_fits(const ProfileTable& table, int32_t profile, const std::vector<nvinfer1::Dims>& input_dims) { | ||||||||
| const auto& bounds = table.bounds[static_cast<size_t>(profile)]; | ||||||||
| for (size_t i = 0; i < input_dims.size(); ++i) { | ||||||||
| if (!dims_fit(input_dims[i], bounds[i])) { | ||||||||
| return false; | ||||||||
| } | ||||||||
| } | ||||||||
| return true; | ||||||||
| } | ||||||||
|
|
||||||||
| // Resolves one thread's profile request against one engine. `index` is read | ||||||||
| // only for ProfileRequest::kPinned. | ||||||||
| inline ProfileSelection select_profile( | ||||||||
| const ProfileTable& table, | ||||||||
| ProfileRequest request, | ||||||||
| int32_t index, | ||||||||
| const std::vector<nvinfer1::Dims>& input_dims, | ||||||||
| int32_t& selected) { | ||||||||
| if (request == ProfileRequest::kUnset) { | ||||||||
| selected = 0; | ||||||||
| return ProfileSelection::kOk; | ||||||||
| } | ||||||||
|
|
||||||||
| if (request == ProfileRequest::kAuto) { | ||||||||
| // Sticky first-fit: keep the loaded profile while it still fits, so shapes | ||||||||
| // that alternate between two equally valid profiles don't thrash the | ||||||||
| // context. Only rescan from 0 once it stops fitting. Overlapping profiles | ||||||||
| // therefore resolve by history, not by lowest index; pin explicitly when | ||||||||
| // that matters. | ||||||||
| if (profile_fits(table, table.active, input_dims)) { | ||||||||
| selected = table.active; | ||||||||
| return ProfileSelection::kOk; | ||||||||
| } | ||||||||
| for (int32_t p = 0; p < table.size(); ++p) { | ||||||||
| if (profile_fits(table, p, input_dims)) { | ||||||||
| selected = p; | ||||||||
| return ProfileSelection::kOk; | ||||||||
| } | ||||||||
| } | ||||||||
| return ProfileSelection::kNoProfileMatchesInputs; | ||||||||
| } | ||||||||
|
|
||||||||
| if (index >= 0 && index < table.size()) { | ||||||||
| selected = index; | ||||||||
| return ProfileSelection::kOk; | ||||||||
| } | ||||||||
|
|
||||||||
| // An engine that accepts exactly one shape has nothing to switch, so a request | ||||||||
| // aimed at its multi-profile siblings in the same method is satisfied by | ||||||||
| // profile 0 rather than failing the whole execution. A dynamic engine that | ||||||||
| // lacks the index is a real mismatch and is reported. | ||||||||
| if (index > 0 && table.size() == 1 && table.all_inputs_static) { | ||||||||
|
Contributor
There was a problem hiding this comment. This tolerance treats two engines that are equally unable to honor the pin differently. Pinning index 1: 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 The two existing runtimes each pick one rule and stick to it:
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 |
||||||||
| selected = 0; | ||||||||
| return ProfileSelection::kOk; | ||||||||
| } | ||||||||
|
|
||||||||
| return ProfileSelection::kRequestedProfileUnavailable; | ||||||||
| } | ||||||||
|
|
||||||||
| } // namespace executorch_backend | ||||||||
| } // namespace torch_tensorrt | ||||||||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This puts a bare
OptimizationProfileSelection.hon the include path of every app that links the backend.strip_include_prefixhere 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:Since this target is in
depsoftensorrt_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
#includework in both build systems, because the header lives insrc/and the CMake build only putscpp/includeon the path (cpp/src/torch_tensorrt/executorch/CMakeLists.txt:37), so switching tostrip_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.htests/cpp/executorch/test_optimization_profile_selection.cppSimplest 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 toincludelike the siblings. It is already effectively public since the test depends on it.