[update] add legacy llm backend

This commit is contained in:
LittleMouse
2025-12-18 12:26:44 +08:00
parent cc9d1bc6bc
commit cde5921d3b
12 changed files with 1341 additions and 59 deletions
@@ -92,6 +92,7 @@ public:
std::string kvcache_path;
int precompute_len = 0;
std::vector<int> _token_ids;
static int ax_init_flage_;
task_callback_t out_callback_;
bool enoutput_;
bool enstream_;
@@ -630,9 +631,38 @@ public:
return port;
}
void _ax_init()
{
if (!ax_init_flage_) {
int ret = AX_SYS_Init();
if (0 != ret) {
fprintf(stderr, "AX_SYS_Init failed! ret = 0x%x\n", ret);
}
AX_ENGINE_NPU_ATTR_T npu_attr;
memset(&npu_attr, 0, sizeof(npu_attr));
ret = AX_ENGINE_Init(&npu_attr);
if (0 != ret) {
fprintf(stderr, "Init ax-engine failed{0x%8x}.\n", ret);
}
}
ax_init_flage_++;
}
void _ax_deinit()
{
if (ax_init_flage_ > 0) {
--ax_init_flage_;
if (!ax_init_flage_) {
AX_ENGINE_Deinit();
AX_SYS_Deinit();
}
}
}
llm_task(const std::string &workid) : tokenizer_server_flage_(false), port_(getNextPort())
{
inference_run_ = std::make_unique<std::thread>(std::bind(&llm_task::run, this));
_ax_init();
}
void start()
@@ -666,10 +696,12 @@ public:
if (lToken2Wav_) {
lToken2Wav_->Deinit();
}
_ax_deinit();
}
};
std::atomic<unsigned int> llm_task::next_port_{8070};
int llm_task::ax_init_flage_ = 0;
#undef CONFIG_AUTO_SET
@@ -18,6 +18,7 @@
#include "cqdm.h"
#include "timer.hpp"
#include "ax_sys_api.h"
#include "ax_engine_api.h"
#include "utils/sampling.hpp"
#include "utils/utils.hpp"
@@ -463,20 +464,11 @@ public:
layer.layer.inference(_attr.prefill_grpid);
auto &input_decoder_k_cache = layer.layer.get_input(decode_grpid, "K_cache");
auto &input_decoder_v_cache = layer.layer.get_input(decode_grpid, "V_cache");
auto &output_k_cache = layer.layer.get_output(_attr.prefill_grpid, "K_cache_out");
auto &output_v_cache = layer.layer.get_output(_attr.prefill_grpid, "V_cache_out");
int kv_offset = (_attr.precompute_len + p * _attr.prefill_token_num) * _attr.kv_cache_size;
memcpy((unsigned short *)input_decoder_k_cache.pVirAddr + kv_offset, (void *)output_k_cache.pVirAddr,
sizeof(unsigned short) * input_num_token * _attr.kv_cache_size);
memcpy((unsigned short *)input_decoder_v_cache.pVirAddr + kv_offset, (void *)output_v_cache.pVirAddr,
sizeof(unsigned short) * input_num_token * _attr.kv_cache_size);
for (int gid = _attr.prefill_grpid + 1; gid < prefill_split_num + 1; gid++) {
auto &input_prefill_k_cache = layer.layer.get_input(gid, "K_cache");
memcpy((unsigned short *)input_prefill_k_cache.pVirAddr + kv_offset,
+34 -2
View File
@@ -67,6 +67,7 @@ public:
std::string kvcache_path;
int precompute_len = 0;
std::vector<int> _token_ids;
static int ax_init_flage_;
task_callback_t out_callback_;
bool enoutput_;
bool enstream_;
@@ -208,7 +209,7 @@ public:
SLOGI("port_=%s model_id=%s content=%s", std::to_string(port_).c_str(),
(base_model + std::string("tokenizer")).c_str(), prompt_.c_str());
std::this_thread::sleep_for(std::chrono::seconds(15));
std::this_thread::sleep_for(std::chrono::seconds(5));
};
auto process_field = [&](std::string &field, const char *name_for_log) -> bool {
@@ -432,9 +433,38 @@ public:
return port;
}
void _ax_init()
{
if (!ax_init_flage_) {
int ret = AX_SYS_Init();
if (0 != ret) {
fprintf(stderr, "AX_SYS_Init failed! ret = 0x%x\n", ret);
}
AX_ENGINE_NPU_ATTR_T npu_attr;
memset(&npu_attr, 0, sizeof(npu_attr));
ret = AX_ENGINE_Init(&npu_attr);
if (0 != ret) {
fprintf(stderr, "Init ax-engine failed{0x%8x}.\n", ret);
}
}
ax_init_flage_++;
}
void _ax_deinit()
{
if (ax_init_flage_ > 0) {
--ax_init_flage_;
if (!ax_init_flage_) {
AX_ENGINE_Deinit();
AX_SYS_Deinit();
}
}
}
llm_task(const std::string &workid) : tokenizer_server_flage_(false), port_(getNextPort())
{
inference_run_ = std::make_unique<std::thread>(std::bind(&llm_task::run, this));
_ax_init();
}
void start()
@@ -469,10 +499,12 @@ public:
if (lLaMa_ctx_) {
lLaMa_ctx_->Deinit();
}
_ax_deinit();
}
};
std::atomic<unsigned int> llm_task::next_port_{8080};
int llm_task::ax_init_flage_ = 0;
#undef CONFIG_AUTO_SET
@@ -527,7 +559,7 @@ public:
void pause(const std::string &work_id, const std::string &object, const std::string &data) override
{
SLOGI("llm_asr::work:%s", data.c_str());
SLOGI("llm_llm::work:%s", data.c_str());
nlohmann::json error_body;
int work_id_num = sample_get_work_id_num(work_id);
@@ -7,12 +7,14 @@
#include "Tokenizer/Tokenizer.hpp"
#include "LLMEmbedSelector.hpp"
#include "ax_model_runner/ax_model_runner_ax650.hpp"
#include "ax_model_runner/legacy/ax_model_runner_ax650.hpp"
#include "ax_cmm_utils.hpp"
#include "cqdm.h"
#include "timer.hpp"
#include "LLMPostprocess.hpp"
#include <ax_sys_api.h>
#include "ax_sys_api.h"
#include "ax_engine_api.h"
#include <arm_neon.h>
#define ALIGN_DOWN(x, a) ((x) & ~((a) - 1))
@@ -78,14 +80,14 @@ private:
LLMAttrType _attr;
struct LLMLayer {
ax_runner_ax650 layer;
ax::legacy::ax_runner_ax650 layer;
std::string filename;
MMap layer_buffer;
std::vector<char> layer_buffer_vec;
};
std::vector<LLMLayer> llama_layers;
ax_runner_ax650 llama_post;
ax::legacy::ax_runner_ax650 llama_post;
int prefill_grpid = 1;
int decode_grpid = 0;
@@ -243,9 +245,9 @@ public:
void Deinit()
{
for (int i = 0; i < _attr.axmodel_num; i++) {
llama_layers[i].layer.deinit();
llama_layers[i].layer.release();
}
llama_post.deinit();
llama_post.release();
embed_selector.Deinit();
}
@@ -1245,9 +1247,6 @@ public:
layer.layer.inference(_attr.prefill_grpid);
auto &input_decoder_k_cache = layer.layer.get_input(decode_grpid, "K_cache");
auto &input_decoder_v_cache = layer.layer.get_input(decode_grpid, "V_cache");
auto &input_prefill_k_cache = layer.layer.get_input(_attr.prefill_grpid, "K_cache");
auto &input_prefill_v_cache = layer.layer.get_input(_attr.prefill_grpid, "V_cache");
@@ -1256,12 +1255,6 @@ public:
int kv_offset = (_attr.precompute_len + p * _attr.prefill_token_num) * _attr.kv_cache_size;
memcpy((unsigned short *)input_decoder_k_cache.pVirAddr + kv_offset, (void *)output_k_cache.pVirAddr,
sizeof(unsigned short) * input_num_token * _attr.kv_cache_size);
memcpy((unsigned short *)input_decoder_v_cache.pVirAddr + kv_offset, (void *)output_v_cache.pVirAddr,
sizeof(unsigned short) * input_num_token * _attr.kv_cache_size);
memcpy((unsigned short *)input_prefill_k_cache.pVirAddr + kv_offset, (void *)output_k_cache.pVirAddr,
sizeof(unsigned short) * input_num_token * _attr.kv_cache_size);
@@ -0,0 +1,177 @@
#pragma once
#include <vector>
#include <string>
#include <map>
#include <stdexcept>
namespace ax::legacy {
typedef enum _color_space_e {
axdl_color_space_unknown,
axdl_color_space_nv12,
axdl_color_space_nv21,
axdl_color_space_bgr,
axdl_color_space_rgb,
} ax_color_space_e;
typedef struct _image_t {
unsigned long long int pPhy;
void *pVir;
unsigned int nSize;
unsigned int nWidth;
unsigned int nHeight;
ax_color_space_e eDtype;
union {
int tStride_H, tStride_W, tStride_C;
};
} ax_image_t;
typedef struct {
std::string sName;
unsigned int nIdx;
std::vector<unsigned int> vShape;
int nSize;
unsigned long phyAddr;
void *pVirAddr;
} ax_runner_tensor_t;
class ax_runner_base {
protected:
std::vector<ax_runner_tensor_t> moutput_tensors;
std::vector<ax_runner_tensor_t> minput_tensors;
std::vector<std::vector<ax_runner_tensor_t>> mgroup_output_tensors;
std::vector<std::vector<ax_runner_tensor_t>> mgroup_input_tensors;
std::map<std::string, ax_runner_tensor_t> map_output_tensors;
std::map<std::string, ax_runner_tensor_t> map_input_tensors;
std::map<std::string, std::vector<ax_runner_tensor_t>> map_group_output_tensors;
std::map<std::string, std::vector<ax_runner_tensor_t>> map_group_input_tensors;
public:
virtual int init(const char *model_file, bool use_mmap = false) = 0;
virtual int init(char *model_buffer, size_t model_size) = 0;
virtual void deinit() = 0;
int get_num_inputs()
{
return minput_tensors.size();
};
int get_num_outputs()
{
return moutput_tensors.size();
};
int get_num_input_groups()
{
return mgroup_input_tensors.size();
};
int get_num_output_groups()
{
return mgroup_output_tensors.size();
};
const ax_runner_tensor_t &get_input(int idx)
{
return minput_tensors[idx];
}
const ax_runner_tensor_t *get_inputs_ptr()
{
return minput_tensors.data();
}
const ax_runner_tensor_t &get_input(std::string name)
{
if (map_input_tensors.size() == 0) {
for (size_t i = 0; i < minput_tensors.size(); i++) {
map_input_tensors[minput_tensors[i].sName] = minput_tensors[i];
}
}
if (map_input_tensors.find(name) == map_input_tensors.end()) {
throw std::runtime_error("input tensor not found: " + name);
}
return map_input_tensors[name];
}
const ax_runner_tensor_t &get_input(int grpid, int idx)
{
return mgroup_input_tensors[grpid][idx];
}
const ax_runner_tensor_t *get_inputs_ptr(int grpid)
{
return mgroup_input_tensors[grpid].data();
}
const ax_runner_tensor_t &get_input(int grpid, std::string name)
{
if (map_group_input_tensors.size() == 0) {
for (size_t i = 0; i < mgroup_input_tensors.size(); i++) {
for (size_t j = 0; j < mgroup_input_tensors[i].size(); j++) {
map_group_input_tensors[mgroup_input_tensors[i][j].sName].push_back(mgroup_input_tensors[i][j]);
}
}
}
if (map_group_input_tensors.find(name) == map_group_input_tensors.end()) {
throw std::runtime_error("input tensor not found: " + name);
}
return map_group_input_tensors[name][grpid];
// return map_input_tensors[name];
}
const ax_runner_tensor_t &get_output(int idx)
{
return moutput_tensors[idx];
}
const ax_runner_tensor_t *get_outputs_ptr()
{
return moutput_tensors.data();
}
const ax_runner_tensor_t &get_output(std::string name)
{
if (map_output_tensors.size() == 0) {
for (size_t i = 0; i < moutput_tensors.size(); i++) {
map_output_tensors[moutput_tensors[i].sName] = moutput_tensors[i];
}
}
if (map_output_tensors.find(name) == map_output_tensors.end()) {
throw std::runtime_error("output tensor not found: " + name);
}
return map_output_tensors[name];
}
const ax_runner_tensor_t &get_output(int grpid, int idx)
{
return mgroup_output_tensors[grpid][idx];
}
const ax_runner_tensor_t *get_outputs_ptr(int grpid)
{
return mgroup_output_tensors[grpid].data();
}
const ax_runner_tensor_t &get_output(int grpid, std::string name)
{
if (map_group_output_tensors.size() == 0) {
for (size_t i = 0; i < mgroup_output_tensors.size(); i++) {
for (size_t j = 0; j < mgroup_output_tensors[i].size(); j++) {
map_group_output_tensors[mgroup_output_tensors[i][j].sName].push_back(mgroup_output_tensors[i][j]);
}
}
}
if (map_group_output_tensors.find(name) == map_group_output_tensors.end()) {
throw std::runtime_error("input tensor not found: " + name);
}
return map_group_output_tensors[name][grpid];
}
virtual int inference() = 0;
virtual int inference(int grpid) = 0;
int operator()()
{
return inference();
}
};
// int ax_cmmcpy(unsigned long long int dst, unsigned long long int src, int size);
} // namespace ax::legacy
@@ -0,0 +1,411 @@
#include "ax_model_runner_ax650.hpp"
#include "string.h"
#include "fstream"
#include "memory"
// #include "utilities/file.hpp"
#include <ax_sys_api.h>
#include <ax_ivps_api.h>
#include <ax_engine_api.h>
#include <fcntl.h>
#include "memory_utils.hpp"
#include "sample_log.h"
#define AX_CMM_ALIGN_SIZE 128
namespace ax::legacy {
const char *AX_CMM_SESSION_NAME = "npu";
typedef enum {
AX_ENGINE_ABST_DEFAULT = 0,
AX_ENGINE_ABST_CACHED = 1,
} AX_ENGINE_ALLOC_BUFFER_STRATEGY_T;
typedef std::pair<AX_ENGINE_ALLOC_BUFFER_STRATEGY_T, AX_ENGINE_ALLOC_BUFFER_STRATEGY_T> INPUT_OUTPUT_ALLOC_STRATEGY;
static void print_io_info(AX_ENGINE_IO_INFO_T *io_info)
{
static std::map<AX_ENGINE_DATA_TYPE_T, const char *> data_type = {
{AX_ENGINE_DT_UNKNOWN, "UNKNOWN"},
{AX_ENGINE_DT_UINT8, "UINT8"},
{AX_ENGINE_DT_UINT16, "UINT16"},
{AX_ENGINE_DT_FLOAT32, "FLOAT32"},
{AX_ENGINE_DT_SINT16, "SINT16"},
{AX_ENGINE_DT_SINT8, "SINT8"},
{AX_ENGINE_DT_SINT32, "SINT32"},
{AX_ENGINE_DT_UINT32, "UINT32"},
{AX_ENGINE_DT_FLOAT64, "FLOAT64"},
{AX_ENGINE_DT_UINT10_PACKED, "UINT10_PACKED"},
{AX_ENGINE_DT_UINT12_PACKED, "UINT12_PACKED"},
{AX_ENGINE_DT_UINT14_PACKED, "UINT14_PACKED"},
{AX_ENGINE_DT_UINT16_PACKED, "UINT16_PACKED"},
};
static std::map<AX_ENGINE_COLOR_SPACE_T, const char *> color_type = {
{AX_ENGINE_CS_FEATUREMAP, "FEATUREMAP"},
{AX_ENGINE_CS_RAW8, "RAW8"},
{AX_ENGINE_CS_RAW10, "RAW10"},
{AX_ENGINE_CS_RAW12, "RAW12"},
{AX_ENGINE_CS_RAW14, "RAW14"},
{AX_ENGINE_CS_RAW16, "RAW16"},
{AX_ENGINE_CS_NV12, "NV12"},
{AX_ENGINE_CS_NV21, "NV21"},
{AX_ENGINE_CS_RGB, "RGB"},
{AX_ENGINE_CS_BGR, "BGR"},
{AX_ENGINE_CS_RGBA, "RGBA"},
{AX_ENGINE_CS_GRAY, "GRAY"},
{AX_ENGINE_CS_YUV444, "YUV444"},
};
printf("\ninput size: %d\n", io_info->nInputSize);
for (uint32_t i = 0; i < io_info->nInputSize; ++i) {
// print shape info,like [batchsize x channel x height x width]
auto &info = io_info->pInputs[i];
printf(" name: \e[1;32m%8s", info.pName);
std::string dt = "unknown";
if (data_type.find(info.eDataType) != data_type.end()) {
dt = data_type[info.eDataType];
printf(" \e[1;34m[%s] ", dt.c_str());
} else {
printf(" \e[1;31m[%s] ", dt.c_str());
}
std::string ct = "unknown";
if (info.pExtraMeta && color_type.find(info.pExtraMeta->eColorSpace) != color_type.end()) {
ct = color_type[info.pExtraMeta->eColorSpace];
printf("\e[1;34m[%s]", ct.c_str());
} else {
printf("\e[1;31m[%s]", ct.c_str());
}
printf(" \n \e[1;31m");
for (AX_U8 s = 0; s < info.nShapeSize; s++) {
printf("%d", info.pShape[s]);
if (s != info.nShapeSize - 1) {
printf(" x ");
}
}
printf("\e[0m\n\n");
}
printf("\noutput size: %d\n", io_info->nOutputSize);
for (uint32_t i = 0; i < io_info->nOutputSize; ++i) {
// print shape info,like [batchsize x channel x height x width]
auto &info = io_info->pOutputs[i];
printf(" name: \e[1;32m%8s \e[1;34m[%s]\e[0m\n \e[1;31m", info.pName, data_type[info.eDataType]);
for (AX_U8 s = 0; s < info.nShapeSize; s++) {
printf("%d", info.pShape[s]);
if (s != info.nShapeSize - 1) {
printf(" x ");
}
}
printf("\e[0m\n\n");
}
}
void free_io_index(AX_ENGINE_IO_BUFFER_T *io_buf, int index)
{
for (int i = 0; i < index; ++i) {
AX_ENGINE_IO_BUFFER_T *pBuf = io_buf + i;
AX_SYS_MemFree(pBuf->phyAddr, pBuf->pVirAddr);
}
}
void free_io(AX_ENGINE_IO_T *io)
{
for (size_t j = 0; j < io->nInputSize; ++j) {
AX_ENGINE_IO_BUFFER_T *pBuf = io->pInputs + j;
AX_SYS_MemFree(pBuf->phyAddr, pBuf->pVirAddr);
}
for (size_t j = 0; j < io->nOutputSize; ++j) {
AX_ENGINE_IO_BUFFER_T *pBuf = io->pOutputs + j;
AX_SYS_MemFree(pBuf->phyAddr, pBuf->pVirAddr);
}
delete[] io->pInputs;
delete[] io->pOutputs;
}
static inline int prepare_io(AX_ENGINE_IO_INFO_T *info, AX_ENGINE_IO_T *io_data, INPUT_OUTPUT_ALLOC_STRATEGY strategy)
{
memset(io_data, 0, sizeof(*io_data));
io_data->pInputs = new AX_ENGINE_IO_BUFFER_T[info->nInputSize];
io_data->nInputSize = info->nInputSize;
auto ret = 0;
for (uint i = 0; i < info->nInputSize; ++i) {
auto meta = info->pInputs[i];
auto buffer = &io_data->pInputs[i];
if (strategy.first == AX_ENGINE_ABST_CACHED) {
ret = AX_SYS_MemAllocCached((AX_U64 *)(&buffer->phyAddr), &buffer->pVirAddr, meta.nSize, AX_CMM_ALIGN_SIZE,
(const AX_S8 *)(AX_CMM_SESSION_NAME));
} else {
ret = AX_SYS_MemAlloc((AX_U64 *)(&buffer->phyAddr), &buffer->pVirAddr, meta.nSize, AX_CMM_ALIGN_SIZE,
(const AX_S8 *)(AX_CMM_SESSION_NAME));
}
if (ret != 0) {
free_io_index(io_data->pInputs, i);
fprintf(stderr, "Allocate input{%d} { phy: %p, vir: %p, size: %lu Bytes }. fail \n", i,
(void *)buffer->phyAddr, buffer->pVirAddr, (long)meta.nSize);
return ret;
}
memset(buffer->pVirAddr, 0, meta.nSize);
// fprintf(stderr, "Allocate input{%d} { phy: %p, vir: %p, size: %lu Bytes }. \n", i, (void*)buffer->phyAddr,
// buffer->pVirAddr, (long)meta.nSize);
}
io_data->pOutputs = new AX_ENGINE_IO_BUFFER_T[info->nOutputSize];
io_data->nOutputSize = info->nOutputSize;
for (uint i = 0; i < info->nOutputSize; ++i) {
auto meta = info->pOutputs[i];
auto buffer = &io_data->pOutputs[i];
buffer->nSize = meta.nSize;
if (strategy.second == AX_ENGINE_ABST_CACHED) {
ret = AX_SYS_MemAllocCached((AX_U64 *)(&buffer->phyAddr), &buffer->pVirAddr, meta.nSize, AX_CMM_ALIGN_SIZE,
(const AX_S8 *)(AX_CMM_SESSION_NAME));
} else {
ret = AX_SYS_MemAlloc((AX_U64 *)(&buffer->phyAddr), &buffer->pVirAddr, meta.nSize, AX_CMM_ALIGN_SIZE,
(const AX_S8 *)(AX_CMM_SESSION_NAME));
}
if (ret != 0) {
fprintf(stderr, "Allocate output{%d} { phy: %p, vir: %p, size: %lu Bytes }. fail \n", i,
(void *)buffer->phyAddr, buffer->pVirAddr, (long)meta.nSize);
free_io_index(io_data->pInputs, io_data->nInputSize);
free_io_index(io_data->pOutputs, i);
return ret;
}
memset(buffer->pVirAddr, 0, meta.nSize);
// fprintf(stderr, "Allocate output{%d} { phy: %p, vir: %p, size: %lu Bytes }.\n", i, (void*)buffer->phyAddr,
// buffer->pVirAddr, (long)meta.nSize);
}
return 0;
}
struct ax_joint_runner_ax650_handle_t {
AX_ENGINE_HANDLE handle;
AX_ENGINE_CONTEXT_T context;
std::vector<AX_ENGINE_IO_INFO_T *> io_info;
std::vector<AX_ENGINE_IO_T> io_data;
// int algo_width, algo_height;
// int algo_colorformat;
};
int ax_runner_ax650::sub_init()
{
// 4. create context
int ret = AX_ENGINE_CreateContext(m_handle->handle);
if (0 != ret) {
ALOGE("AX_ENGINE_CreateContext");
return ret;
}
ret = AX_ENGINE_CreateContextV2(m_handle->handle, &m_handle->context);
if (0 != ret) {
ALOGE("AX_ENGINE_CreateContextV2");
return ret;
}
// fprintf(stdout, "Engine creating context is done.\n");
// 5. set io
AX_U32 io_count = 0;
ret = AX_ENGINE_GetGroupIOInfoCount(m_handle->handle, &io_count);
if (0 != ret) {
ALOGE("AX_ENGINE_GetGroupIOInfoCount");
return ret;
}
// ALOGI("io_count=%d", io_count);
m_handle->io_info.resize(io_count);
m_handle->io_data.resize(io_count);
mgroup_input_tensors.resize(io_count);
mgroup_output_tensors.resize(io_count);
// fprintf(stdout, "Engine get io info is done. \n");
// 6. alloc io
if (!_parepare_io) {
for (size_t grpid = 0; grpid < io_count; grpid++) {
AX_ENGINE_IO_INFO_T *io_info = nullptr;
ret = AX_ENGINE_GetGroupIOInfo(m_handle->handle, grpid, &io_info);
if (0 != ret) {
ALOGE("AX_ENGINE_GetIOInfo");
return ret;
}
// print_io_info(io_info);
m_handle->io_info[grpid] = io_info;
ret = prepare_io(m_handle->io_info[grpid], &m_handle->io_data[grpid],
std::make_pair(AX_ENGINE_ABST_DEFAULT, AX_ENGINE_ABST_CACHED));
if (0 != ret) {
ALOGE("prepare_io grpid=%d", grpid);
return ret;
}
}
for (size_t grpid = 0; grpid < io_count; grpid++) {
auto &io_info = m_handle->io_info[grpid];
auto &io_data = m_handle->io_data[grpid];
for (size_t i = 0; i < io_info->nOutputSize; i++) {
ax_runner_tensor_t tensor;
tensor.nIdx = i;
tensor.sName = std::string(io_info->pOutputs[i].pName);
tensor.nSize = io_info->pOutputs[i].nSize;
for (size_t j = 0; j < io_info->pOutputs[i].nShapeSize; j++) {
tensor.vShape.push_back(io_info->pOutputs[i].pShape[j]);
}
tensor.phyAddr = io_data.pOutputs[i].phyAddr;
tensor.pVirAddr = io_data.pOutputs[i].pVirAddr;
mgroup_output_tensors[grpid].push_back(tensor);
}
for (size_t i = 0; i < io_info->nInputSize; i++) {
ax_runner_tensor_t tensor;
tensor.nIdx = i;
tensor.sName = std::string(io_info->pInputs[i].pName);
tensor.nSize = io_info->pInputs[i].nSize;
for (size_t j = 0; j < io_info->pInputs[i].nShapeSize; j++) {
tensor.vShape.push_back(io_info->pInputs[i].pShape[j]);
}
tensor.phyAddr = io_data.pInputs[i].phyAddr;
tensor.pVirAddr = io_data.pInputs[i].pVirAddr;
mgroup_input_tensors[grpid].push_back(tensor);
}
}
moutput_tensors = mgroup_output_tensors[0];
minput_tensors = mgroup_input_tensors[0];
_parepare_io = true;
} else {
}
return ret;
}
int ax_runner_ax650::init(const char *model_file, bool use_mmap)
{
if (use_mmap) {
MMap model_buffer(model_file);
if (!model_buffer.data()) {
ALOGE("mmap");
return -1;
}
auto ret = init((char *)model_buffer.data(), model_buffer.size());
model_buffer.close_file();
return ret;
} else {
char *model_buffer;
size_t len;
if (!read_file(model_file, &model_buffer, &len)) {
ALOGE("read_file");
return -1;
}
auto ret = init(model_buffer, len);
delete[] model_buffer;
return ret;
}
}
int ax_runner_ax650::init(char *model_buffer, size_t model_size)
{
if (!m_handle) {
m_handle = new ax_joint_runner_ax650_handle_t;
}
// static bool b_init = false;
// if (!b_init) {
// // 1. init engine
// AX_ENGINE_NPU_ATTR_T npu_attr;
// memset(&npu_attr, 0, sizeof(npu_attr));
// npu_attr.eHardMode = AX_ENGINE_VIRTUAL_NPU_DISABLE;
// AX_SYS_Init();
// auto ret = AX_ENGINE_Init(&npu_attr);
// if (0 != ret) {
// return ret;
// }
// b_init = true;
// }
// 3. create handle
int ret = AX_ENGINE_CreateHandle(&m_handle->handle, model_buffer, model_size);
if (0 != ret) {
ALOGE("AX_ENGINE_CreateHandle");
return ret;
}
// fprintf(stdout, "Engine creating handle is done.\n");
return sub_init();
}
void ax_runner_ax650::release()
{
if (m_handle && m_handle->handle) {
for (size_t i = 0; i < m_handle->io_data.size(); i++) {
/* code */
free_io(&m_handle->io_data[i]);
}
AX_ENGINE_DestroyHandle(m_handle->handle);
m_handle->handle = nullptr;
}
if (m_handle) {
delete m_handle;
m_handle = nullptr;
}
moutput_tensors.clear();
minput_tensors.clear();
map_input_tensors.clear();
map_output_tensors.clear();
mgroup_output_tensors.clear();
mgroup_input_tensors.clear();
map_group_input_tensors.clear();
map_group_output_tensors.clear();
// AX_ENGINE_Deinit();
}
void ax_runner_ax650::deinit()
{
if (m_handle && m_handle->handle) {
// free_io(&m_handle->io_data);
// mtensors.clear();
// minput_tensors.clear();
// map_input_tensors.clear();
// map_tensors.clear();
AX_ENGINE_DestroyHandle(m_handle->handle);
m_handle->handle = nullptr;
// delete m_handle;
// m_handle = nullptr;
}
// AX_ENGINE_Deinit();
}
int ax_runner_ax650::inference()
{
int ret = AX_ENGINE_RunSync(m_handle->handle, &m_handle->io_data[0]);
for (size_t i = 0; i < get_num_outputs(); i++) {
auto &tensor = get_output(i);
AX_SYS_MinvalidateCache(tensor.phyAddr, tensor.pVirAddr, tensor.nSize);
}
return ret;
}
int ax_runner_ax650::inference(int grpid)
{
int ret = AX_ENGINE_RunGroupIOSync(m_handle->handle, m_handle->context, grpid, &m_handle->io_data[grpid]);
for (size_t i = 0; i < get_num_outputs(); i++) {
auto &tensor = get_output(grpid, i);
AX_SYS_MinvalidateCache(tensor.phyAddr, tensor.pVirAddr, tensor.nSize);
}
return ret;
}
} // namespace ax::legacy
@@ -0,0 +1,25 @@
#pragma once
#include "ax_model_runner.hpp"
namespace ax::legacy {
class ax_runner_ax650 : public ax_runner_base {
protected:
struct ax_joint_runner_ax650_handle_t *m_handle = nullptr;
bool _parepare_io = false;
int sub_init();
public:
int init(const char *model_file, bool use_mmap = false) override;
int init(char *model_buffer, size_t model_size) override;
void release();
void deinit() override;
int inference() override;
int inference(int grpid) override;
};
} // namespace ax::legacy
@@ -82,6 +82,7 @@ public:
std::string kvcache_path;
int precompute_len = 0;
std::vector<int> _token_ids;
static int ax_init_flage_;
task_callback_t out_callback_;
bool enoutput_;
bool enstream_;
@@ -572,6 +573,7 @@ public:
llm_task(const std::string &workid) : tokenizer_server_flage_(false), port_(getNextPort())
{
_ax_init();
}
void start()
@@ -582,6 +584,34 @@ public:
{
}
void _ax_init()
{
if (!ax_init_flage_) {
int ret = AX_SYS_Init();
if (0 != ret) {
fprintf(stderr, "AX_SYS_Init failed! ret = 0x%x\n", ret);
}
AX_ENGINE_NPU_ATTR_T npu_attr;
memset(&npu_attr, 0, sizeof(npu_attr));
ret = AX_ENGINE_Init(&npu_attr);
if (0 != ret) {
fprintf(stderr, "Init ax-engine failed{0x%8x}.\n", ret);
}
}
ax_init_flage_++;
}
void _ax_deinit()
{
if (ax_init_flage_ > 0) {
--ax_init_flage_;
if (!ax_init_flage_) {
AX_ENGINE_Deinit();
AX_SYS_Deinit();
}
}
}
~llm_task()
{
stop();
@@ -598,10 +628,12 @@ public:
if (qwen_) {
qwen_->Deinit();
}
_ax_deinit();
}
};
std::atomic<unsigned int> llm_task::next_port_{8090};
int llm_task::ax_init_flage_ = 0;
#undef CONFIG_AUTO_SET
@@ -11,11 +11,13 @@
#include "Tokenizer/Tokenizer.hpp"
#include "LLMEmbedSelector.hpp"
#include "ax_model_runner/ax_model_runner_ax650.hpp"
#include "ax_model_runner/legacy/ax_model_runner_ax650.hpp"
#include "ax_cmm_utils.hpp"
#include "cqdm.h"
#include "timer.hpp"
#include "opencv2/opencv.hpp"
#include "ax_sys_api.h"
#include "ax_engine_api.h"
#include "LLMPostprocess.hpp"
#define ALIGN_DOWN(x, a) ((x) & ~((a) - 1))
@@ -96,16 +98,16 @@ private:
LLMAttrType _attr;
struct LLMLayer {
ax_runner_ax650 layer;
ax::legacy::ax_runner_ax650 layer;
std::string filename;
MMap layer_buffer;
std::vector<char> layer_buffer_vec;
};
std::vector<LLMLayer> llama_layers;
ax_runner_ax650 llama_post;
ax::legacy::ax_runner_ax650 llama_post;
ax_runner_ax650 vpm_encoder, vpm_resampler;
ax::legacy::ax_runner_ax650 vpm_encoder, vpm_resampler;
int prefill_grpid = 1;
int decode_grpid = 0;
@@ -285,11 +287,11 @@ public:
void Deinit()
{
for (int i = 0; i < _attr.axmodel_num; i++) {
llama_layers[i].layer.deinit();
llama_layers[i].layer.release();
}
llama_post.deinit();
vpm_encoder.deinit();
vpm_resampler.deinit();
llama_post.release();
vpm_encoder.release();
vpm_resampler.release();
embed_selector.Deinit();
}
@@ -978,9 +980,6 @@ public:
layer.layer.inference(prefill_grpid);
auto &input_decoder_k_cache = layer.layer.get_input(decode_grpid, "K_cache");
auto &input_decoder_v_cache = layer.layer.get_input(decode_grpid, "V_cache");
auto &input_prefill_k_cache = layer.layer.get_input(prefill_grpid, "K_cache");
auto &input_prefill_v_cache = layer.layer.get_input(prefill_grpid, "V_cache");
@@ -989,12 +988,6 @@ public:
int kv_offset = (p * _attr.prefill_token_num) * _attr.kv_cache_size;
memcpy((unsigned short *)input_decoder_k_cache.pVirAddr + kv_offset, (void *)output_k_cache.pVirAddr,
sizeof(unsigned short) * _attr.prefill_token_num * _attr.kv_cache_size);
memcpy((unsigned short *)input_decoder_v_cache.pVirAddr + kv_offset, (void *)output_v_cache.pVirAddr,
sizeof(unsigned short) * _attr.prefill_token_num * _attr.kv_cache_size);
memcpy((unsigned short *)input_prefill_k_cache.pVirAddr + kv_offset, (void *)output_k_cache.pVirAddr,
sizeof(unsigned short) * _attr.prefill_token_num * _attr.kv_cache_size);
@@ -1560,9 +1553,6 @@ public:
layer.layer.inference(_attr.prefill_grpid);
auto &input_decoder_k_cache = layer.layer.get_input(decode_grpid, "K_cache");
auto &input_decoder_v_cache = layer.layer.get_input(decode_grpid, "V_cache");
auto &input_prefill_k_cache = layer.layer.get_input(_attr.prefill_grpid, "K_cache");
auto &input_prefill_v_cache = layer.layer.get_input(_attr.prefill_grpid, "V_cache");
@@ -1571,12 +1561,6 @@ public:
int kv_offset = (_attr.precompute_len + p * _attr.prefill_token_num) * _attr.kv_cache_size;
memcpy((unsigned short *)input_decoder_k_cache.pVirAddr + kv_offset, (void *)output_k_cache.pVirAddr,
sizeof(unsigned short) * input_num_token * _attr.kv_cache_size);
memcpy((unsigned short *)input_decoder_v_cache.pVirAddr + kv_offset, (void *)output_v_cache.pVirAddr,
sizeof(unsigned short) * input_num_token * _attr.kv_cache_size);
memcpy((unsigned short *)input_prefill_k_cache.pVirAddr + kv_offset, (void *)output_k_cache.pVirAddr,
sizeof(unsigned short) * input_num_token * _attr.kv_cache_size);
@@ -2265,20 +2249,11 @@ public:
layer.layer.inference(_attr.prefill_grpid);
auto &input_decoder_k_cache = layer.layer.get_input(decode_grpid, "K_cache");
auto &input_decoder_v_cache = layer.layer.get_input(decode_grpid, "V_cache");
auto &output_k_cache = layer.layer.get_output(_attr.prefill_grpid, "K_cache_out");
auto &output_v_cache = layer.layer.get_output(_attr.prefill_grpid, "V_cache_out");
int kv_offset = (_attr.precompute_len + p * _attr.prefill_token_num) * _attr.kv_cache_size;
memcpy((unsigned short *)input_decoder_k_cache.pVirAddr + kv_offset, (void *)output_k_cache.pVirAddr,
sizeof(unsigned short) * input_num_token * _attr.kv_cache_size);
memcpy((unsigned short *)input_decoder_v_cache.pVirAddr + kv_offset, (void *)output_v_cache.pVirAddr,
sizeof(unsigned short) * input_num_token * _attr.kv_cache_size);
for (int gid = _attr.prefill_grpid + 1; gid < prefill_split_num + 1; gid++) {
auto &input_prefill_k_cache = layer.layer.get_input(gid, "K_cache");
@@ -0,0 +1,177 @@
#pragma once
#include <vector>
#include <string>
#include <map>
#include <stdexcept>
namespace ax::legacy {
typedef enum _color_space_e {
axdl_color_space_unknown,
axdl_color_space_nv12,
axdl_color_space_nv21,
axdl_color_space_bgr,
axdl_color_space_rgb,
} ax_color_space_e;
typedef struct _image_t {
unsigned long long int pPhy;
void *pVir;
unsigned int nSize;
unsigned int nWidth;
unsigned int nHeight;
ax_color_space_e eDtype;
union {
int tStride_H, tStride_W, tStride_C;
};
} ax_image_t;
typedef struct {
std::string sName;
unsigned int nIdx;
std::vector<unsigned int> vShape;
int nSize;
unsigned long phyAddr;
void *pVirAddr;
} ax_runner_tensor_t;
class ax_runner_base {
protected:
std::vector<ax_runner_tensor_t> moutput_tensors;
std::vector<ax_runner_tensor_t> minput_tensors;
std::vector<std::vector<ax_runner_tensor_t>> mgroup_output_tensors;
std::vector<std::vector<ax_runner_tensor_t>> mgroup_input_tensors;
std::map<std::string, ax_runner_tensor_t> map_output_tensors;
std::map<std::string, ax_runner_tensor_t> map_input_tensors;
std::map<std::string, std::vector<ax_runner_tensor_t>> map_group_output_tensors;
std::map<std::string, std::vector<ax_runner_tensor_t>> map_group_input_tensors;
public:
virtual int init(const char *model_file, bool use_mmap = false) = 0;
virtual int init(char *model_buffer, size_t model_size) = 0;
virtual void deinit() = 0;
int get_num_inputs()
{
return minput_tensors.size();
};
int get_num_outputs()
{
return moutput_tensors.size();
};
int get_num_input_groups()
{
return mgroup_input_tensors.size();
};
int get_num_output_groups()
{
return mgroup_output_tensors.size();
};
const ax_runner_tensor_t &get_input(int idx)
{
return minput_tensors[idx];
}
const ax_runner_tensor_t *get_inputs_ptr()
{
return minput_tensors.data();
}
const ax_runner_tensor_t &get_input(std::string name)
{
if (map_input_tensors.size() == 0) {
for (size_t i = 0; i < minput_tensors.size(); i++) {
map_input_tensors[minput_tensors[i].sName] = minput_tensors[i];
}
}
if (map_input_tensors.find(name) == map_input_tensors.end()) {
throw std::runtime_error("input tensor not found: " + name);
}
return map_input_tensors[name];
}
const ax_runner_tensor_t &get_input(int grpid, int idx)
{
return mgroup_input_tensors[grpid][idx];
}
const ax_runner_tensor_t *get_inputs_ptr(int grpid)
{
return mgroup_input_tensors[grpid].data();
}
const ax_runner_tensor_t &get_input(int grpid, std::string name)
{
if (map_group_input_tensors.size() == 0) {
for (size_t i = 0; i < mgroup_input_tensors.size(); i++) {
for (size_t j = 0; j < mgroup_input_tensors[i].size(); j++) {
map_group_input_tensors[mgroup_input_tensors[i][j].sName].push_back(mgroup_input_tensors[i][j]);
}
}
}
if (map_group_input_tensors.find(name) == map_group_input_tensors.end()) {
throw std::runtime_error("input tensor not found: " + name);
}
return map_group_input_tensors[name][grpid];
// return map_input_tensors[name];
}
const ax_runner_tensor_t &get_output(int idx)
{
return moutput_tensors[idx];
}
const ax_runner_tensor_t *get_outputs_ptr()
{
return moutput_tensors.data();
}
const ax_runner_tensor_t &get_output(std::string name)
{
if (map_output_tensors.size() == 0) {
for (size_t i = 0; i < moutput_tensors.size(); i++) {
map_output_tensors[moutput_tensors[i].sName] = moutput_tensors[i];
}
}
if (map_output_tensors.find(name) == map_output_tensors.end()) {
throw std::runtime_error("output tensor not found: " + name);
}
return map_output_tensors[name];
}
const ax_runner_tensor_t &get_output(int grpid, int idx)
{
return mgroup_output_tensors[grpid][idx];
}
const ax_runner_tensor_t *get_outputs_ptr(int grpid)
{
return mgroup_output_tensors[grpid].data();
}
const ax_runner_tensor_t &get_output(int grpid, std::string name)
{
if (map_group_output_tensors.size() == 0) {
for (size_t i = 0; i < mgroup_output_tensors.size(); i++) {
for (size_t j = 0; j < mgroup_output_tensors[i].size(); j++) {
map_group_output_tensors[mgroup_output_tensors[i][j].sName].push_back(mgroup_output_tensors[i][j]);
}
}
}
if (map_group_output_tensors.find(name) == map_group_output_tensors.end()) {
throw std::runtime_error("input tensor not found: " + name);
}
return map_group_output_tensors[name][grpid];
}
virtual int inference() = 0;
virtual int inference(int grpid) = 0;
int operator()()
{
return inference();
}
};
// int ax_cmmcpy(unsigned long long int dst, unsigned long long int src, int size);
} // namespace ax::legacy
@@ -0,0 +1,411 @@
#include "ax_model_runner_ax650.hpp"
#include "string.h"
#include "fstream"
#include "memory"
// #include "utilities/file.hpp"
#include <ax_sys_api.h>
#include <ax_ivps_api.h>
#include <ax_engine_api.h>
#include <fcntl.h>
#include "memory_utils.hpp"
#include "sample_log.h"
#define AX_CMM_ALIGN_SIZE 128
namespace ax::legacy {
const char *AX_CMM_SESSION_NAME = "npu";
typedef enum {
AX_ENGINE_ABST_DEFAULT = 0,
AX_ENGINE_ABST_CACHED = 1,
} AX_ENGINE_ALLOC_BUFFER_STRATEGY_T;
typedef std::pair<AX_ENGINE_ALLOC_BUFFER_STRATEGY_T, AX_ENGINE_ALLOC_BUFFER_STRATEGY_T> INPUT_OUTPUT_ALLOC_STRATEGY;
static void print_io_info(AX_ENGINE_IO_INFO_T *io_info)
{
static std::map<AX_ENGINE_DATA_TYPE_T, const char *> data_type = {
{AX_ENGINE_DT_UNKNOWN, "UNKNOWN"},
{AX_ENGINE_DT_UINT8, "UINT8"},
{AX_ENGINE_DT_UINT16, "UINT16"},
{AX_ENGINE_DT_FLOAT32, "FLOAT32"},
{AX_ENGINE_DT_SINT16, "SINT16"},
{AX_ENGINE_DT_SINT8, "SINT8"},
{AX_ENGINE_DT_SINT32, "SINT32"},
{AX_ENGINE_DT_UINT32, "UINT32"},
{AX_ENGINE_DT_FLOAT64, "FLOAT64"},
{AX_ENGINE_DT_UINT10_PACKED, "UINT10_PACKED"},
{AX_ENGINE_DT_UINT12_PACKED, "UINT12_PACKED"},
{AX_ENGINE_DT_UINT14_PACKED, "UINT14_PACKED"},
{AX_ENGINE_DT_UINT16_PACKED, "UINT16_PACKED"},
};
static std::map<AX_ENGINE_COLOR_SPACE_T, const char *> color_type = {
{AX_ENGINE_CS_FEATUREMAP, "FEATUREMAP"},
{AX_ENGINE_CS_RAW8, "RAW8"},
{AX_ENGINE_CS_RAW10, "RAW10"},
{AX_ENGINE_CS_RAW12, "RAW12"},
{AX_ENGINE_CS_RAW14, "RAW14"},
{AX_ENGINE_CS_RAW16, "RAW16"},
{AX_ENGINE_CS_NV12, "NV12"},
{AX_ENGINE_CS_NV21, "NV21"},
{AX_ENGINE_CS_RGB, "RGB"},
{AX_ENGINE_CS_BGR, "BGR"},
{AX_ENGINE_CS_RGBA, "RGBA"},
{AX_ENGINE_CS_GRAY, "GRAY"},
{AX_ENGINE_CS_YUV444, "YUV444"},
};
printf("\ninput size: %d\n", io_info->nInputSize);
for (uint32_t i = 0; i < io_info->nInputSize; ++i) {
// print shape info,like [batchsize x channel x height x width]
auto &info = io_info->pInputs[i];
printf(" name: \e[1;32m%8s", info.pName);
std::string dt = "unknown";
if (data_type.find(info.eDataType) != data_type.end()) {
dt = data_type[info.eDataType];
printf(" \e[1;34m[%s] ", dt.c_str());
} else {
printf(" \e[1;31m[%s] ", dt.c_str());
}
std::string ct = "unknown";
if (info.pExtraMeta && color_type.find(info.pExtraMeta->eColorSpace) != color_type.end()) {
ct = color_type[info.pExtraMeta->eColorSpace];
printf("\e[1;34m[%s]", ct.c_str());
} else {
printf("\e[1;31m[%s]", ct.c_str());
}
printf(" \n \e[1;31m");
for (AX_U8 s = 0; s < info.nShapeSize; s++) {
printf("%d", info.pShape[s]);
if (s != info.nShapeSize - 1) {
printf(" x ");
}
}
printf("\e[0m\n\n");
}
printf("\noutput size: %d\n", io_info->nOutputSize);
for (uint32_t i = 0; i < io_info->nOutputSize; ++i) {
// print shape info,like [batchsize x channel x height x width]
auto &info = io_info->pOutputs[i];
printf(" name: \e[1;32m%8s \e[1;34m[%s]\e[0m\n \e[1;31m", info.pName, data_type[info.eDataType]);
for (AX_U8 s = 0; s < info.nShapeSize; s++) {
printf("%d", info.pShape[s]);
if (s != info.nShapeSize - 1) {
printf(" x ");
}
}
printf("\e[0m\n\n");
}
}
void free_io_index(AX_ENGINE_IO_BUFFER_T *io_buf, int index)
{
for (int i = 0; i < index; ++i) {
AX_ENGINE_IO_BUFFER_T *pBuf = io_buf + i;
AX_SYS_MemFree(pBuf->phyAddr, pBuf->pVirAddr);
}
}
void free_io(AX_ENGINE_IO_T *io)
{
for (size_t j = 0; j < io->nInputSize; ++j) {
AX_ENGINE_IO_BUFFER_T *pBuf = io->pInputs + j;
AX_SYS_MemFree(pBuf->phyAddr, pBuf->pVirAddr);
}
for (size_t j = 0; j < io->nOutputSize; ++j) {
AX_ENGINE_IO_BUFFER_T *pBuf = io->pOutputs + j;
AX_SYS_MemFree(pBuf->phyAddr, pBuf->pVirAddr);
}
delete[] io->pInputs;
delete[] io->pOutputs;
}
static inline int prepare_io(AX_ENGINE_IO_INFO_T *info, AX_ENGINE_IO_T *io_data, INPUT_OUTPUT_ALLOC_STRATEGY strategy)
{
memset(io_data, 0, sizeof(*io_data));
io_data->pInputs = new AX_ENGINE_IO_BUFFER_T[info->nInputSize];
io_data->nInputSize = info->nInputSize;
auto ret = 0;
for (uint i = 0; i < info->nInputSize; ++i) {
auto meta = info->pInputs[i];
auto buffer = &io_data->pInputs[i];
if (strategy.first == AX_ENGINE_ABST_CACHED) {
ret = AX_SYS_MemAllocCached((AX_U64 *)(&buffer->phyAddr), &buffer->pVirAddr, meta.nSize, AX_CMM_ALIGN_SIZE,
(const AX_S8 *)(AX_CMM_SESSION_NAME));
} else {
ret = AX_SYS_MemAlloc((AX_U64 *)(&buffer->phyAddr), &buffer->pVirAddr, meta.nSize, AX_CMM_ALIGN_SIZE,
(const AX_S8 *)(AX_CMM_SESSION_NAME));
}
if (ret != 0) {
free_io_index(io_data->pInputs, i);
fprintf(stderr, "Allocate input{%d} { phy: %p, vir: %p, size: %lu Bytes }. fail \n", i,
(void *)buffer->phyAddr, buffer->pVirAddr, (long)meta.nSize);
return ret;
}
memset(buffer->pVirAddr, 0, meta.nSize);
// fprintf(stderr, "Allocate input{%d} { phy: %p, vir: %p, size: %lu Bytes }. \n", i, (void*)buffer->phyAddr,
// buffer->pVirAddr, (long)meta.nSize);
}
io_data->pOutputs = new AX_ENGINE_IO_BUFFER_T[info->nOutputSize];
io_data->nOutputSize = info->nOutputSize;
for (uint i = 0; i < info->nOutputSize; ++i) {
auto meta = info->pOutputs[i];
auto buffer = &io_data->pOutputs[i];
buffer->nSize = meta.nSize;
if (strategy.second == AX_ENGINE_ABST_CACHED) {
ret = AX_SYS_MemAllocCached((AX_U64 *)(&buffer->phyAddr), &buffer->pVirAddr, meta.nSize, AX_CMM_ALIGN_SIZE,
(const AX_S8 *)(AX_CMM_SESSION_NAME));
} else {
ret = AX_SYS_MemAlloc((AX_U64 *)(&buffer->phyAddr), &buffer->pVirAddr, meta.nSize, AX_CMM_ALIGN_SIZE,
(const AX_S8 *)(AX_CMM_SESSION_NAME));
}
if (ret != 0) {
fprintf(stderr, "Allocate output{%d} { phy: %p, vir: %p, size: %lu Bytes }. fail \n", i,
(void *)buffer->phyAddr, buffer->pVirAddr, (long)meta.nSize);
free_io_index(io_data->pInputs, io_data->nInputSize);
free_io_index(io_data->pOutputs, i);
return ret;
}
memset(buffer->pVirAddr, 0, meta.nSize);
// fprintf(stderr, "Allocate output{%d} { phy: %p, vir: %p, size: %lu Bytes }.\n", i, (void*)buffer->phyAddr,
// buffer->pVirAddr, (long)meta.nSize);
}
return 0;
}
struct ax_joint_runner_ax650_handle_t {
AX_ENGINE_HANDLE handle;
AX_ENGINE_CONTEXT_T context;
std::vector<AX_ENGINE_IO_INFO_T *> io_info;
std::vector<AX_ENGINE_IO_T> io_data;
// int algo_width, algo_height;
// int algo_colorformat;
};
int ax_runner_ax650::sub_init()
{
// 4. create context
int ret = AX_ENGINE_CreateContext(m_handle->handle);
if (0 != ret) {
ALOGE("AX_ENGINE_CreateContext");
return ret;
}
ret = AX_ENGINE_CreateContextV2(m_handle->handle, &m_handle->context);
if (0 != ret) {
ALOGE("AX_ENGINE_CreateContextV2");
return ret;
}
// fprintf(stdout, "Engine creating context is done.\n");
// 5. set io
AX_U32 io_count = 0;
ret = AX_ENGINE_GetGroupIOInfoCount(m_handle->handle, &io_count);
if (0 != ret) {
ALOGE("AX_ENGINE_GetGroupIOInfoCount");
return ret;
}
// ALOGI("io_count=%d", io_count);
m_handle->io_info.resize(io_count);
m_handle->io_data.resize(io_count);
mgroup_input_tensors.resize(io_count);
mgroup_output_tensors.resize(io_count);
// fprintf(stdout, "Engine get io info is done. \n");
// 6. alloc io
if (!_parepare_io) {
for (size_t grpid = 0; grpid < io_count; grpid++) {
AX_ENGINE_IO_INFO_T *io_info = nullptr;
ret = AX_ENGINE_GetGroupIOInfo(m_handle->handle, grpid, &io_info);
if (0 != ret) {
ALOGE("AX_ENGINE_GetIOInfo");
return ret;
}
// print_io_info(io_info);
m_handle->io_info[grpid] = io_info;
ret = prepare_io(m_handle->io_info[grpid], &m_handle->io_data[grpid],
std::make_pair(AX_ENGINE_ABST_DEFAULT, AX_ENGINE_ABST_CACHED));
if (0 != ret) {
ALOGE("prepare_io grpid=%d", grpid);
return ret;
}
}
for (size_t grpid = 0; grpid < io_count; grpid++) {
auto &io_info = m_handle->io_info[grpid];
auto &io_data = m_handle->io_data[grpid];
for (size_t i = 0; i < io_info->nOutputSize; i++) {
ax_runner_tensor_t tensor;
tensor.nIdx = i;
tensor.sName = std::string(io_info->pOutputs[i].pName);
tensor.nSize = io_info->pOutputs[i].nSize;
for (size_t j = 0; j < io_info->pOutputs[i].nShapeSize; j++) {
tensor.vShape.push_back(io_info->pOutputs[i].pShape[j]);
}
tensor.phyAddr = io_data.pOutputs[i].phyAddr;
tensor.pVirAddr = io_data.pOutputs[i].pVirAddr;
mgroup_output_tensors[grpid].push_back(tensor);
}
for (size_t i = 0; i < io_info->nInputSize; i++) {
ax_runner_tensor_t tensor;
tensor.nIdx = i;
tensor.sName = std::string(io_info->pInputs[i].pName);
tensor.nSize = io_info->pInputs[i].nSize;
for (size_t j = 0; j < io_info->pInputs[i].nShapeSize; j++) {
tensor.vShape.push_back(io_info->pInputs[i].pShape[j]);
}
tensor.phyAddr = io_data.pInputs[i].phyAddr;
tensor.pVirAddr = io_data.pInputs[i].pVirAddr;
mgroup_input_tensors[grpid].push_back(tensor);
}
}
moutput_tensors = mgroup_output_tensors[0];
minput_tensors = mgroup_input_tensors[0];
_parepare_io = true;
} else {
}
return ret;
}
int ax_runner_ax650::init(const char *model_file, bool use_mmap)
{
if (use_mmap) {
MMap model_buffer(model_file);
if (!model_buffer.data()) {
ALOGE("mmap");
return -1;
}
auto ret = init((char *)model_buffer.data(), model_buffer.size());
model_buffer.close_file();
return ret;
} else {
char *model_buffer;
size_t len;
if (!read_file(model_file, &model_buffer, &len)) {
ALOGE("read_file");
return -1;
}
auto ret = init(model_buffer, len);
delete[] model_buffer;
return ret;
}
}
int ax_runner_ax650::init(char *model_buffer, size_t model_size)
{
if (!m_handle) {
m_handle = new ax_joint_runner_ax650_handle_t;
}
// static bool b_init = false;
// if (!b_init) {
// // 1. init engine
// AX_ENGINE_NPU_ATTR_T npu_attr;
// memset(&npu_attr, 0, sizeof(npu_attr));
// npu_attr.eHardMode = AX_ENGINE_VIRTUAL_NPU_DISABLE;
// AX_SYS_Init();
// auto ret = AX_ENGINE_Init(&npu_attr);
// if (0 != ret) {
// return ret;
// }
// b_init = true;
// }
// 3. create handle
int ret = AX_ENGINE_CreateHandle(&m_handle->handle, model_buffer, model_size);
if (0 != ret) {
ALOGE("AX_ENGINE_CreateHandle");
return ret;
}
// fprintf(stdout, "Engine creating handle is done.\n");
return sub_init();
}
void ax_runner_ax650::release()
{
if (m_handle && m_handle->handle) {
for (size_t i = 0; i < m_handle->io_data.size(); i++) {
/* code */
free_io(&m_handle->io_data[i]);
}
AX_ENGINE_DestroyHandle(m_handle->handle);
m_handle->handle = nullptr;
}
if (m_handle) {
delete m_handle;
m_handle = nullptr;
}
moutput_tensors.clear();
minput_tensors.clear();
map_input_tensors.clear();
map_output_tensors.clear();
mgroup_output_tensors.clear();
mgroup_input_tensors.clear();
map_group_input_tensors.clear();
map_group_output_tensors.clear();
// AX_ENGINE_Deinit();
}
void ax_runner_ax650::deinit()
{
if (m_handle && m_handle->handle) {
// free_io(&m_handle->io_data);
// mtensors.clear();
// minput_tensors.clear();
// map_input_tensors.clear();
// map_tensors.clear();
AX_ENGINE_DestroyHandle(m_handle->handle);
m_handle->handle = nullptr;
// delete m_handle;
// m_handle = nullptr;
}
// AX_ENGINE_Deinit();
}
int ax_runner_ax650::inference()
{
int ret = AX_ENGINE_RunSync(m_handle->handle, &m_handle->io_data[0]);
for (size_t i = 0; i < get_num_outputs(); i++) {
auto &tensor = get_output(i);
AX_SYS_MinvalidateCache(tensor.phyAddr, tensor.pVirAddr, tensor.nSize);
}
return ret;
}
int ax_runner_ax650::inference(int grpid)
{
int ret = AX_ENGINE_RunGroupIOSync(m_handle->handle, m_handle->context, grpid, &m_handle->io_data[grpid]);
for (size_t i = 0; i < get_num_outputs(); i++) {
auto &tensor = get_output(grpid, i);
AX_SYS_MinvalidateCache(tensor.phyAddr, tensor.pVirAddr, tensor.nSize);
}
return ret;
}
} // namespace ax::legacy
@@ -0,0 +1,25 @@
#pragma once
#include "ax_model_runner.hpp"
namespace ax::legacy {
class ax_runner_ax650 : public ax_runner_base {
protected:
struct ax_joint_runner_ax650_handle_t *m_handle = nullptr;
bool _parepare_io = false;
int sub_init();
public:
int init(const char *model_file, bool use_mmap = false) override;
int init(char *model_buffer, size_t model_size) override;
void release();
void deinit() override;
int inference() override;
int inference(int grpid) override;
};
} // namespace ax::legacy