Files
nvidia--tensorrt/samples/common/sampleTuning.h
T
wehub-resource-sync c8a779b1bb
Docker Image CI / build-ubuntu2004 (push) Waiting to run
chore: import upstream snapshot with attribution
2026-07-13 13:36:55 +08:00

357 lines
16 KiB
C++

/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef TRT_SAMPLE_TUNING_H
#define TRT_SAMPLE_TUNING_H
#include "bigInt.h"
#include "logger.h"
#include "sampleOptions.h"
#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <unordered_map>
#include <vector>
namespace sample
{
//! \brief Split a pipe-delimited string into trimmed tokens.
//!
//! Given a string like "val1|val2|val3", returns {"val1", "val2", "val3"}.
//! Each token is trimmed of leading/trailing whitespace. Empty tokens
//! (after trimming) are skipped.
//!
//! \param[in] str The pipe-delimited string to split.
//! \return Vector of trimmed token strings.
std::vector<std::string> splitPipeDelimited(std::string const& str);
// ============================================================================
// Build Route Expression Parser and Expander
// ============================================================================
//
// These classes implement parsing and expansion of build route expressions
// for the --tuneBuildRoutes feature in trtexec. The expression syntax supports:
// - Variable values: "-knob=[val1|val2|val3]"
// - Fixed values: "-knob=fixed_value"
//
// Two expansion modes are supported:
// - Full (exhaustive): All combinations enumerated (exponential in number of knobs)
// - Fast (one-off): Baseline (all defaults) + one-knob-different variations (linear)
//
// Example:
// Expression: "-opt1=[on|off] -opt2=[0|1|2]"
// Full expansion (6 combinations):
// [0]: -opt1=on -opt2=0
// [1]: -opt1=on -opt2=1
// [2]: -opt1=on -opt2=2
// [3]: -opt1=off -opt2=0
// [4]: -opt1=off -opt2=1
// [5]: -opt1=off -opt2=2
// Fast expansion (with defaults opt1=off, opt2=2):
// [0]: -opt1=off -opt2=2 (baseline)
// [1]: -opt1=off -opt2=0 (opt2 changed)
// [2]: -opt1=off -opt2=1 (opt2 changed)
// [3]: -opt1=on -opt2=2 (opt1 changed)
// ============================================================================
//! \struct BuildRouteKnobDef
//! \brief Represents a single knob/option definition from the JSON configuration.
struct BuildRouteKnobDef
{
std::string mOption; //!< Option name, e.g., "-kgen:tiling"
std::string mAllowedValues; //!< Allowed values string, e.g., "-kgen:tiling=[0|1|2]"
std::string mDefaultValue; //!< Default value
std::string mHelp; //!< Help text
std::vector<std::string> mValues; //!< Parsed allowed values (empty if unbounded)
bool mIsBounded{false}; //!< True if values form a closed set
};
//! \struct BuildRouteParsedExpr
//! \brief Represents a parsed expression from user input.
struct BuildRouteParsedExpr
{
std::string mKnobName; //!< Knob name, e.g., "-kgen:tiling"
std::vector<std::string> mValues; //!< Values to expand, e.g., ["0", "1", "2"]
bool mIsFixed{false}; //!< True if this is a fixed value (no expansion)
};
//! \class BuildRouteKnobDatabase
//! \brief Database of knob definitions loaded from getAllBuildRoutes() JSON output.
//!
//! This class parses the JSON output from IBuilderConfig::getAllBuildRoutes()
//! and provides lookup and validation functions for knob definitions.
class BuildRouteKnobDatabase
{
public:
//! \brief Default constructor.
BuildRouteKnobDatabase() = default;
//! \brief Destructor.
~BuildRouteKnobDatabase() = default;
// Non-copyable, movable
BuildRouteKnobDatabase(BuildRouteKnobDatabase const&) = delete;
BuildRouteKnobDatabase& operator=(BuildRouteKnobDatabase const&) = delete;
BuildRouteKnobDatabase(BuildRouteKnobDatabase&&) = default;
BuildRouteKnobDatabase& operator=(BuildRouteKnobDatabase&&) = default;
//! \brief Load knob definitions from a JSON string.
//! \param[in] jsonStr JSON string from getAllBuildRoutes().
//! \return True if loading succeeded, false otherwise.
bool loadFromJsonString(std::string const& jsonStr);
//! \brief Check if a knob exists in the database.
//! \param[in] knobName The knob name to check.
//! \return True if the knob exists.
bool hasKnob(std::string const& knobName) const;
//! \brief Get the knob definition for a given knob name.
//! \param[in] knobName The knob name to look up.
//! \return Pointer to the KnobDef, or nullptr if not found.
BuildRouteKnobDef const* getKnob(std::string const& knobName) const;
//! \brief Validate that all values are allowed for a knob.
//! \param[in] knobName The knob name.
//! \param[in] values The values to validate.
//! \return True if all values are valid.
bool validateValues(std::string const& knobName, std::vector<std::string> const& values) const;
//! \brief Check if a knob has bounded (finite) values.
//! \param[in] knobName The knob name.
//! \return True if the knob has bounded values.
bool isBounded(std::string const& knobName) const;
//! \brief Get the default value for a knob.
//! \param[in] knobName The knob name.
//! \return The default value string, or empty string if not found.
std::string getDefaultValue(std::string const& knobName) const;
//! \brief Get the tuner version string extracted from the JSON.
//! \return Tuner version (e.g. "2.17.83"), or "unknown" if not available.
std::string const& getTunerVersion() const
{
return mTunerVersion;
}
//! \brief Build the default build route string from all knob defaults.
//! Returns a space-separated string of "knob=default_value" pairs,
//! preserving the insertion order from the original JSON.
//! Example: "-conv_use_long_w=on -reshape_ppg=on -transpose_ppg=on ..."
std::string buildDefaultPath() const;
private:
//! \brief Parse allowed values from an allowed_values string.
//! \param[in] allowedStr The allowed_values string from JSON.
//! \return Vector of parsed values (empty if unbounded).
static std::vector<std::string> parseAllowedValues(std::string const& allowedStr);
std::unordered_map<std::string, BuildRouteKnobDef> mKnobs; //!< Map of knob names to definitions.
std::vector<std::string> mKnobOrder; //!< Insertion order of knob names.
std::string mTunerVersion{"unknown"}; //!< Tuner version from JSON.
};
//! \class BuildRouteExprParser
//! \brief Parser for build route expression strings.
//!
//! Parses expressions like "-opt1=[a|b] -opt2=fixed -opt3=[0|1|2]"
//! into a list of BuildRouteParsedExpr structures for expansion.
class BuildRouteExprParser
{
public:
//! \brief Construct a parser with a reference to a knob database.
//! \param[in] db Reference to the BuildRouteKnobDatabase for validation.
explicit BuildRouteExprParser(BuildRouteKnobDatabase const& db);
//! \brief Destructor.
~BuildRouteExprParser() = default;
// Non-copyable, non-movable (holds reference)
BuildRouteExprParser(BuildRouteExprParser const&) = delete;
BuildRouteExprParser& operator=(BuildRouteExprParser const&) = delete;
BuildRouteExprParser(BuildRouteExprParser&&) = delete;
BuildRouteExprParser& operator=(BuildRouteExprParser&&) = delete;
//! \brief Parse an input string into a list of expressions.
//! \param[in] input The input string containing expressions.
//! \return Optional vector of BuildRouteParsedExpr, or nullopt on error.
std::optional<std::vector<BuildRouteParsedExpr>> parse(std::string const& input) const;
//! \brief Get the last error message.
//! \return The error message string.
std::string const& getError() const noexcept;
private:
//! \brief Parse a single expression.
//! \param[in] expr The expression string.
//! \return Optional BuildRouteParsedExpr, or nullopt on error.
std::optional<BuildRouteParsedExpr> parseExpr(std::string const& expr) const;
//! \brief Tokenize input into individual expressions.
//! \param[in] input The input string.
//! \return Vector of token strings.
std::vector<std::string> tokenize(std::string const& input) const;
BuildRouteKnobDatabase const& mDb; //!< Reference to the knob database.
mutable std::string mError; //!< Last error message.
};
//! \struct TuningContext
//! \brief Context for build route tuning when --tuneBuildRoutes is specified.
//!
//! Uses lazy expansion: build route strings are computed on-demand from
//! the parsed expressions and an index, without storing all strings in memory.
//! This supports arbitrarily large expansion spaces (full mode with BigInt count).
struct TuningContext
{
//! Parsed expressions from --tuneBuildRoutes.
std::vector<BuildRouteParsedExpr> parsedExprs;
//! Default values for each knob (from knob database, needed for fast mode).
std::vector<std::string> defaultValues;
//! Total number of configurations.
BigInt totalCount;
//! Search algorithm (determines expansion strategy).
TuningSearchAlgorithm searchAlgorithm{TuningSearchAlgorithm::kFAST};
//! Tuner version string from the knob database JSON (e.g. "2.17.83").
std::string tunerVersion{"unknown"};
//! Default build route string (all knobs at their default values).
std::string defaultBuildRoute;
//! \brief Compute total configuration count based on searchAlgorithm.
//! - Exhaustive: product of all value list sizes.
//! - Fast: 1 (baseline) + sum of non-default values per variable knob.
BigInt count() const;
//! \brief Get the build route string for a given index (lazy).
//! Throws std::out_of_range if index is out of range.
//! - Exhaustive: reverse mixed-radix decomposition.
//! - Fast: baseline at index 0, one-off variations at index 1..N.
std::string getPathAtIndex(BigInt const& index) const;
};
// ============================================================================
// Mixed Search Algorithm Support
// ============================================================================
//
// Mixed search runs in two phases:
// Phase 1 (fast scan): baseline + one-off variations (same as --tuningSearch=fast).
// After each one-off iteration, if it passed accuracy, didn't crash, and was
// faster than the baseline, the knob is recorded as a "positive" knob.
// Phase 2 (full combinatorial): if >1 positive knobs found, run 2^n combinations
// of the top-N positive knobs. Other knobs stay at baseline defaults.
//
// Total iterations: fast_count + (2^n if n > 1, else 0).
//! Maximum number of positive knobs to carry into phase 2 of mixed search.
constexpr int32_t kMixedSearchMaxPositiveKnobs = 10;
//! \struct MixedSearchKnobResult
//! \brief Records a single knob variation that improved performance in phase 1.
//! Used to select which knobs enter phase 2 of mixed search.
struct MixedSearchKnobResult
{
int32_t knobIndex; //!< Index into TuningContext::parsedExprs
std::string value; //!< The non-default value that was tested
double gpuTimeMs; //!< GPU time achieved (lower is better)
};
//! \brief Build a TuningContext for phase 2 of mixed search.
//!
//! Constructs a new exhaustive-mode TuningContext where only the positive knobs
//! are variable (with their full value lists from phase 1) and all other knobs
//! are fixed to their baseline default values.
//!
//! \param phase1Context The original TuningContext from phase 1 (fast mode).
//! \param positiveKnobs The knobs that showed improvement in phase 1.
//! \return A new TuningContext configured for exhaustive search over positive knobs.
TuningContext buildMixedPhase2Context(
TuningContext const& phase1Context, std::vector<MixedSearchKnobResult> const& positiveKnobs);
//! \brief Identify which knob was varied at a given fast-mode iteration index.
//!
//! In fast mode, index 0 is the baseline (all defaults). Indices 1..N are one-off
//! variations, iterating knobs right-to-left and skipping default values.
//! This function reverses that mapping: given a one-off index (must be > 0),
//! returns the (knobIndex, value) pair, or nullopt if the index is out of range.
//!
//! Example: with 2 binary knobs (defaults "on", "on"), values [on|off] each:
//! index 1 → knob 1 changed to "off" → returns {1, "off"}
//! index 2 → knob 0 changed to "off" → returns {0, "off"}
//!
//! \param ctx The TuningContext (must be fast or mixed mode).
//! \param index The one-off iteration index (must be > 0).
//! \return (knobIndex, value) pair, or nullopt if index is invalid.
std::optional<std::pair<int32_t, std::string>> identifyVariedKnob(TuningContext const& ctx, BigInt const& index);
//! \brief Collect a positive knob from an iteration result (used in mixed search).
//!
//! If the iteration did not crash, achieved positive GPU time, and was faster than the
//! baseline, identifies which knob was varied (using identifyVariedKnob) and appends
//! it to the positiveKnobs vector. Called from both the live tuning loop (phase 1) and
//! from the --continue cache reconstruction path.
//!
//! Example: if iteration 3 changed knob 1 to "off" and achieved 1.2ms vs baseline 1.5ms,
//! this appends {knobIndex=1, value="off", gpuTimeMs=1.2} to positiveKnobs.
//!
//! \param[in] crashed Whether the iteration crashed.
//! \param[in] gpuTimeMs GPU time achieved by this iteration.
//! \param[in] baselineGpuTimeMs Baseline GPU time (index 0).
//! \param[in] index The iteration index (must be > 0 for one-off variations).
//! \param[in] ctx The TuningContext (fast/mixed mode).
//! \param[in,out] positiveKnobs Vector to append the positive knob result to.
void collectPositiveKnobFromResult(bool crashed, double gpuTimeMs, double baselineGpuTimeMs, BigInt const& index,
TuningContext const& ctx, std::vector<MixedSearchKnobResult>& positiveKnobs);
// Exit codes for child process to distinguish failure modes (used when --tuneBuildRoutes expands to multiple configs)
constexpr int32_t kChildExitSuccess = 0;
constexpr int32_t kChildExitFailure = 1; // runOnceBuildAndInfer returned failure
constexpr int32_t kChildExitStdException = 100; // Caught std::exception
constexpr int32_t kChildExitUnknownException = 101; // Caught unknown exception
//! \struct TuningIterationResult
//! \brief Captures results from a single tuning iteration for the tuning cache file.
struct TuningIterationResult
{
double gpuTimeMs{0.0}; //!< Mean GPU compute time in milliseconds
bool accuracyFailed{false}; //!< True if accuracy exceeded the threshold
std::unordered_map<std::string, double> accuracyLossValues; //!< Per-tensor accuracy values
};
//! \brief True if `arg` is a tuning-only parent flag that should be stripped from
//! the child argv. Matches both bare flag and flag=value forms. Tuning-only means
//! the parent loop interprets it; the child runs a plain single-route build.
[[nodiscard]] bool isTuningOnlyArg(char const* arg);
//! \brief Build a child argv for one tuning iteration: copies argv with tuning-only
//! flags removed and appends `--setBuildRoute=<route>`, `--saveEngine=<enginePath>`,
//! `--tuningResultFile=<resultJsonPath>`. String storage is owned by `storage` so
//! the returned `char*` pointers stay valid until the caller's execvp() completes.
//! The returned vector is nullptr-terminated, ready for execvp().
[[nodiscard]] std::vector<char*> buildTuningChildArgv(int32_t argc, char** argv, std::string const& route,
std::string const& enginePath, std::string const& resultJsonPath, std::vector<std::string>& storage);
} // namespace sample
#endif // TRT_SAMPLE_TUNING_H