mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
add caching support for dynamic input models (#4702)
This commit is contained in:
parent
fff0b41fcb
commit
75ad7be336
2 changed files with 77 additions and 32 deletions
|
|
@ -16,6 +16,11 @@
|
|||
#include "cuda_runtime_api.h"
|
||||
#include "gsl/gsl"
|
||||
#include <experimental/filesystem>
|
||||
#include <unordered_map>
|
||||
#include <utility>
|
||||
#include <limits>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
|
||||
#define CUDA_RETURN_IF_ERROR(expr) \
|
||||
ORT_RETURN_IF_ERROR(CUDA_CALL(expr) \
|
||||
|
|
@ -40,6 +45,14 @@ std::string GetEnginePath(const ::std::string& root, const std::string& name) {
|
|||
return path.string();
|
||||
}
|
||||
}
|
||||
|
||||
std::string GetVecHash(const ::std::vector<int> & vec) {
|
||||
std::size_t ret = 0;
|
||||
for (auto& i : vec) {
|
||||
ret ^= std::hash<uint32_t>()(i);
|
||||
}
|
||||
return std::to_string(ret);
|
||||
}
|
||||
} // namespace
|
||||
|
||||
namespace google {
|
||||
|
|
@ -252,8 +265,7 @@ void ToGraphProtoInternal(const onnxruntime::Provider_GraphViewer& graph, Provid
|
|||
}
|
||||
|
||||
// Check if cycle exists in the graph after partitioning
|
||||
bool FindCycleHelper(int i, const std::list<int>* adjacency_map,
|
||||
bool visited[], bool* st, std::vector<int>& cycles) {
|
||||
bool FindCycleHelper(int i, const std::list<int>* adjacency_map, bool visited[], bool* st, std::vector<int>& cycles) {
|
||||
if (!visited[i]) {
|
||||
visited[i] = true;
|
||||
st[i] = true;
|
||||
|
|
@ -748,19 +760,22 @@ common::Status TensorrtExecutionProvider::Provider_Compile(const std::vector<onn
|
|||
|
||||
// Initialize shape range for dynamic shape tensors
|
||||
bool has_dynamic_shape = false;
|
||||
std::vector<int> input_shapes;
|
||||
for (unsigned int i = 0, end = num_inputs; i < end; ++i) {
|
||||
auto input = trt_network->getInput(i);
|
||||
const std::string& input_name = input->getName();
|
||||
nvinfer1::Dims dims = input->getDimensions();
|
||||
int nb_dims = dims.nbDims;
|
||||
|
||||
if (input->isShapeTensor()) {
|
||||
// Shape tensor
|
||||
input_shape_ranges[input_name][0] = std::make_pair(INT_MAX, INT_MIN);
|
||||
has_dynamic_shape = true;
|
||||
for (int i = 0; i < nb_dims; i++)
|
||||
input_shapes.push_back(1); // dummy value
|
||||
} else {
|
||||
// Execution tensor
|
||||
for (int j = 0, end = nb_dims; j < end; ++j) {
|
||||
input_shapes.push_back(dims.d[j]); // could be neg.
|
||||
if (dims.d[j] == -1) {
|
||||
input_shape_ranges[input_name][j] = std::make_pair(INT_MAX, INT_MIN);
|
||||
has_dynamic_shape = true;
|
||||
|
|
@ -781,7 +796,9 @@ common::Status TensorrtExecutionProvider::Provider_Compile(const std::vector<onn
|
|||
tensorrt_ptr::unique_pointer<nvinfer1::ICudaEngine> trt_engine;
|
||||
tensorrt_ptr::unique_pointer<nvinfer1::IExecutionContext> trt_context;
|
||||
if (!has_dynamic_shape) {
|
||||
std::ifstream planFile(GetEnginePath(engine_cache_path_, trt_node_name_with_precision), std::ios::binary | std::ios::in);
|
||||
std::string trt_node_name_with_precision_shape = trt_node_name_with_precision + "_" + GetVecHash(input_shapes);
|
||||
std::string cached_path = GetEnginePath(engine_cache_path_, trt_node_name_with_precision_shape);
|
||||
std::ifstream planFile(cached_path, std::ios::binary | std::ios::in);
|
||||
if (planFile && engine_cache_enable_) {
|
||||
planFile.seekg(0, std::ios::end);
|
||||
int engine_size = planFile.tellg();
|
||||
|
|
@ -790,24 +807,19 @@ common::Status TensorrtExecutionProvider::Provider_Compile(const std::vector<onn
|
|||
planFile.read((char*)engine_buf.get(), engine_size);
|
||||
planFile.close();
|
||||
trt_engine = tensorrt_ptr::unique_pointer<nvinfer1::ICudaEngine>(runtime_->deserializeCudaEngine(engine_buf.get(), engine_size, nullptr));
|
||||
LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Deserialized " + GetEnginePath(engine_cache_path_, trt_node_name_with_precision);
|
||||
if (trt_engine == nullptr) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL,
|
||||
"TensorRT EP could not deserialize engine from " + trt_node_name_with_precision + ".engine");
|
||||
}
|
||||
LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] DeSerialized " + cached_path;
|
||||
} else {
|
||||
trt_engine = tensorrt_ptr::unique_pointer<nvinfer1::ICudaEngine>(trt_builder->buildEngineWithConfig(*trt_network, *trt_config));
|
||||
if (trt_engine == nullptr) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL,
|
||||
"TensorRT EP could not build engine for fused node: " + fused_node->Name());
|
||||
}
|
||||
|
||||
if (engine_cache_enable_) {
|
||||
nvinfer1::IHostMemory* serializedModel = trt_engine->serialize();
|
||||
std::ofstream file(GetEnginePath(engine_cache_path_, trt_node_name_with_precision), std::ios::binary | std::ios::out);
|
||||
std::ofstream file(cached_path, std::ios::binary | std::ios::out);
|
||||
file.write(reinterpret_cast<char*>(serializedModel->data()), serializedModel->size());
|
||||
serializedModel->destroy();
|
||||
LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Serialized " + GetEnginePath(engine_cache_path_, trt_node_name_with_precision);
|
||||
LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Serialized " + cached_path;
|
||||
}
|
||||
}
|
||||
trt_context = tensorrt_ptr::unique_pointer<nvinfer1::IExecutionContext>(trt_engine->createExecutionContext());
|
||||
|
|
@ -859,7 +871,7 @@ common::Status TensorrtExecutionProvider::Provider_Compile(const std::vector<onn
|
|||
&engines_[context->node_name], &contexts_[context->node_name], builders_[context->node_name].get(),
|
||||
networks_[context->node_name].get(), input_info_[context->node_name], output_info_[context->node_name],
|
||||
input_shape_ranges_[context->node_name], &tensorrt_mu_, &fp16_enable_,
|
||||
&max_workspace_size_};
|
||||
&max_workspace_size_, trt_node_name_with_precision, engine_cache_enable_, engine_cache_path_, runtime_};
|
||||
*state = p.release();
|
||||
return 0;
|
||||
};
|
||||
|
|
@ -884,12 +896,11 @@ common::Status TensorrtExecutionProvider::Provider_Compile(const std::vector<onn
|
|||
auto trt_context = trt_state->context->get();
|
||||
int num_inputs = input_indexes.size();
|
||||
int num_outputs = output_indexes.size();
|
||||
|
||||
// Update shape ranges
|
||||
bool engine_update = false;
|
||||
std::unordered_map<std::string, bool> dimension_update;
|
||||
std::unordered_map<std::string, std::vector<int32_t>> tensor_shape_values;
|
||||
nvinfer1::IOptimizationProfile* trt_profile = nullptr;
|
||||
std::vector<int> input_shapes;
|
||||
for (int i = 0, end = num_inputs; i < end; ++i) {
|
||||
auto input = trt_state->network->getInput(i);
|
||||
const std::string& input_name = input->getName();
|
||||
|
|
@ -921,6 +932,7 @@ common::Status TensorrtExecutionProvider::Provider_Compile(const std::vector<onn
|
|||
int32_t* input = new int32_t[shape_size];
|
||||
CUDA_RETURN_IF_ERROR(cudaMemcpy(input, ort.GetTensorData<int32_t>(input_tensor), shape_size * sizeof(int32_t), cudaMemcpyDeviceToHost));
|
||||
for (int j = 0; j < shape_size; ++j) {
|
||||
input_shapes.push_back(input[j]);
|
||||
tensor_shape_values[input_name][j] = input[j];
|
||||
}
|
||||
delete[] input;
|
||||
|
|
@ -930,6 +942,7 @@ common::Status TensorrtExecutionProvider::Provider_Compile(const std::vector<onn
|
|||
int64_t* input = new int64_t[shape_size];
|
||||
CUDA_RETURN_IF_ERROR(cudaMemcpy(input, ort.GetTensorData<int64_t>(input_tensor), shape_size * sizeof(int64_t), cudaMemcpyDeviceToHost));
|
||||
for (int j = 0; j < shape_size; ++j) {
|
||||
input_shapes.push_back(input[j]);
|
||||
tensor_shape_values[input_name][j] = static_cast<int32_t>(input[j]);
|
||||
}
|
||||
delete[] input;
|
||||
|
|
@ -988,12 +1001,11 @@ common::Status TensorrtExecutionProvider::Provider_Compile(const std::vector<onn
|
|||
trt_profile->setShapeValues(input_name.c_str(), nvinfer1::OptProfileSelector::kOPT, &shapes_opt[0], shape_size);
|
||||
trt_profile->setShapeValues(input_name.c_str(), nvinfer1::OptProfileSelector::kMAX, &shapes_max[0], shape_size);
|
||||
}
|
||||
|
||||
} else //execution tensor
|
||||
{
|
||||
} else { // execution tensor
|
||||
nvinfer1::Dims dims_min(dims), dims_opt(dims), dims_max(dims);
|
||||
for (int j = 0, end = nb_dims; j < end; ++j) {
|
||||
const auto& tensor_shape = tensor_shapes[j];
|
||||
input_shapes.push_back(tensor_shape);
|
||||
if (shape_range.find(j) != shape_range.end()) {
|
||||
dims_min.d[j] = shape_range[j].first;
|
||||
dims_opt.d[j] = shape_range[j].second;
|
||||
|
|
@ -1036,22 +1048,51 @@ common::Status TensorrtExecutionProvider::Provider_Compile(const std::vector<onn
|
|||
// Regenerate engine
|
||||
// Only one profile is generated, so no need to explicitly set optimization profile
|
||||
if (engine_update) {
|
||||
auto trt_config = tensorrt_ptr::unique_pointer<nvinfer1::IBuilderConfig>(trt_builder->createBuilderConfig());
|
||||
trt_config->setMaxWorkspaceSize(*(trt_state->max_workspace_size_ptr));
|
||||
trt_config->addOptimizationProfile(trt_profile);
|
||||
if (*(trt_state->fp16_enable_ptr) && trt_builder->platformHasFastFp16()) {
|
||||
trt_config->setFlag(nvinfer1::BuilderFlag::kFP16);
|
||||
std::string trt_node_name_with_precision_shape = trt_state->trt_node_name_with_precision + "_" + GetVecHash(input_shapes);
|
||||
std::string cached_path = GetEnginePath(trt_state->engine_cache_path, trt_node_name_with_precision_shape);
|
||||
std::ifstream planFile(cached_path, std::ios::binary | std::ios::in);
|
||||
if (planFile && trt_state->engine_cache_enable) {
|
||||
planFile.seekg(0, std::ios::end);
|
||||
int engine_size = planFile.tellg();
|
||||
planFile.seekg(0, std::ios::beg);
|
||||
std::unique_ptr<char[]> engine_buf{new char[engine_size]};
|
||||
planFile.read((char*)engine_buf.get(), engine_size);
|
||||
planFile.close();
|
||||
|
||||
auto runtime_ = trt_state->runtime;
|
||||
trt_state->engine->reset();
|
||||
*(trt_state->engine) = tensorrt_ptr::unique_pointer<nvinfer1::ICudaEngine>(
|
||||
runtime_->deserializeCudaEngine(engine_buf.get(), engine_size, nullptr));
|
||||
if (trt_state->engine->get() == nullptr) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP Failed to Build Engine.");
|
||||
}
|
||||
trt_engine = trt_state->engine->get();
|
||||
|
||||
} else {
|
||||
auto trt_config = tensorrt_ptr::unique_pointer<nvinfer1::IBuilderConfig>(trt_builder->createBuilderConfig());
|
||||
trt_config->setMaxWorkspaceSize(*(trt_state->max_workspace_size_ptr));
|
||||
trt_config->addOptimizationProfile(trt_profile);
|
||||
if (*(trt_state->fp16_enable_ptr) && trt_builder->platformHasFastFp16()) {
|
||||
trt_config->setFlag(nvinfer1::BuilderFlag::kFP16);
|
||||
}
|
||||
trt_state->context->reset();
|
||||
trt_state->engine->reset();
|
||||
*(trt_state->engine) = tensorrt_ptr::unique_pointer<nvinfer1::ICudaEngine>(
|
||||
trt_builder->buildEngineWithConfig(*trt_state->network, *trt_config));
|
||||
|
||||
if (trt_state->engine->get() == nullptr) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP Failed to Build Engine.");
|
||||
}
|
||||
trt_engine = trt_state->engine->get();
|
||||
if (trt_state->engine_cache_enable) {
|
||||
nvinfer1::IHostMemory* serializedModel = trt_engine->serialize();
|
||||
std::ofstream file(cached_path, std::ios::binary | std::ios::out);
|
||||
file.write(reinterpret_cast<char*>(serializedModel->data()), serializedModel->size());
|
||||
serializedModel->destroy();
|
||||
LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Serialized " + cached_path;
|
||||
}
|
||||
}
|
||||
trt_state->context->reset();
|
||||
trt_state->engine->reset();
|
||||
*(trt_state->engine) = tensorrt_ptr::unique_pointer<nvinfer1::ICudaEngine>(
|
||||
trt_builder->buildEngineWithConfig(*trt_state->network, *trt_config));
|
||||
|
||||
if (trt_state->engine->get() == nullptr) {
|
||||
return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP failed to build engine.");
|
||||
}
|
||||
trt_engine = trt_state->engine->get();
|
||||
|
||||
*(trt_state->context) = tensorrt_ptr::unique_pointer<nvinfer1::IExecutionContext>(
|
||||
trt_state->engine->get()->createExecutionContext());
|
||||
if (trt_state->context->get() == nullptr) {
|
||||
|
|
|
|||
|
|
@ -74,6 +74,10 @@ struct TensorrtFuncState {
|
|||
OrtMutex* tensorrt_mu_ptr = nullptr;
|
||||
bool* fp16_enable_ptr = nullptr;
|
||||
size_t* max_workspace_size_ptr = nullptr;
|
||||
std::string trt_node_name_with_precision;
|
||||
bool engine_cache_enable;
|
||||
std::string engine_cache_path;
|
||||
nvinfer1::IRuntime* runtime = nullptr;
|
||||
};
|
||||
|
||||
// Logical device representation.
|
||||
|
|
|
|||
Loading…
Reference in a new issue