Files
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

583 lines
17 KiB
C++

/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2026 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_OPTIONS_H
#define TRT_SAMPLE_OPTIONS_H
#include <array>
#include <iostream>
#include <optional>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
#include "NvInfer.h"
#if ENABLE_UNIFIED_BUILDER
#include "safeCommon.h"
#endif
namespace sample
{
// Build default params
constexpr int32_t defaultAvgTiming{8};
constexpr int32_t defaultMaxAuxStreams{-1};
constexpr int32_t defaultBuilderOptimizationLevel{-1};
constexpr int32_t defaultTilingOptimizationLevel{static_cast<int32_t>(nvinfer1::TilingOptimizationLevel::kNONE)};
constexpr int32_t defaultMaxTactics{-1};
// System default params
constexpr int32_t defaultDevice{0};
// Inference default params
constexpr int32_t defaultBatch{1};
constexpr int32_t batchNotProvided{0};
constexpr int32_t defaultStreams{1};
constexpr int32_t defaultIterations{10};
constexpr int32_t defaultOptProfileIndex{0};
constexpr float defaultWarmUp{200.F};
constexpr float defaultDuration{3.F};
constexpr float defaultSleep{};
constexpr float defaultIdle{};
constexpr float defaultPersistentCacheRatio{0};
// Reporting default params
constexpr int32_t defaultAvgRuns{10};
constexpr std::array<float, 3> defaultPercentiles{90, 95, 99};
enum class ModelFormat
{
kANY,
kONNX
};
enum class SparsityFlag
{
kDISABLE,
kENABLE,
kFORCE
};
enum class TimingCacheMode
{
kDISABLE,
kLOCAL,
kGLOBAL
};
enum class MemoryAllocationStrategy
{
kSTATIC, //< Allocate device memory based on max size across all profiles.
kPROFILE, //< Allocate device memory based on max size of the current profile.
kRUNTIME, //< Allocate device memory based on the current input shapes.
};
enum class AccuracyValidationAlgorithm
{
kL0, //< L0 accuracy algorithm
kL1, //< L1 accuracy algorithm
kL2, //< L2 accuracy algorithm
kLInf, //< LInf accuracy algorithm
kCosineSimilarity, //< Cosine similarity accuracy algorithm
};
enum class TuningSearchAlgorithm
{
kFAST, //< Fast searching algorithm: baseline + one-off variations (linear in #knobs)
kEXHAUSTIVE, //< Exhaustive: all combinations enumerated (product of knob value-counts)
kMIXED, //< Two-phase: fast scan, then exhaustive over knobs that improved performance
};
//! \brief Convert a TuningSearchAlgorithm enum to its CLI / cache-file string.
//! kFAST -> "fast"
//! kEXHAUSTIVE -> "full"
//! kMIXED -> "mixed"
inline std::string toString(TuningSearchAlgorithm algo)
{
switch (algo)
{
case TuningSearchAlgorithm::kFAST: return "fast";
case TuningSearchAlgorithm::kEXHAUSTIVE: return "full";
case TuningSearchAlgorithm::kMIXED: return "mixed";
}
return "unknown";
}
//!
//! \enum RuntimeMode
//!
//! \brief Used to dictate which TensorRT runtime library to dynamically load.
//!
enum class RuntimeMode
{
//! Maps to libnvinfer.so or nvinfer.dll
kFULL,
//! Maps to libnvinfer_dispatch.so or nvinfer_dispatch.dll
kDISPATCH,
//! Maps to libnvinfer_lean.so or nvinfer_lean.dll
kLEAN,
//! Maps to libnvinfer_safe.so or nvinfer_safe.dll
kSAFE,
};
inline std::ostream& operator<<(std::ostream& os, RuntimeMode const mode)
{
switch (mode)
{
case RuntimeMode::kFULL:
{
os << "full";
break;
}
case RuntimeMode::kDISPATCH:
{
os << "dispatch";
break;
}
case RuntimeMode::kLEAN:
{
os << "lean";
break;
}
case RuntimeMode::kSAFE:
{
os << "safe";
break;
}
}
return os;
}
using Arguments = std::unordered_multimap<std::string, std::pair<std::string, int32_t>>;
//! An IO format specification.
struct IOFormat
{
nvinfer1::TensorFormats formats{};
};
using ShapeRange = std::array<std::vector<int64_t>, nvinfer1::EnumMax<nvinfer1::OptProfileSelector>()>;
using LayerDeviceTypes = std::unordered_map<std::string, nvinfer1::DeviceType>;
using DecomposableAttentions = std::unordered_map<std::string, bool>;
using StringSet = std::unordered_set<std::string>;
class WeightStreamingBudget
{
public:
static constexpr int64_t kDISABLE{-2};
static constexpr int64_t kAUTOMATIC{-1};
int64_t bytes{kDISABLE};
double percent{static_cast<double>(100.0)};
bool isDisabled()
{
return bytes == kDISABLE && percent == kDISABLE;
}
};
class Options
{
public:
virtual ~Options() = default;
virtual void parse(Arguments& arguments) = 0;
};
class BaseModelOptions : public Options
{
public:
ModelFormat format{ModelFormat::kANY};
std::string model;
void parse(Arguments& arguments) override;
static void help(std::ostream& out);
};
class ModelOptions : public Options
{
public:
BaseModelOptions baseModel;
std::string prototxt;
std::vector<std::string> outputs;
void parse(Arguments& arguments) override;
static void help(std::ostream& out);
};
constexpr nvinfer1::TempfileControlFlags getTempfileControlDefaults()
{
using F = nvinfer1::TempfileControlFlag;
return (1U << static_cast<uint32_t>(F::kALLOW_TEMPORARY_FILES))
| (1U << static_cast<uint32_t>(F::kALLOW_IN_MEMORY_FILES));
}
class BuildOptions : public Options
{
public:
// Unit in MB.
double workspace{-1.0};
// Unit in MB.
double dlaSRAM{-1.0};
// Unit in MB.
double dlaLocalDRAM{-1.0};
// Unit in MB.
double dlaGlobalDRAM{-1.0};
// Unit in KB.
double tacticSharedMem{-1.0};
int32_t avgTiming{defaultAvgTiming};
bool tf32{true};
bool stronglyTyped{true};
bool directIO{false};
LayerDeviceTypes layerDeviceTypes;
DecomposableAttentions decomposableAttentions;
StringSet debugTensors;
bool markUnfusedTensorsAsDebugTensors{false};
StringSet debugTensorStates;
bool safe{false};
bool consistency{false};
bool dumpKernelText{false};
bool buildDLAStandalone{false};
bool allowGPUFallback{false};
bool skipInference{false};
bool save{false};
bool saveAllEngines{false}; //!< Save per-iteration engines as <engine>.iter<N> during tuning
bool load{false};
bool asyncFileReader{false};
bool refittable{false};
bool stripWeights{false};
bool versionCompatible{false};
bool pluginInstanceNorm{false};
bool enableUInt8AsymmetricQuantizationDLA{false};
bool reportCapabilityDLA{false};
bool adjustForDLA{false};
bool enablePluginOverride{false};
bool excludeLeanRuntime{false};
bool disableCompilationCache{false};
bool enableMonitorMemory{false};
bool cpuOnly{false};
int32_t builderOptimizationLevel{defaultBuilderOptimizationLevel};
int32_t maxTactics{defaultMaxTactics};
SparsityFlag sparsity{SparsityFlag::kDISABLE};
nvinfer1::ProfilingVerbosity profilingVerbosity{nvinfer1::ProfilingVerbosity::kLAYER_NAMES_ONLY};
std::string engine;
using ShapeProfile = std::unordered_map<std::string, ShapeRange>;
std::vector<ShapeProfile> optProfiles;
std::vector<IOFormat> inputFormats;
std::vector<IOFormat> outputFormats;
nvinfer1::TacticSources enabledTactics{0};
nvinfer1::TacticSources disabledTactics{0};
TimingCacheMode timingCacheMode{TimingCacheMode::kLOCAL};
std::string timingCacheFile{};
bool errorOnTimingCacheMiss{false};
// C++11 does not automatically generate hash function for enum class.
// Use int32_t to support C++11 compilers.
std::unordered_map<int32_t, bool> previewFeatures;
nvinfer1::HardwareCompatibilityLevel hardwareCompatibilityLevel{nvinfer1::HardwareCompatibilityLevel::kNONE};
nvinfer1::RuntimePlatform runtimePlatform{nvinfer1::RuntimePlatform::kSAME_AS_BUILD};
std::string tempdir{};
nvinfer1::TempfileControlFlags tempfileControls{getTempfileControlDefaults()};
RuntimeMode useRuntime{RuntimeMode::kFULL};
std::string leanDLLPath{};
std::string buildRoute{}; //!< --setBuildRoute=<route> string; passed to IBuilderConfig::setBuildRoute()
int32_t maxAuxStreams{defaultMaxAuxStreams};
bool getPlanVersionOnly{false};
bool allowWeightStreaming{false};
int32_t tilingOptimizationLevel{defaultTilingOptimizationLevel};
int64_t l2LimitForTiling{-1};
bool distributiveIndependence{false};
std::string remoteAutoTuningConfig{};
void parse(Arguments& arguments) override;
static void help(std::ostream& out);
};
class SystemOptions : public Options
{
public:
int32_t device{defaultDevice};
int32_t DLACore{-1};
bool enableStaticPlugins{true};
bool ignoreParsedPluginLibs{false};
std::vector<std::string> plugins;
std::vector<std::string> setPluginsToSerialize;
std::vector<std::string> dynamicPlugins;
#if ENABLE_UNIFIED_BUILDER
std::vector<samplesSafeCommon::SafetyPluginLibraryArgument> safetyPlugins;
#endif
void parse(Arguments& arguments) override;
static void help(std::ostream& out, bool enableStaticPlugins = true);
};
class InferenceOptions : public Options
{
public:
int32_t batch{batchNotProvided};
int32_t iterations{defaultIterations};
int32_t infStreams{defaultStreams};
int32_t optProfileIndex{defaultOptProfileIndex};
float warmup{defaultWarmUp};
float duration{defaultDuration};
float sleep{defaultSleep};
float idle{defaultIdle};
float persistentCacheRatio{defaultPersistentCacheRatio};
float atol{1e-5}; // Element-wise accuracy threshold absolute tolerance
float rtol{1e-5}; // Element-wise accuracy threshold relative tolerance
float accuracyThresholdEndToEnd{}; // End-to-end accuracy threshold should not have default value
// because it depends on the model and accuracy validation algorithm
bool overlap{true};
bool includeTransfers{false};
bool useManaged{false};
bool spin{true};
bool threads{false};
bool graph{true};
bool timeDeserialize{false};
bool timeRefit{false};
bool setOptProfile{false};
// Reference pairs for accuracy validation with multiple test cases.
// Each pair contains (inputMap, refOutputMap) indexed by pair number.
// refPairs[0] is used as the primary pair (for --loadInputs/--loadRefOutputs without --refPair).
using RefPair = std::pair<std::unordered_map<std::string, std::string>,
std::unordered_map<std::string, std::string>>; // pair of (inputMap, refOutputMap)
std::vector<RefPair> refPairs{1}; // Initialize with one empty pair
AccuracyValidationAlgorithm accuracyValidationAlgorithm{AccuracyValidationAlgorithm::kL0};
using ShapeProfile = std::unordered_map<std::string, std::vector<int64_t>>;
ShapeProfile shapes;
nvinfer1::ProfilingVerbosity nvtxVerbosity{nvinfer1::ProfilingVerbosity::kLAYER_NAMES_ONLY};
MemoryAllocationStrategy memoryAllocationStrategy{MemoryAllocationStrategy::kSTATIC};
std::unordered_map<std::string, std::string> debugTensorFileNames;
std::vector<std::string> dumpAlldebugTensorFormats;
WeightStreamingBudget weightStreamingBudget;
std::string refitOnnxModel;
void parse(Arguments& arguments) override;
static void help(std::ostream& out);
};
class ReportingOptions : public Options
{
public:
bool verbose{false};
int32_t avgs{defaultAvgRuns};
std::vector<float> percentiles{defaultPercentiles.begin(), defaultPercentiles.end()};
bool refit{false};
bool output{false};
bool dumpRawBindings{false};
bool profile{false};
bool layerInfo{false};
bool optProfileInfo{false};
std::string exportTimes;
std::string exportOutput;
std::string exportProfile;
std::string exportLayerInfo;
void parse(Arguments& arguments) override;
static void help(std::ostream& out);
};
class SafeBuilderOptions : public Options
{
public:
std::string serialized{};
std::string onnxModelFile{};
bool help{false};
bool verbose{false};
std::vector<IOFormat> inputFormats;
std::vector<IOFormat> outputFormats;
bool int8{false};
bool fp8{false};
bool int4{false};
std::vector<std::string> plugins;
bool consistency{false};
bool standard{false};
TimingCacheMode timingCacheMode{TimingCacheMode::kLOCAL};
std::string timingCacheFile{};
SparsityFlag sparsity{SparsityFlag::kDISABLE};
int32_t avgTiming{defaultAvgTiming};
void parse(Arguments& arguments) override;
static void printHelp(std::ostream& out);
};
//! \brief Options for `--tuneBuildRoutes`-driven autotuning of build routes.
//!
//! The tuning loop is implemented in samples/trtexec/trtexec.cpp by forking
//! the trtexec process for each route; each child runs `runOnceBuildAndInfer`
//! with `--setBuildRoute=<route>` injected by the parent.
class TuningOptions : public Options
{
public:
std::string tuningCacheFile{"best_config.json"};
std::string tuningExpr{}; //!< --tuneBuildRoutes
std::string tuningExprFile{}; //!< --tuneBuildRouteFile
TuningSearchAlgorithm tuningSearchAlgorithm{TuningSearchAlgorithm::kFAST};
int64_t timeout{-1}; //!< --tuningTimeOut (s); -1 = no timeout
bool helpBuildRoute{false}; //!< --helpBuildRoute (short-circuit)
std::string helpBuildRouteKnob{}; //!< --helpBuildRoute=<knob> filter
bool continueFromCache{false}; //!< --continue
bool dryRun{false}; //!< --dryRun (enumerate, don't build)
//! \brief Hidden parent->child IPC channel.
//!
//! When set, runOnceBuildAndInfer writes a small JSON to this path containing
//! gpu_time_ms, accuracy_failed, and per-tensor accuracy_loss before returning.
//! Injected into the child's argv by the tuning loop; never shown in --help.
std::string tuningResultFile{}; //!< --tuningResultFile=<path>
void parse(Arguments& arguments) override;
static void help(std::ostream& out);
};
class AllOptions : public Options
{
public:
ModelOptions model;
BuildOptions build;
SystemOptions system;
InferenceOptions inference;
ReportingOptions reporting;
TuningOptions tuning;
bool helps{false};
void parse(Arguments& arguments) override;
static void help(std::ostream& out, bool enableStaticPlugins = true);
};
class TaskInferenceOptions : public Options
{
public:
std::string engine;
int32_t device{defaultDevice};
int32_t DLACore{-1};
int32_t batch{batchNotProvided};
bool graph{true};
float persistentCacheRatio{defaultPersistentCacheRatio};
void parse(Arguments& arguments) override;
static void help(std::ostream& out);
};
Arguments argsToArgumentsMap(int32_t argc, char* argv[]);
bool parseHelp(Arguments& arguments);
void helpHelp(std::ostream& out);
// Functions to print options
std::ostream& operator<<(std::ostream& os, BaseModelOptions const& options);
std::ostream& operator<<(std::ostream& os, IOFormat const& format);
std::ostream& operator<<(std::ostream& os, ShapeRange const& dims);
std::ostream& operator<<(std::ostream& os, ModelOptions const& options);
std::ostream& operator<<(std::ostream& os, BuildOptions const& options);
std::ostream& operator<<(std::ostream& os, SystemOptions const& options);
std::ostream& operator<<(std::ostream& os, InferenceOptions const& options);
std::ostream& operator<<(std::ostream& os, ReportingOptions const& options);
std::ostream& operator<<(std::ostream& os, AllOptions const& options);
std::ostream& operator<<(std::ostream& os, SafeBuilderOptions const& options);
std::ostream& operator<<(std::ostream& os, nvinfer1::DataType dtype);
std::ostream& operator<<(std::ostream& os, nvinfer1::DeviceType devType);
inline std::ostream& operator<<(std::ostream& os, nvinfer1::Dims const& dims)
{
for (int32_t i = 0; i < dims.nbDims; ++i)
{
os << (i ? "x" : "") << dims.d[i];
}
return os;
}
inline std::ostream& operator<<(std::ostream& os, const nvinfer1::WeightsRole role)
{
switch (role)
{
case nvinfer1::WeightsRole::kKERNEL:
{
os << "Kernel";
break;
}
case nvinfer1::WeightsRole::kBIAS:
{
os << "Bias";
break;
}
case nvinfer1::WeightsRole::kSHIFT:
{
os << "Shift";
break;
}
case nvinfer1::WeightsRole::kSCALE:
{
os << "Scale";
break;
}
case nvinfer1::WeightsRole::kCONSTANT:
{
os << "Constant";
break;
}
case nvinfer1::WeightsRole::kANY:
{
os << "Any";
break;
}
}
return os;
}
inline std::ostream& operator<<(std::ostream& os, std::vector<int64_t> const& vec)
{
for (int32_t i = 0, e = static_cast<int32_t>(vec.size()); i < e; ++i)
{
os << (i ? "x" : "") << vec[i];
}
return os;
}
} // namespace sample
#endif // TRT_SAMPLES_OPTIONS_H