138 lines
5.6 KiB
C++
138 lines
5.6 KiB
C++
#pragma once
|
|
#include "inference/backend_config.hpp"
|
|
|
|
#include <onnxruntime/onnxruntime_cxx_api.h>
|
|
#include <filesystem>
|
|
#include <iostream>
|
|
#include <string>
|
|
#include <unordered_map>
|
|
#include <vector>
|
|
|
|
// Private to the ORT backend (backends/ort_backend.cpp). Detects the best
|
|
// available ORT execution provider and applies it to a SessionOptions.
|
|
// Priority order: TensorRT EP > CUDA > ROCm > CPU.
|
|
//
|
|
// Note: "TensorRT" here is ORT's TensorRT *execution provider*, distinct from
|
|
// the raw-TensorRT backend (backends/trt_backend.cpp). This header never reaches
|
|
// core translation units.
|
|
//
|
|
// Detection is conservative: GetAvailableProviders() confirms ORT was compiled
|
|
// with the provider, then AppendExecutionProvider_* is attempted inside a
|
|
// try/catch so a missing runtime library degrades gracefully to the next tier.
|
|
|
|
enum class OrtProvider { CPU, CUDA, ROCm, TensorRT };
|
|
|
|
inline OrtProvider detect_ort_provider() {
|
|
auto available = Ort::GetAvailableProviders();
|
|
for (const auto& p : available) {
|
|
#ifdef SAE_ORT_WITH_TRT_EP
|
|
if (p == "TensorrtExecutionProvider") return OrtProvider::TensorRT;
|
|
if (p == "CUDAExecutionProvider") return OrtProvider::CUDA;
|
|
#endif
|
|
if (p == "ROCMExecutionProvider") return OrtProvider::ROCm;
|
|
}
|
|
return OrtProvider::CPU;
|
|
}
|
|
|
|
inline const char* provider_name(OrtProvider p) {
|
|
switch (p) {
|
|
case OrtProvider::TensorRT: return "TensorRT-EP";
|
|
case OrtProvider::CUDA: return "CUDA";
|
|
case OrtProvider::ROCm: return "ROCm";
|
|
default: return "CPU";
|
|
}
|
|
}
|
|
|
|
// If trt_cfg.ort_cache_dir is set, configure ORT to write/read a pre-optimized
|
|
// .ort model for model_path. Must be called before AppendExecutionProvider_*.
|
|
inline void apply_ort_model_cache(Ort::SessionOptions& opts,
|
|
const std::string& model_path,
|
|
const BackendConfig& trt_cfg) {
|
|
if (trt_cfg.ort_cache_dir.empty()) return;
|
|
std::filesystem::create_directories(trt_cfg.ort_cache_dir);
|
|
const std::string stem =
|
|
std::filesystem::path(model_path).stem().string();
|
|
const std::string cache_path =
|
|
trt_cfg.ort_cache_dir + "/" + stem + ".ort";
|
|
opts.SetOptimizedModelFilePath(cache_path.c_str());
|
|
}
|
|
|
|
// Apply the given provider to opts. Falls back to CPU on failure and returns the
|
|
// provider that was actually applied.
|
|
inline OrtProvider apply_ort_provider(Ort::SessionOptions& opts,
|
|
OrtProvider provider,
|
|
const char* label,
|
|
const BackendConfig& trt_cfg = {}) {
|
|
#ifdef SAE_ORT_WITH_TRT_EP
|
|
if (provider == OrtProvider::TensorRT) {
|
|
try {
|
|
std::filesystem::create_directories(trt_cfg.cache_dir);
|
|
|
|
std::unordered_map<std::string, std::string> kv = {
|
|
{"device_id", "0"},
|
|
{"trt_max_workspace_size", "2147483648"},
|
|
{"trt_fp16_enable", trt_cfg.fp16 ? "1" : "0"},
|
|
{"trt_int8_enable", trt_cfg.int8 ? "1" : "0"},
|
|
{"trt_engine_cache_enable", "1"},
|
|
{"trt_engine_cache_path", trt_cfg.cache_dir},
|
|
};
|
|
|
|
if (!trt_cfg.input_name.empty() && !trt_cfg.profile_min.empty()) {
|
|
kv["trt_profile_min_shapes"] =
|
|
trt_cfg.input_name + ":" + trt_cfg.profile_min;
|
|
kv["trt_profile_opt_shapes"] =
|
|
trt_cfg.input_name + ":" + trt_cfg.profile_opt;
|
|
kv["trt_profile_max_shapes"] =
|
|
trt_cfg.input_name + ":" + trt_cfg.profile_max;
|
|
}
|
|
|
|
Ort::TensorRTProviderOptions trt_v2;
|
|
trt_v2.Update(kv);
|
|
opts.AppendExecutionProvider_TensorRT_V2(*trt_v2);
|
|
|
|
std::cerr << "[" << label << "] TensorRT EP"
|
|
<< (trt_cfg.fp16 ? " FP16" : "")
|
|
<< (trt_cfg.int8 ? " INT8" : "")
|
|
<< " cache=" << trt_cfg.cache_dir
|
|
<< (trt_cfg.profile_min.empty() ? "" :
|
|
" profile=" + trt_cfg.profile_min
|
|
+ "/" + trt_cfg.profile_opt
|
|
+ "/" + trt_cfg.profile_max)
|
|
<< "\n";
|
|
return OrtProvider::TensorRT;
|
|
} catch (const Ort::Exception& e) {
|
|
std::cerr << "[" << label << "] TensorRT EP unavailable ("
|
|
<< e.what() << "), trying CUDA\n";
|
|
provider = OrtProvider::CUDA;
|
|
}
|
|
}
|
|
#endif // SAE_ORT_WITH_TRT_EP
|
|
if (provider == OrtProvider::CUDA) {
|
|
try {
|
|
OrtCUDAProviderOptions cuda{};
|
|
cuda.device_id = 0;
|
|
opts.AppendExecutionProvider_CUDA(cuda);
|
|
std::cerr << "[" << label << "] CUDA provider\n";
|
|
return OrtProvider::CUDA;
|
|
} catch (const Ort::Exception& e) {
|
|
std::cerr << "[" << label << "] CUDA unavailable ("
|
|
<< e.what() << "), trying ROCm\n";
|
|
provider = OrtProvider::ROCm;
|
|
}
|
|
}
|
|
if (provider == OrtProvider::ROCm) {
|
|
try {
|
|
OrtROCMProviderOptions rocm{};
|
|
rocm.device_id = 0;
|
|
opts.AppendExecutionProvider_ROCM(rocm);
|
|
std::cerr << "[" << label << "] ROCm provider\n";
|
|
return OrtProvider::ROCm;
|
|
} catch (const Ort::Exception& e) {
|
|
std::cerr << "[" << label << "] ROCm unavailable ("
|
|
<< e.what() << "), falling back to CPU\n";
|
|
}
|
|
}
|
|
std::cerr << "[" << label << "] CPU provider\n";
|
|
return OrtProvider::CPU;
|
|
}
|