Trim use-can to buildable runtime sources and standalone Orin package

native/trt_engine.cpp

@@ -0,0 +1,106 @@
+#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