| #include "ax_engine.hpp" |
|
|
| #include <ax_engine_api.h> |
| #include <ax_sys_api.h> |
|
|
| #include <cstring> |
| #include <fstream> |
| #include <stdexcept> |
|
|
| namespace kantts { |
|
|
| namespace { |
|
|
| std::vector<char> ReadBinary(const std::string& path) { |
| std::ifstream f(path, std::ios::binary); |
| if (!f) throw std::runtime_error("cannot open " + path); |
| return std::vector<char>(std::istreambuf_iterator<char>(f), std::istreambuf_iterator<char>()); |
| } |
|
|
| void Check(int ret, const char* msg) { |
| if (ret != 0) throw std::runtime_error(msg); |
| } |
|
|
| int g_refcount = 0; |
|
|
| } |
|
|
| void AxRuntimeInit() { |
| if (g_refcount++ > 0) return; |
| Check(AX_SYS_Init(), "AX_SYS_Init failed"); |
| AX_ENGINE_NPU_ATTR_T attr; |
| std::memset(&attr, 0, sizeof(attr)); |
| attr.eHardMode = AX_ENGINE_VIRTUAL_NPU_DISABLE; |
| Check(AX_ENGINE_Init(&attr), "AX_ENGINE_Init failed"); |
| } |
|
|
| void AxRuntimeDeinit() { |
| if (--g_refcount > 0) return; |
| AX_ENGINE_Deinit(); |
| AX_SYS_Deinit(); |
| } |
|
|
| struct ModelSession::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; |
| std::map<std::string, AX_U32> input_index; |
| std::map<std::string, AX_U32> output_index; |
|
|
| explicit Impl(const std::string& path) : model(ReadBinary(path)) { |
| AX_ENGINE_HANDLE_EXTRA_T extra; |
| std::memset(&extra, 0, sizeof(extra)); |
| Check(AX_ENGINE_CreateHandleV2(&handle, model.data(), |
| static_cast<AX_U32>(model.size()), &extra), |
| "AX_ENGINE_CreateHandleV2 failed"); |
| Check(AX_ENGINE_CreateContextV2(handle, &context), "AX_ENGINE_CreateContextV2 failed"); |
| Check(AX_ENGINE_GetIOInfo(handle, &info), "AX_ENGINE_GetIOInfo failed"); |
| 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_SYS_MemAllocCached(&inputs[i].phyAddr, &inputs[i].pVirAddr, |
| inputs[i].nSize, 128, (AX_S8*)"kantts_in"), |
| "input alloc failed"); |
| input_index[info->pInputs[i].pName] = i; |
| } |
| 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_SYS_MemAllocCached(&outputs[i].phyAddr, &outputs[i].pVirAddr, |
| outputs[i].nSize, 128, (AX_S8*)"kantts_out"), |
| "output alloc failed"); |
| output_index[info->pOutputs[i].pName] = i; |
| } |
| } |
|
|
| ~Impl() { |
| for (auto& x : inputs) if (x.phyAddr) AX_SYS_MemFree(x.phyAddr, x.pVirAddr); |
| for (auto& x : outputs) if (x.phyAddr) AX_SYS_MemFree(x.phyAddr, x.pVirAddr); |
| if (handle) AX_ENGINE_DestroyHandle(handle); |
| } |
| }; |
|
|
| ModelSession::ModelSession(const std::string& model_path) : impl_(new Impl(model_path)) {} |
| ModelSession::~ModelSession() { delete impl_; } |
|
|
| void ModelSession::SetInput(const std::string& name, const void* data, size_t bytes) { |
| auto it = impl_->input_index.find(name); |
| if (it == impl_->input_index.end()) throw std::runtime_error("no input named " + name); |
| auto& buf = impl_->inputs[it->second]; |
| if (bytes > buf.nSize) throw std::runtime_error("input too large " + name); |
| std::memcpy(buf.pVirAddr, data, bytes); |
| AX_SYS_MflushCache(buf.phyAddr, buf.pVirAddr, buf.nSize); |
| } |
|
|
| void ModelSession::Run() { |
| Check(AX_ENGINE_RunSyncV2(impl_->handle, impl_->context, &impl_->io), |
| "AX_ENGINE_RunSyncV2 failed"); |
| } |
|
|
| size_t ModelSession::OutputBytes(const std::string& name) const { |
| auto it = impl_->output_index.find(name); |
| if (it == impl_->output_index.end()) throw std::runtime_error("no output named " + name); |
| return impl_->outputs[it->second].nSize; |
| } |
|
|
| void ModelSession::GetOutput(const std::string& name, void* out, size_t bytes) const { |
| auto it = impl_->output_index.find(name); |
| if (it == impl_->output_index.end()) throw std::runtime_error("no output named " + name); |
| auto& buf = impl_->outputs[it->second]; |
| if (bytes > buf.nSize) bytes = buf.nSize; |
| AX_SYS_MinvalidateCache(buf.phyAddr, buf.pVirAddr, buf.nSize); |
| std::memcpy(out, buf.pVirAddr, bytes); |
| } |
|
|
| std::vector<int64_t> ModelSession::OutputShape(const std::string& name) const { |
| auto it = impl_->output_index.find(name); |
| if (it == impl_->output_index.end()) throw std::runtime_error("no output named " + name); |
| auto& t = impl_->info->pOutputs[it->second]; |
| std::vector<int64_t> shape(t.nShapeSize); |
| for (AX_U32 i = 0; i < t.nShapeSize; ++i) shape[i] = t.pShape[i]; |
| return shape; |
| } |
|
|
| } |
|
|
