107 lines
4.1 KiB
C++
107 lines
4.1 KiB
C++
#include "trt_engine.hpp"
|
|
|
|
#include <algorithm>
|
|
|
|
namespace fangjia {
|
|
|
|
TrtEngine::TrtEngine(std::filesystem::path const& enginePath, TrtLogger& logger)
|
|
: enginePath_(enginePath) {
|
|
auto bytes = readBinary(enginePath);
|
|
runtime_.reset(nvinfer1::createInferRuntime(logger));
|
|
if (!runtime_) throw std::runtime_error("createInferRuntime failed");
|
|
engine_.reset(runtime_->deserializeCudaEngine(bytes.data(), bytes.size()));
|
|
if (!engine_) throw std::runtime_error("Cannot deserialize " + enginePath.string());
|
|
context_.reset(engine_->createExecutionContext());
|
|
if (!context_) throw std::runtime_error("createExecutionContext failed");
|
|
for (int index = 0; index < engine_->getNbIOTensors(); ++index) {
|
|
buffers_.emplace(engine_->getIOTensorName(index), Buffer{});
|
|
}
|
|
}
|
|
|
|
TrtEngine::~TrtEngine() {
|
|
for (auto& [name, item] : buffers_) {
|
|
(void) name;
|
|
if (item.owned) cudaFree(item.owned);
|
|
}
|
|
}
|
|
|
|
void TrtEngine::setInputShape(std::string const& name, nvinfer1::Dims const& shape) {
|
|
if (!context_->setInputShape(name.c_str(), shape)) {
|
|
throw std::runtime_error("TensorRT rejected input shape for " + name);
|
|
}
|
|
}
|
|
|
|
void TrtEngine::prepare() {
|
|
for (int index = 0; index < engine_->getNbIOTensors(); ++index) {
|
|
std::string const name = engine_->getIOTensorName(index);
|
|
auto const dims = context_->getTensorShape(name.c_str());
|
|
for (int axis = 0; axis < dims.nbDims; ++axis) {
|
|
if (dims.d[axis] < 0) {
|
|
throw std::runtime_error(
|
|
"TensorRT dimensions are incomplete for " + enginePath_.string() +
|
|
":" + name);
|
|
}
|
|
}
|
|
auto& item = buffer(name);
|
|
item.required = volume(dims) * dataTypeSize(engine_->getTensorDataType(name.c_str()));
|
|
if (item.capacity < item.required) {
|
|
if (item.owned) CUDA_CHECK(cudaFree(item.owned));
|
|
CUDA_CHECK(cudaMalloc(&item.owned, item.required));
|
|
item.capacity = item.required;
|
|
}
|
|
item.active = item.owned;
|
|
}
|
|
}
|
|
|
|
bool TrtEngine::enqueue(cudaStream_t stream) {
|
|
for (int index = 0; index < engine_->getNbIOTensors(); ++index) {
|
|
std::string const name = engine_->getIOTensorName(index);
|
|
auto& item = buffer(name);
|
|
if (!item.active) throw std::runtime_error("No device buffer for " + name);
|
|
if (!context_->setTensorAddress(name.c_str(), item.active)) {
|
|
throw std::runtime_error("setTensorAddress failed for " + name);
|
|
}
|
|
}
|
|
return context_->enqueueV3(stream);
|
|
}
|
|
|
|
void* TrtEngine::device(std::string const& name) { return buffer(name).active; }
|
|
std::size_t TrtEngine::bytes(std::string const& name) const { return buffer(name).required; }
|
|
nvinfer1::Dims TrtEngine::shape(std::string const& name) const {
|
|
return context_->getTensorShape(name.c_str());
|
|
}
|
|
nvinfer1::DataType TrtEngine::dataType(std::string const& name) const {
|
|
return engine_->getTensorDataType(name.c_str());
|
|
}
|
|
|
|
std::vector<std::string> TrtEngine::inputs() const {
|
|
std::vector<std::string> result;
|
|
for (int index = 0; index < engine_->getNbIOTensors(); ++index) {
|
|
char const* name = engine_->getIOTensorName(index);
|
|
if (engine_->getTensorIOMode(name) == nvinfer1::TensorIOMode::kINPUT) result.emplace_back(name);
|
|
}
|
|
return result;
|
|
}
|
|
|
|
std::vector<std::string> TrtEngine::outputs() const {
|
|
std::vector<std::string> result;
|
|
for (int index = 0; index < engine_->getNbIOTensors(); ++index) {
|
|
char const* name = engine_->getIOTensorName(index);
|
|
if (engine_->getTensorIOMode(name) == nvinfer1::TensorIOMode::kOUTPUT) result.emplace_back(name);
|
|
}
|
|
return result;
|
|
}
|
|
|
|
TrtEngine::Buffer& TrtEngine::buffer(std::string const& name) {
|
|
auto iterator = buffers_.find(name);
|
|
if (iterator == buffers_.end()) throw std::runtime_error("Unknown tensor " + name);
|
|
return iterator->second;
|
|
}
|
|
TrtEngine::Buffer const& TrtEngine::buffer(std::string const& name) const {
|
|
auto iterator = buffers_.find(name);
|
|
if (iterator == buffers_.end()) throw std::runtime_error("Unknown tensor " + name);
|
|
return iterator->second;
|
|
}
|
|
|
|
} // namespace fangjia
|