File size: 5,438 Bytes
139f25e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | #include "model_runner.hpp"
#include <ax_engine_api.h>
#include <ax_sys_api.h>
#include <algorithm>
#include <cstring>
#include <fstream>
#include <iterator>
#include <stdexcept>
namespace {
std::vector<char> read_binary(const std::string& path) {
std::ifstream file(path, std::ios::binary);
if (!file) {
throw std::runtime_error("failed to open " + path);
}
return std::vector<char>(
std::istreambuf_iterator<char>(file),
std::istreambuf_iterator<char>());
}
void check_ax(int ret, const char* message) {
if (ret != 0) {
throw std::runtime_error(message);
}
}
} // namespace
struct ModelRunner::Impl {
AX_ENGINE_HANDLE handle = nullptr;
AX_ENGINE_CONTEXT_T context = nullptr;
AX_ENGINE_IO_INFO_T* info = nullptr;
AX_ENGINE_IO_T io {};
std::vector<AX_ENGINE_IO_BUFFER_T> inputs;
std::vector<AX_ENGINE_IO_BUFFER_T> outputs;
std::vector<char> model;
explicit Impl(const std::string& model_path, const std::string& model_name)
: model(read_binary(model_path)) {
check_ax(AX_SYS_Init(), "AX_SYS_Init failed");
AX_ENGINE_NPU_ATTR_T npu_attr;
std::memset(&npu_attr, 0, sizeof(npu_attr));
if (AX_ENGINE_GetVNPUAttr(&npu_attr) != 0) {
npu_attr.eHardMode = AX_ENGINE_VIRTUAL_NPU_DISABLE;
}
check_ax(AX_ENGINE_Init(&npu_attr), "AX_ENGINE_Init failed");
AX_ENGINE_HANDLE_EXTRA_T extra;
std::memset(&extra, 0, sizeof(extra));
extra.pName = const_cast<AX_S8*>(reinterpret_cast<const AX_S8*>(model_name.c_str()));
check_ax(
AX_ENGINE_CreateHandleV2(
&handle, model.data(), static_cast<AX_U32>(model.size()), &extra),
"AX_ENGINE_CreateHandleV2 failed");
check_ax(AX_ENGINE_CreateContextV2(handle, &context), "AX_ENGINE_CreateContextV2 failed");
check_ax(AX_ENGINE_GetIOInfo(handle, &info), "AX_ENGINE_GetIOInfo failed");
if (!info || info->nInputSize < 1 || info->nOutputSize < 1) {
throw std::runtime_error("model has no input or output tensors");
}
inputs.resize(info->nInputSize);
outputs.resize(info->nOutputSize);
io.pInputs = inputs.data();
io.nInputSize = info->nInputSize;
io.pOutputs = outputs.data();
io.nOutputSize = info->nOutputSize;
for (AX_U32 i = 0; i < info->nInputSize; ++i) {
std::memset(&inputs[i], 0, sizeof(inputs[i]));
inputs[i].nSize = info->pInputs[i].nSize;
check_ax(
AX_SYS_MemAllocCached(
&inputs[i].phyAddr, &inputs[i].pVirAddr, inputs[i].nSize, 128,
reinterpret_cast<const AX_S8*>("model_input")),
"AX_SYS_MemAllocCached failed");
}
for (AX_U32 i = 0; i < info->nOutputSize; ++i) {
std::memset(&outputs[i], 0, sizeof(outputs[i]));
outputs[i].nSize = info->pOutputs[i].nSize;
check_ax(
AX_SYS_MemAllocCached(
&outputs[i].phyAddr, &outputs[i].pVirAddr, outputs[i].nSize, 128,
reinterpret_cast<const AX_S8*>("model_output")),
"AX_SYS_MemAllocCached failed");
}
}
~Impl() {
for (auto& item : inputs) {
if (item.phyAddr) AX_SYS_MemFree(item.phyAddr, item.pVirAddr);
}
for (auto& item : outputs) {
if (item.phyAddr) AX_SYS_MemFree(item.phyAddr, item.pVirAddr);
}
if (handle) AX_ENGINE_DestroyHandle(handle);
AX_ENGINE_Deinit();
AX_SYS_Deinit();
}
};
ModelRunner::ModelRunner(const std::string& model_path, const std::string& model_name)
: impl_(new Impl(model_path, model_name)) {}
ModelRunner::~ModelRunner() {
delete impl_;
}
size_t ModelRunner::NumInputs() const {
return impl_->info->nInputSize;
}
size_t ModelRunner::NumOutputs() const {
return impl_->info->nOutputSize;
}
size_t ModelRunner::InputBytes(size_t index) const {
return impl_->inputs.at(index).nSize;
}
size_t ModelRunner::OutputBytes(size_t index) const {
return impl_->outputs.at(index).nSize;
}
std::vector<std::vector<float>> ModelRunner::Run(const std::vector<std::vector<float>>& inputs) {
if (inputs.size() != NumInputs()) {
throw std::runtime_error("input count mismatch");
}
for (size_t i = 0; i < inputs.size(); ++i) {
if (inputs[i].size() * sizeof(float) > impl_->inputs[i].nSize) {
throw std::runtime_error("input tensor is larger than model input buffer");
}
std::memcpy(impl_->inputs[i].pVirAddr, inputs[i].data(), inputs[i].size() * sizeof(float));
AX_SYS_MflushCache(impl_->inputs[i].phyAddr, impl_->inputs[i].pVirAddr,
impl_->inputs[i].nSize);
}
check_ax(AX_ENGINE_RunSyncV2(impl_->handle, impl_->context, &impl_->io),
"AX_ENGINE_RunSyncV2 failed");
std::vector<std::vector<float>> outputs(NumOutputs());
for (size_t i = 0; i < outputs.size(); ++i) {
AX_SYS_MinvalidateCache(impl_->outputs[i].phyAddr, impl_->outputs[i].pVirAddr,
impl_->outputs[i].nSize);
const size_t count = impl_->outputs[i].nSize / sizeof(float);
const auto* src = static_cast<const float*>(impl_->outputs[i].pVirAddr);
outputs[i].assign(src, src + count);
}
return outputs;
}
|