50 lines
1.4 KiB
C++
50 lines
1.4 KiB
C++
#pragma once
|
|
|
|
#include "common.hpp"
|
|
|
|
#include <unordered_map>
|
|
|
|
namespace fangjia {
|
|
|
|
class TrtEngine {
|
|
public:
|
|
TrtEngine(std::filesystem::path const& enginePath, TrtLogger& logger);
|
|
~TrtEngine();
|
|
|
|
TrtEngine(TrtEngine const&) = delete;
|
|
TrtEngine& operator=(TrtEngine const&) = delete;
|
|
|
|
void setInputShape(std::string const& name, nvinfer1::Dims const& shape);
|
|
void prepare();
|
|
void bindExternal(std::string const& name, void* pointer, std::size_t bytes);
|
|
void useOwned(std::string const& name);
|
|
bool enqueue(cudaStream_t stream);
|
|
|
|
void* device(std::string const& name);
|
|
std::size_t bytes(std::string const& name) const;
|
|
nvinfer1::Dims shape(std::string const& name) const;
|
|
nvinfer1::DataType dataType(std::string const& name) const;
|
|
std::vector<std::string> inputs() const;
|
|
std::vector<std::string> outputs() const;
|
|
|
|
private:
|
|
struct Buffer {
|
|
void* owned{};
|
|
void* active{};
|
|
std::size_t capacity{};
|
|
std::size_t required{};
|
|
bool external{};
|
|
};
|
|
|
|
Buffer& buffer(std::string const& name);
|
|
Buffer const& buffer(std::string const& name) const;
|
|
|
|
TrtUniquePtr<nvinfer1::IRuntime> runtime_;
|
|
TrtUniquePtr<nvinfer1::ICudaEngine> engine_;
|
|
TrtUniquePtr<nvinfer1::IExecutionContext> context_;
|
|
std::filesystem::path enginePath_;
|
|
std::unordered_map<std::string, Buffer> buffers_;
|
|
};
|
|
|
|
} // namespace fangjia
|