Implemented executorch's multi-optimization profile #4441
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
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,210 @@ | ||
| /* | ||
| * 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. | ||
| */ | ||
|
|
||
| /** | ||
| * @file OptimizationProfileSelection.h | ||
| * @brief 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 { | ||
|
|
||
| /// @brief The [min, max] dim envelope one optimization profile allows for one input. | ||
| struct InputProfileBounds { | ||
| nvinfer1::Dims min{}; ///< Smallest shape the profile accepts. | ||
| nvinfer1::Dims max{}; ///< Largest shape the profile accepts. | ||
| }; | ||
|
|
||
| /// @brief Everything a profile decision depends on, read from the engine once at init(). | ||
| struct ProfileTable { | ||
| /** | ||
| * @brief Bounds 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; | ||
| /// @brief The profile currently loaded into the execution context. | ||
| int32_t active = 0; | ||
|
|
||
| /// @return The engine's optimization profile count. | ||
| int32_t size() const { | ||
| return static_cast<int32_t>(bounds.size()); | ||
| } | ||
| }; | ||
|
|
||
| /// @brief 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. | ||
| }; | ||
|
|
||
| /** | ||
| * @brief The outcome of one profile decision. | ||
| * | ||
| * Its own enum rather than executorch's Error so that this header stays | ||
| * independent of executorch and can be tested separately. | ||
| * | ||
| * Two axes: whether execution continues, and which message the caller prints. | ||
| * The two failure values stay apart rather than being merged and re-derived from | ||
| * the request kind, because the empty-table guard in select_profile() returns | ||
| * kNoProfileMatchesInputs for every request kind -- so one merged value would put | ||
| * the message back at the mercy of which branches each request can reach. | ||
| */ | ||
| enum class ProfileSelection { | ||
| kOk, ///< The selected profile is usable. | ||
| /** | ||
| * @brief Succeeded, but the pin could not be honored and profile 0 was used instead. | ||
| * | ||
| * Distinct from kOk so the caller can warn that the pin did nothing here. | ||
| */ | ||
| kPinIgnoredSingleProfile, | ||
| /** | ||
| * @brief A pinned index this engine does not have and cannot substitute for. | ||
| * | ||
| * Fatal, on the same split the standard runtime uses in | ||
| * TRTEngine::set_active_profile_with_stream: an engine that had profiles to | ||
| * choose between must not quietly run on one the caller did not ask for. | ||
| * OptimizationProfileGuard cannot validate anything -- it never sees an | ||
| * engine, by design -- so execute() is the only place an ExecuTorch caller's | ||
| * bad index can be caught at all. | ||
| */ | ||
| kRequestedProfileUnavailable, | ||
| /// @brief Auto-selection ran out of profiles. | ||
| kNoProfileMatchesInputs, | ||
| }; | ||
|
|
||
| /** | ||
| * @brief Whether one input shape falls inside one profile's envelope. | ||
| * | ||
| * @param dims The runtime shape of the input. | ||
| * @param bounds The profile's [min, max] envelope for that input. | ||
| * @return true when the ranks match and every extent is in range. | ||
| */ | ||
| 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; | ||
| } | ||
|
|
||
| /** | ||
| * @brief Whether one profile accepts every input shape of an execution. | ||
| * | ||
| * A profile index or bounds row that cannot describe these inputs answers "does | ||
| * not fit" rather than indexing past the end of the table. Neither is reachable | ||
| * from the backend, which builds one row of num_inputs bounds per profile and | ||
| * only ever stores an index this function approved -- but this is installed | ||
| * public API, and on the auto path the answer also happens to be the useful one: | ||
| * a ProfileTable carrying a stale ProfileTable::active rescans from 0 instead of | ||
| * crashing. | ||
| * | ||
| * @param table The engine's profile bounds. | ||
| * @param profile Index of the profile to test. | ||
| * @param input_dims Runtime shape of each input, in binding order. | ||
| * @return true when the profile accepts all of them. | ||
| */ | ||
| inline bool profile_fits(const ProfileTable& table, int32_t profile, const std::vector<nvinfer1::Dims>& input_dims) { | ||
| if (profile < 0 || profile >= table.size()) { | ||
|
Contributor
There was a problem hiding this comment. Nothing exercises this guard or the rank check above it. I deleted each in turn and all eighteen tests still passed, including under a hardened standard library, so it is a real gap rather than a test that forgets to assert. A negative pinned index reaching Separately, the comment above promises a fit only when the ranks match, but the code only checks the min bound's rank, never the max. |
||
| return false; | ||
| } | ||
| const auto& bounds = table.bounds[static_cast<size_t>(profile)]; | ||
| if (bounds.size() < input_dims.size()) { | ||
| return false; | ||
| } | ||
| for (size_t i = 0; i < input_dims.size(); ++i) { | ||
| if (!dims_fit(input_dims[i], bounds[i])) { | ||
| return false; | ||
| } | ||
| } | ||
| return true; | ||
| } | ||
|
|
||
| /** | ||
| * @brief Resolves one thread's profile request against one engine. | ||
| * | ||
| * @param table The engine's profile bounds and currently loaded profile. | ||
| * @param request What the calling thread asked for. | ||
| * @param index The pinned profile index; read only for ProfileRequest::kPinned. | ||
| * @param input_dims Runtime shape of each input, in binding order. | ||
| * @param[out] selected The profile to run, written only when the result allows | ||
| * execution to continue (ProfileSelection::kOk or | ||
| * ProfileSelection::kPinIgnoredSingleProfile). | ||
| * @return What the caller should do, and which message it should print. | ||
| */ | ||
| inline ProfileSelection select_profile( | ||
| const ProfileTable& table, | ||
| ProfileRequest request, | ||
| int32_t index, | ||
| const std::vector<nvinfer1::Dims>& input_dims, | ||
| int32_t& selected) { | ||
| // init() rejects an engine reporting no profiles, so this is unreachable in the | ||
| // backend. Checked here so the policy is safe to call on its own rather than on | ||
| // the strength of a guard in another translation unit. | ||
| if (table.bounds.empty()) { | ||
| return ProfileSelection::kNoProfileMatchesInputs; | ||
| } | ||
|
|
||
| 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)) { | ||
|
Contributor
There was a problem hiding this comment.
The guard above checks only the outer vector, so if (table.bounds.empty()) {
return ProfileSelection::kNoProfileMatchesInputs;
}
...
if (profile_fits(table, table.active, input_dims)) {
Neither is reachable through Raising it because of the comment right above the guard:
That is the bar this header sets for itself, and it is installed public API, so "on its |
||
| 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; | ||
| } | ||
|
|
||
| // A single-profile engine has no choice to get wrong: profile 0 is the only | ||
| // thing it can run, whether or not its shapes are dynamic. So a pin aimed at a | ||
| // multi-profile sibling in the same method must not fail it. An engine with | ||
| // several profiles is different -- substituting one would be a guess -- so an | ||
| // index it lacks stays an error there. | ||
| if (index > 0 && table.size() == 1) { | ||
| selected = 0; | ||
| return ProfileSelection::kPinIgnoredSingleProfile; | ||
| } | ||
|
|
||
| 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.
On a multi-profile engine this turns a hard failure into a warning.
Before, an out-of-range index skipped the
<= 1guard, reachedsetOptimizationProfileAsync, gotfalseback, andTORCHTRT_CHECKthrew. Now itwarns and returns, so
set_active_profile(99)on a 2-profile engine goes from raisingto silently continuing on whatever profile was already loaded, which means silently
mistuned kernels.
Making the single-profile case non-silent is a genuine improvement. Would you consider
keeping the throw for an out-of-range index on a multi-profile engine, and warning only
where the engine could not have done anything differently (one profile, any nonzero
index)?
Reachability is limited:
set_optimization_profilevalidates first and raisesValueError, so only a caller driving the engine directly can hit this. I also checkedthe warning cannot spam a hot loop, since every per-call caller is gated on
num_optimization_profiles > 1and passes an index it already checked withprofile_fits.One small thing while you are here: this fixed the
.cppcomment that pointed atTorchTensorRTModule.resolve_profile_index, but the identical reference survives atTRTEngine.h:300. That name has never existed as code (git log -Sfinds it only inthose two comments); the real validator is
TorchTensorRTModule.set_optimization_profile.