123 lines
3.7 KiB
C++
123 lines
3.7 KiB
C++
#pragma once
|
|
|
|
#include <NvInfer.h>
|
|
#include <NvInferVersion.h>
|
|
#include <cuda_runtime.h>
|
|
|
|
#include <cstring>
|
|
#include <cstdint>
|
|
#include <filesystem>
|
|
#include <fstream>
|
|
#include <iostream>
|
|
#include <memory>
|
|
#include <stdexcept>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
namespace fangjia {
|
|
|
|
inline void checkCuda(cudaError_t status, char const* expression, char const* file, int line) {
|
|
if (status != cudaSuccess) {
|
|
throw std::runtime_error(
|
|
std::string("CUDA failure: ") + expression + " at " + file + ":" +
|
|
std::to_string(line) + ": " + cudaGetErrorString(status));
|
|
}
|
|
}
|
|
|
|
#define CUDA_CHECK(expr) ::fangjia::checkCuda((expr), #expr, __FILE__, __LINE__)
|
|
|
|
class TrtLogger final : public nvinfer1::ILogger {
|
|
public:
|
|
explicit TrtLogger(Severity minimum = Severity::kWARNING) : minimum_(minimum) {}
|
|
|
|
void log(Severity severity, char const* message) noexcept override {
|
|
if (severity <= minimum_) {
|
|
std::cerr << "[TensorRT] " << message << '\n';
|
|
}
|
|
}
|
|
|
|
private:
|
|
Severity minimum_;
|
|
};
|
|
|
|
template <typename T>
|
|
struct TrtDeleter {
|
|
void operator()(T* object) const noexcept { delete object; }
|
|
};
|
|
|
|
template <typename T>
|
|
using TrtUniquePtr = std::unique_ptr<T, TrtDeleter<T>>;
|
|
|
|
inline std::vector<std::uint8_t> readBinary(std::filesystem::path const& path) {
|
|
std::ifstream stream(path, std::ios::binary | std::ios::ate);
|
|
if (!stream) {
|
|
throw std::runtime_error("Cannot open " + path.string());
|
|
}
|
|
auto const size = stream.tellg();
|
|
if(size<=0)throw std::runtime_error("Empty or unreadable engine: "+path.string());
|
|
stream.seekg(0);
|
|
std::vector<std::uint8_t> bytes(static_cast<std::size_t>(size));
|
|
if(!stream.read(reinterpret_cast<char*>(bytes.data()), size))throw std::runtime_error("Truncated engine: "+path.string());
|
|
return bytes;
|
|
}
|
|
|
|
inline void writeBinary(std::filesystem::path const& path, void const* data, std::size_t size) {
|
|
std::filesystem::create_directories(path.parent_path());
|
|
std::ofstream stream(path, std::ios::binary);
|
|
if (!stream) {
|
|
throw std::runtime_error("Cannot write " + path.string());
|
|
}
|
|
stream.write(static_cast<char const*>(data), static_cast<std::streamsize>(size));
|
|
}
|
|
|
|
inline std::size_t volume(nvinfer1::Dims const& dims) {
|
|
std::size_t value = 1;
|
|
for (int index = 0; index < dims.nbDims; ++index) {
|
|
if (dims.d[index] < 0) {
|
|
throw std::runtime_error("Unresolved dynamic tensor dimension");
|
|
}
|
|
value *= static_cast<std::size_t>(dims.d[index]);
|
|
}
|
|
return value;
|
|
}
|
|
|
|
inline std::size_t dataTypeSize(nvinfer1::DataType type) {
|
|
switch (type) {
|
|
case nvinfer1::DataType::kFLOAT: return 4;
|
|
case nvinfer1::DataType::kHALF: return 2;
|
|
case nvinfer1::DataType::kINT8: return 1;
|
|
case nvinfer1::DataType::kINT32: return 4;
|
|
#if NV_TENSORRT_MAJOR >= 10
|
|
case nvinfer1::DataType::kINT64: return 8;
|
|
#endif
|
|
case nvinfer1::DataType::kBOOL: return 1;
|
|
default: throw std::runtime_error("Unsupported TensorRT data type");
|
|
}
|
|
}
|
|
|
|
class CudaEventTimer {
|
|
public:
|
|
CudaEventTimer() {
|
|
CUDA_CHECK(cudaEventCreate(&start_));
|
|
CUDA_CHECK(cudaEventCreate(&end_));
|
|
}
|
|
~CudaEventTimer() {
|
|
cudaEventDestroy(start_);
|
|
cudaEventDestroy(end_);
|
|
}
|
|
void start(cudaStream_t stream) { CUDA_CHECK(cudaEventRecord(start_, stream)); }
|
|
float stop(cudaStream_t stream) {
|
|
CUDA_CHECK(cudaEventRecord(end_, stream));
|
|
CUDA_CHECK(cudaEventSynchronize(end_));
|
|
float milliseconds = 0.0F;
|
|
CUDA_CHECK(cudaEventElapsedTime(&milliseconds, start_, end_));
|
|
return milliseconds;
|
|
}
|
|
|
|
private:
|
|
cudaEvent_t start_{};
|
|
cudaEvent_t end_{};
|
|
};
|
|
|
|
} // namespace fangjia
|