diff --git a/doc/projects_llm_framework_doc/llm_cosyvoice2.md b/doc/projects_llm_framework_doc/llm_cosyvoice2.md new file mode 100644 index 0000000..7f3a7ab --- /dev/null +++ b/doc/projects_llm_framework_doc/llm_cosyvoice2.md @@ -0,0 +1,223 @@ +# llm_cosy_voice + +使用 npu 加速的文字转语音单元，用于提供文字转语音服务，可使用语音克隆，用于提供多语言转语音服务。 + +## setup + +配置单元工作。 + +发送 json： + +```json +cosy_voice +{ + "request_id": "2", + "work_id": "cosy_voice", + "action": "setup", + "object": "cosy_voice.setup", + "data": { + "model": "CosyVoice2-0.5B-ax650", + "response_format": "file", + "input": "tts.utf-8", + "enoutput": false + } +} +``` + + +- request_id：参考基本数据解释。 +- work_id：配置单元时，为 `cosy_voice`。 +- action：调用的方法为 `setup`。 +- object：传输的数据类型为 `cosy_voice.setup`。 +- model：使用的模型为 `CosyVoice2-0.5B-ax650` 模型。 +- prompt_files：要克隆的音频信息文件。 +- response_format：返回结果为 `sys.pcm`, 系统音频数据，并直接发送到 llm-audio 模块进行播放。返回结果为 `file`, 生成的音频写 wav 文件，可用 `prompt_data` 指定路径或文件名。 +- input：输入的为 `tts.utf-8`,代表的是从用户输入。 +- enoutput：是否起用用户结果输出。 + +响应 json： + +```json +{ + "created": 1761791627, + "data": "None", + "error": { + "code": 0, + "message": "" + }, + "object": "None", + "request_id": "2", + "work_id": "cosy_voice.1000" +} +``` + +- created：消息创建时间，unix 时间。 +- work_id：返回成功创建的 work_id 单元。 + +## inference + +### 流式输入 + +```json +{ + "request_id": "2", + "work_id": "cosy_voice.1000", + "action": "inference", + "object": "cosy_voice.utf-8.stream", + "data": { + "delta": "今天天气真好！", + "index": 0, + "finish": true + } +} +``` +- object：传输的数据类型为 `cosy_voice.utf-8.stream` 代表的是从用户 utf-8 的流式输入 +- delta：流式输入的分段数据 +- index：流式输入的分段索引 +- finish:流式输入是否完成的标志位 + +### 非流式输入 + +```json +{ + "request_id": "2", + "work_id": "cosy_voice.1000", + "action": "inference", + "object": "cosy_voice.utf-8", + "data": "今天天气真好！" +} +``` +- object：传输的数据类型为 `cosy_voice.utf-8` 代表的是从用户 utf-8 的非流式输入 +- data：非流式输入的数据 + +## pause + +暂停单元工作。 + +发送 json： + +```json +{ + "request_id": "5", + "work_id": "cosy_voice.1000", + "action": "pause" +} +``` + +响应 json： + +```json +{ + "created": 1761791706, + "data": "None", + "error": { + "code": 0, + "message": "" + }, + "object": "None", + "request_id": "5", + "work_id": "cosy_voice.1000" +} +``` + +error::code 为 0 表示执行成功。 + +## exit + +单元退出。 + +发送 json： + +```json +{ + "request_id": "7", + "work_id": "cosy_voice.1000", + "action": "exit" +} +``` + +响应 json： + +```json +{ + "created": 1761791854, + "data": "None", + "error": { + "code": 0, + "message": "" + }, + "object": "None", + "request_id": "7", + "work_id": "cosy_voice.1000" +} +``` + +error::code 为 0 表示执行成功。 + +## taskinfo + +获取任务列表。 + +发送 json： + +```json +{ + "request_id": "2", + "work_id": "cosy_voice", + "action": "taskinfo" +} +``` + +响应 json： + +```json +{ + "created": 1761791739, + "data": [ + "cosy_voice.1000" + ], + "error": { + "code": 0, + "message": "" + }, + "object": "llm.tasklist", + "request_id": "2", + "work_id": "cosy_voice" +} +``` + +获取任务运行参数。 + +```json +{ + "request_id": "2", + "work_id": "cosy_voice.1000", + "action": "taskinfo" +} +``` + +响应 json： + +```json +{ + "created": 1761791761, + "data": { + "enoutput": false, + "inputs": [ + "tts.utf-8" + ], + "model": "CosyVoice2-0.5B-ax650", + "response_format": "sys.pcm" + }, + "error": { + "code": 0, + "message": "" + }, + "object": "cosy_voice.taskinfo", + "request_id": "2", + "work_id": "cosy_voice.1000" +} +``` + +> **注意：work_id 是按照单元的初始化注册顺序增加的，并不是固定的索引值。** +> **同类型单元不能配置多个单元同时工作，否则会产生未知错误。例如 tts 和 melo tts 不能同时拍起用工作。** diff --git a/ext_components/tokenizer/Kconfig b/ext_components/tokenizer/Kconfig new file mode 100644 index 0000000..1cf1138 --- /dev/null +++ b/ext_components/tokenizer/Kconfig @@ -0,0 +1,5 @@ +menuconfig AX_TOKENIZER_ENABLED + bool "Enable tokenizer support" + default n + help + enable tokenizer support \ No newline at end of file diff --git a/ext_components/tokenizer/SConstruct b/ext_components/tokenizer/SConstruct new file mode 100644 index 0000000..e393f25 --- /dev/null +++ b/ext_components/tokenizer/SConstruct @@ -0,0 +1,52 @@ +# component2/SConscript +Import("env") +import os +from pathlib import Path + +with open(env["PROJECT_TOOL_S"]) as f: + exec(f.read()) + +_SDK_PATH = os.path.normpath( + os.environ.get("SDK_PATH", str(Path(os.getcwd()) / ".." / "..")) +) + +env["GIT_REPO_LISTS"]["tokenizer"] = { + "url": "https://github.com/ZHEQIUSHUI/tokenizer.git", + "commit": "83f41d4b5b9a135c167d44fcdf2a0c56ebacca6d", + "path": str(Path(_SDK_PATH) / "github_source" / "tokenizer"), +} + +if "CONFIG_AX_TOKENIZER_ENABLED" in os.environ: + check_component("tokenizer") + SRCS = [] + INCLUDE = [] + PRIVATE_INCLUDE = [] + REQUIREMENTS = [] + STATIC_LIB = [] + DYNAMIC_LIB = [] + DEFINITIONS = [] + DEFINITIONS_PRIVATE = [] + LDFLAGS = [] + LINK_SEARCH_PATH = [] + + INCLUDE += [ + os.path.join(env["GIT_REPO_LISTS"]["tokenizer"]["path"], "include"), + ] + print("AX-TOKENIZER INCLUDE:", INCLUDE) + + env["COMPONENTS"].append( + { + "target": os.path.basename(env["component_dir"]), + "SRCS": SRCS, + "INCLUDE": INCLUDE, + "PRIVATE_INCLUDE": PRIVATE_INCLUDE, + "REQUIREMENTS": REQUIREMENTS, + "STATIC_LIB": STATIC_LIB, + "DYNAMIC_LIB": DYNAMIC_LIB, + "DEFINITIONS": DEFINITIONS, + "DEFINITIONS_PRIVATE": DEFINITIONS_PRIVATE, + "LDFLAGS": LDFLAGS, + "LINK_SEARCH_PATH": LINK_SEARCH_PATH, + "REGISTER": "static", + } + ) \ No newline at end of file diff --git a/projects/llm_framework/main_cosy_voice/src/runner/LLM.hpp b/projects/llm_framework/main_cosy_voice/src/runner/LLM.hpp index 52994ca..bb18782 100644 --- a/projects/llm_framework/main_cosy_voice/src/runner/LLM.hpp +++ b/projects/llm_framework/main_cosy_voice/src/runner/LLM.hpp @@ -246,10 +246,10 @@ public: void Deinit() { for (int i = 0; i < _attr.axmodel_num; i++) { - llama_layers[i].layer.release(); + llama_layers[i].layer.deinit(); } - llama_post.release(); - llm_decoder.release(); + llama_post.deinit(); + llm_decoder.deinit(); embed_selector.Deinit(); llm_embed_selector.Deinit(); speech_embed_selector.Deinit(); diff --git a/projects/llm_framework/main_cosy_voice/src/runner/Token2wav.hpp b/projects/llm_framework/main_cosy_voice/src/runner/Token2wav.hpp index 0eb0450..e1daffa 100644 --- a/projects/llm_framework/main_cosy_voice/src/runner/Token2wav.hpp +++ b/projects/llm_framework/main_cosy_voice/src/runner/Token2wav.hpp @@ -145,15 +145,15 @@ public: void Deinit() { - flow_encoder_28.release(); - flow_encoder_53.release(); - flow_encoder_78.release(); - flow_encoder_50_final.release(); - flow_estimator_200.release(); - flow_estimator_250.release(); - flow_estimator_300.release(); - hift_p2_50_first.release(); - hift_p2_58.release(); + flow_encoder_28.deinit(); + flow_encoder_53.deinit(); + flow_encoder_78.deinit(); + flow_encoder_50_final.deinit(); + flow_estimator_200.deinit(); + flow_estimator_250.deinit(); + flow_estimator_300.deinit(); + hift_p2_50_first.deinit(); + hift_p2_58.deinit(); flow_embed_selector.Deinit(); } diff --git a/projects/llm_framework/main_cosy_voice/src/runner/ax_model_runner/ax_model_runner.hpp b/projects/llm_framework/main_cosy_voice/src/runner/ax_model_runner/ax_model_runner.hpp index c9ed4fa..c4bcb2a 100644 --- a/projects/llm_framework/main_cosy_voice/src/runner/ax_model_runner/ax_model_runner.hpp +++ b/projects/llm_framework/main_cosy_voice/src/runner/ax_model_runner/ax_model_runner.hpp @@ -4,31 +4,7 @@ #include #include -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 -{ +typedef struct { std::string sName; unsigned int nIdx; std::vector vShape; @@ -37,8 +13,7 @@ typedef struct void *pVirAddr; } ax_runner_tensor_t; -class ax_runner_base -{ +class ax_runner_base { protected: std::vector moutput_tensors; std::vector minput_tensors; @@ -52,106 +27,124 @@ protected: std::map> map_group_output_tensors; std::map> map_group_input_tensors; + void build_tensor_maps() + { + map_input_tensors.clear(); + for (const auto &t : minput_tensors) map_input_tensors[t.sName] = t; + + map_output_tensors.clear(); + for (const auto &t : moutput_tensors) map_output_tensors[t.sName] = t; + + map_group_input_tensors.clear(); + for (const auto &grp : mgroup_input_tensors) { + for (const auto &t : grp) map_group_input_tensors[t.sName].push_back(t); + } + + map_group_output_tensors.clear(); + for (const auto &grp : mgroup_output_tensors) { + for (const auto &t : grp) map_group_output_tensors[t.sName].push_back(t); + } + } + public: + virtual ~ax_runner_base() + { + } + virtual int init(const char *model_file, bool use_mmap = false) = 0; - virtual int init(char *model_buffer, size_t model_size) = 0; + virtual int init(char *model_buffer, size_t model_size) = 0; + virtual void deinit() = 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) + int get_num_inputs() { - 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 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(); + }; - return map_input_tensors[name]; + 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(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) + const ax_runner_tensor_t &get_input(const 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]; + auto it = map_input_tensors.find(name); + if (it == map_input_tensors.end()) throw std::runtime_error("input tensor not found: " + name); + return it->second; } - 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) + const ax_runner_tensor_t &get_input(int grpid, int idx) { - 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]; + 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_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) + const ax_runner_tensor_t &get_input(int grpid, const 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]; + auto it = map_group_input_tensors.find(name); + if (it == map_group_input_tensors.end()) throw std::runtime_error("input tensor not found: " + name); + if (grpid < 0 || grpid >= (int)it->second.size()) + throw std::runtime_error("group id out of range for: " + name); + return it->second[grpid]; } - virtual int inference() = 0; + 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(const std::string &name) + { + auto it = map_output_tensors.find(name); + if (it == map_output_tensors.end()) throw std::runtime_error("output tensor not found: " + name); + return it->second; + } + + 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, const std::string &name) + { + auto it = map_group_output_tensors.find(name); + if (it == map_group_output_tensors.end()) throw std::runtime_error("output tensor not found: " + name); + if (grpid < 0 || grpid >= (int)it->second.size()) + throw std::runtime_error("group id out of range for: " + name); + return it->second[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); +}; \ No newline at end of file diff --git a/projects/llm_framework/main_cosy_voice/src/runner/ax_model_runner/ax_model_runner_ax650.cpp b/projects/llm_framework/main_cosy_voice/src/runner/ax_model_runner/ax_model_runner_ax650.cpp index a974f0c..363c3d5 100644 --- a/projects/llm_framework/main_cosy_voice/src/runner/ax_model_runner/ax_model_runner_ax650.cpp +++ b/projects/llm_framework/main_cosy_voice/src/runner/ax_model_runner/ax_model_runner_ax650.cpp @@ -1,8 +1,9 @@ #include "ax_model_runner_ax650.hpp" -#include "string.h" -#include "fstream" -#include "memory" -// #include "utilities/file.hpp" +#include +#include +#include +#include +#include #include #include #include @@ -11,332 +12,211 @@ #include "sample_log.h" #define AX_CMM_ALIGN_SIZE 128 - const char *AX_CMM_SESSION_NAME = "npu"; -typedef enum -{ +typedef enum { AX_ENGINE_ABST_DEFAULT = 0, - AX_ENGINE_ABST_CACHED = 1, + AX_ENGINE_ABST_CACHED = 1, } AX_ENGINE_ALLOC_BUFFER_STRATEGY_T; -typedef std::pair INPUT_OUTPUT_ALLOC_STRATEGY; +struct ax_runner_ax650_handle_t { + AX_ENGINE_HANDLE handle = nullptr; + AX_ENGINE_CONTEXT_T context = 0; + std::vector io_info; + std::vector io_data; +}; -static void print_io_info(AX_ENGINE_IO_INFO_T *io_info) -{ - static std::map 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 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) +static int prepare_io_struct_only(AX_ENGINE_IO_INFO_T *info, AX_ENGINE_IO_T *io_data) { memset(io_data, 0, sizeof(*io_data)); - io_data->pInputs = new AX_ENGINE_IO_BUFFER_T[info->nInputSize]; + io_data->pInputs = new AX_ENGINE_IO_BUFFER_T[info->nInputSize]; io_data->nInputSize = info->nInputSize; + memset(io_data->pInputs, 0, sizeof(AX_ENGINE_IO_BUFFER_T) * 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->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); - } + memset(io_data->pOutputs, 0, sizeof(AX_ENGINE_IO_BUFFER_T) * info->nOutputSize); + + for (uint i = 0; i < info->nInputSize; ++i) io_data->pInputs[i].nSize = info->pInputs[i].nSize; + for (uint i = 0; i < info->nOutputSize; ++i) io_data->pOutputs[i].nSize = info->pOutputs[i].nSize; return 0; } -struct ax_joint_runner_ax650_handle_t +static int prepare_io_with_alloc( + AX_ENGINE_IO_INFO_T *info, AX_ENGINE_IO_T *io_data, + std::pair strategy, + std::vector skip_alloc_names = {}) { - AX_ENGINE_HANDLE handle; - AX_ENGINE_CONTEXT_T context; - std::vector io_info; - std::vector io_data; + int ret = prepare_io_struct_only(info, io_data); + if (ret != 0) return ret; - // int algo_width, algo_height; - // int algo_colorformat; -}; + for (uint i = 0; i < info->nInputSize; ++i) { + auto &buffer = io_data->pInputs[i]; + if (std::find(skip_alloc_names.begin(), skip_alloc_names.end(), info->pInputs[i].pName) != + skip_alloc_names.end()) { + continue; + } + if (strategy.first == AX_ENGINE_ABST_CACHED) { + ret = AX_SYS_MemAllocCached((AX_U64 *)(&buffer.phyAddr), &buffer.pVirAddr, buffer.nSize, AX_CMM_ALIGN_SIZE, + (const AX_S8 *)(AX_CMM_SESSION_NAME)); + } else { + ret = AX_SYS_MemAlloc((AX_U64 *)(&buffer.phyAddr), &buffer.pVirAddr, buffer.nSize, AX_CMM_ALIGN_SIZE, + (const AX_S8 *)(AX_CMM_SESSION_NAME)); + } + if (ret != 0) { + ALOGE("Alloc input[%d] failed", i); + return ret; + } + memset(buffer.pVirAddr, 0, buffer.nSize); + } + + for (uint i = 0; i < info->nOutputSize; ++i) { + auto &buffer = io_data->pOutputs[i]; + if (std::find(skip_alloc_names.begin(), skip_alloc_names.end(), info->pOutputs[i].pName) != + skip_alloc_names.end()) { + continue; + } + if (strategy.second == AX_ENGINE_ABST_CACHED) { + ret = AX_SYS_MemAllocCached((AX_U64 *)(&buffer.phyAddr), &buffer.pVirAddr, buffer.nSize, AX_CMM_ALIGN_SIZE, + (const AX_S8 *)(AX_CMM_SESSION_NAME)); + } else { + ret = AX_SYS_MemAlloc((AX_U64 *)(&buffer.phyAddr), &buffer.pVirAddr, buffer.nSize, AX_CMM_ALIGN_SIZE, + (const AX_S8 *)(AX_CMM_SESSION_NAME)); + } + if (ret != 0) { + ALOGE("Alloc output[%d] failed", i); + return ret; + } + memset(buffer.pVirAddr, 0, buffer.nSize); + } + return 0; +} int ax_runner_ax650::sub_init() { - // 4. create context + if (!m_handle) return -1; + int ret = AX_ENGINE_CreateContext(m_handle->handle); - if (0 != ret) - { - ALOGE("AX_ENGINE_CreateContext"); - return ret; - } + if (ret != 0) 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"); + if (ret != 0) return ret; - // 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); + ret = AX_ENGINE_GetGroupIOInfoCount(m_handle->handle, &io_count); + if (ret != 0) return ret; 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"); + std::vector skip_alloc_names = {"K_cache", "V_cache"}; + 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 (ret != 0) return ret; + m_handle->io_info[grpid] = io_info; - // 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); + if (grpid == 0) { + ret = prepare_io_with_alloc(io_info, &m_handle->io_data[grpid], + {AX_ENGINE_ABST_DEFAULT, AX_ENGINE_ABST_CACHED}); + } else if (grpid == io_count - 1) { + ret = prepare_io_with_alloc(io_info, &m_handle->io_data[grpid], + {AX_ENGINE_ABST_DEFAULT, AX_ENGINE_ABST_CACHED}, skip_alloc_names); + } else { + ret = prepare_io_struct_only(io_info, &m_handle->io_data[grpid]); + } + if (ret != 0) return ret; + } - 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; + if (io_count > 2) { + auto &first_io_data = m_handle->io_data[0]; + auto &first_io_info = m_handle->io_info[0]; + auto &last_io_data = m_handle->io_data[io_count - 1]; + auto &last_io_info = m_handle->io_info[io_count - 1]; + for (uint i = 0; i < last_io_data.nInputSize; ++i) { + if (std::find(skip_alloc_names.begin(), skip_alloc_names.end(), last_io_info->pInputs[i].pName) != + skip_alloc_names.end()) { + for (uint j = 0; j < first_io_data.nInputSize; ++j) { + if (first_io_info->pInputs[j].pName == last_io_info->pInputs[i].pName) { + last_io_data.pInputs[i].phyAddr = first_io_data.pInputs[j].phyAddr; + last_io_data.pInputs[i].pVirAddr = first_io_data.pInputs[j].pVirAddr; + } + } } } - for (size_t grpid = 0; grpid < io_count; grpid++) - { + for (size_t grpid = 1; grpid < io_count - 1; 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); + + size_t min_inputs = std::min(io_info->nInputSize, last_io_data.nInputSize); + for (size_t i = 0; i < min_inputs; i++) { + io_data.pInputs[i].phyAddr = last_io_data.pInputs[i].phyAddr; + io_data.pInputs[i].pVirAddr = last_io_data.pInputs[i].pVirAddr; } - 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); + size_t min_outputs = std::min(io_info->nOutputSize, last_io_data.nOutputSize); + for (size_t i = 0; i < min_outputs; i++) { + io_data.pOutputs[i].phyAddr = last_io_data.pOutputs[i].phyAddr; + io_data.pOutputs[i].pVirAddr = last_io_data.pOutputs[i].pVirAddr; } } - - moutput_tensors = mgroup_output_tensors[0]; - minput_tensors = mgroup_input_tensors[0]; - - _parepare_io = true; - } - else - { } - 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 = io_info->pOutputs[i].pName ? std::string(io_info->pOutputs[i].pName) : ""; + tensor.nSize = io_info->pOutputs[i].nSize; + tensor.phyAddr = io_data.pOutputs[i].phyAddr; + tensor.pVirAddr = io_data.pOutputs[i].pVirAddr; + for (size_t j = 0; j < io_info->pOutputs[i].nShapeSize; j++) { + tensor.vShape.push_back(io_info->pOutputs[i].pShape[j]); + } + 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 = io_info->pInputs[i].pName ? std::string(io_info->pInputs[i].pName) : ""; + tensor.nSize = io_info->pInputs[i].nSize; + tensor.phyAddr = io_data.pInputs[i].phyAddr; + tensor.pVirAddr = io_data.pInputs[i].pVirAddr; + for (size_t j = 0; j < io_info->pInputs[i].nShapeSize; j++) { + tensor.vShape.push_back(io_info->pInputs[i].pShape[j]); + } + mgroup_input_tensors[grpid].push_back(tensor); + } + } + + if (!mgroup_output_tensors.empty()) moutput_tensors = mgroup_output_tensors[0]; + if (!mgroup_input_tensors.empty()) minput_tensors = mgroup_input_tensors[0]; + + build_tensor_maps(); + + return 0; } int ax_runner_ax650::init(const char *model_file, bool use_mmap) { - if (use_mmap) - { + if (use_mmap) { MMap model_buffer(model_file); - if (!model_buffer.data()) - { - ALOGE("mmap"); - return -1; - } + if (!model_buffer.data()) 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; - } + } else { + char *model_buffer = nullptr; + size_t len = 0; + if (!read_file(model_file, &model_buffer, &len)) return -1; auto ret = init(model_buffer, len); delete[] model_buffer; return ret; @@ -345,96 +225,84 @@ int ax_runner_ax650::init(const char *model_file, bool use_mmap) 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"); + if (m_handle) deinit(); + m_handle = new ax_runner_ax650_handle_t; + int ret = AX_ENGINE_CreateHandle(&m_handle->handle, model_buffer, model_size); + if (0 != ret) { + ALOGE("AX_ENGINE_CreateHandle failed: 0x%x", ret); + delete m_handle; + m_handle = nullptr; return ret; } - // fprintf(stdout, "Engine creating handle is done.\n"); - return sub_init(); } -void ax_runner_ax650::release() +void ax_runner_ax650::deinit() { - 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]); + if (!m_handle) return; + + std::unordered_set freed_phy_addrs; + + for (size_t g = 0; g < m_handle->io_data.size(); ++g) { + auto &io = m_handle->io_data[g]; + + if (io.pInputs) { + for (size_t j = 0; j < io.nInputSize; ++j) { + AX_ENGINE_IO_BUFFER_T *pBuf = io.pInputs + j; + if (pBuf->phyAddr != 0) { + if (freed_phy_addrs.find(pBuf->phyAddr) == freed_phy_addrs.end()) { + AX_SYS_MemFree(pBuf->phyAddr, pBuf->pVirAddr); + freed_phy_addrs.insert(pBuf->phyAddr); + } + } + } + delete[] io.pInputs; + io.pInputs = nullptr; } - AX_ENGINE_DestroyHandle(m_handle->handle); - m_handle->handle = nullptr; + if (io.pOutputs) { + for (size_t j = 0; j < io.nOutputSize; ++j) { + AX_ENGINE_IO_BUFFER_T *pBuf = io.pOutputs + j; + if (pBuf->phyAddr != 0) { + if (freed_phy_addrs.find(pBuf->phyAddr) == freed_phy_addrs.end()) { + AX_SYS_MemFree(pBuf->phyAddr, pBuf->pVirAddr); + freed_phy_addrs.insert(pBuf->phyAddr); + } + } + } + delete[] io.pOutputs; + io.pOutputs = nullptr; + } } - if (m_handle) - { - delete m_handle; - m_handle = nullptr; + if (m_handle->handle) { + AX_ENGINE_DestroyHandle(m_handle->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() { + if (!m_handle) return -1; + for (size_t i = 0; i < get_num_inputs(); i++) { + auto &tensor = get_input(i); + AX_SYS_MflushCache(tensor.phyAddr, tensor.pVirAddr, tensor.nSize); + } + int ret = AX_ENGINE_RunSync(m_handle->handle, &m_handle->io_data[0]); - for (size_t i = 0; i < get_num_outputs(); i++) - { + + for (size_t i = 0; i < get_num_outputs(); i++) { auto &tensor = get_output(i); AX_SYS_MinvalidateCache(tensor.phyAddr, tensor.pVirAddr, tensor.nSize); } @@ -443,11 +311,18 @@ int ax_runner_ax650::inference() int ax_runner_ax650::inference(int grpid) { + if (!m_handle) return -1; + if (grpid < 0 || grpid >= (int)m_handle->io_data.size()) return -1; + + for (size_t i = 0; i < mgroup_input_tensors[grpid].size(); i++) { + auto &tensor = mgroup_input_tensors[grpid][i]; + AX_SYS_MflushCache(tensor.phyAddr, tensor.pVirAddr, tensor.nSize); + } + 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); + for (size_t i = 0; i < mgroup_output_tensors[grpid].size(); i++) { + auto &tensor = mgroup_output_tensors[grpid][i]; AX_SYS_MinvalidateCache(tensor.phyAddr, tensor.pVirAddr, tensor.nSize); } return ret; diff --git a/projects/llm_framework/main_cosy_voice/src/runner/ax_model_runner/ax_model_runner_ax650.hpp b/projects/llm_framework/main_cosy_voice/src/runner/ax_model_runner/ax_model_runner_ax650.hpp index 351fb04..62b6538 100644 --- a/projects/llm_framework/main_cosy_voice/src/runner/ax_model_runner/ax_model_runner_ax650.hpp +++ b/projects/llm_framework/main_cosy_voice/src/runner/ax_model_runner/ax_model_runner_ax650.hpp @@ -1,20 +1,23 @@ #pragma once #include "ax_model_runner.hpp" -class ax_runner_ax650 : public ax_runner_base -{ +struct ax_runner_ax650_handle_t; + +class ax_runner_ax650 : public ax_runner_base { protected: - struct ax_joint_runner_ax650_handle_t *m_handle = nullptr; - - bool _parepare_io = false; - + struct ax_runner_ax650_handle_t *m_handle = nullptr; int sub_init(); public: + ax_runner_ax650() = default; + virtual ~ax_runner_ax650() + { + deinit(); + } + 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; diff --git a/projects/llm_framework/main_llm/src/runner/LLM.hpp b/projects/llm_framework/main_llm/src/runner/LLM.hpp index 9de1f1f..a5fb872 100644 --- a/projects/llm_framework/main_llm/src/runner/LLM.hpp +++ b/projects/llm_framework/main_llm/src/runner/LLM.hpp @@ -243,9 +243,9 @@ public: void Deinit() { for (int i = 0; i < _attr.axmodel_num; i++) { - llama_layers[i].layer.release(); + llama_layers[i].layer.deinit(); } - llama_post.release(); + llama_post.deinit(); embed_selector.Deinit(); } @@ -686,9 +686,9 @@ public: void Deinit() { for (int i = 0; i < _attr.axmodel_num; i++) { - llama_layers[i].layer.release(); + llama_layers[i].layer.deinit(); } - llama_post.release(); + llama_post.deinit(); embed_selector.Deinit(); } diff --git a/projects/llm_framework/main_llm/src/runner/ax_model_runner/ax_model_runner.hpp b/projects/llm_framework/main_llm/src/runner/ax_model_runner/ax_model_runner.hpp index c9ed4fa..c4bcb2a 100644 --- a/projects/llm_framework/main_llm/src/runner/ax_model_runner/ax_model_runner.hpp +++ b/projects/llm_framework/main_llm/src/runner/ax_model_runner/ax_model_runner.hpp @@ -4,31 +4,7 @@ #include #include -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 -{ +typedef struct { std::string sName; unsigned int nIdx; std::vector vShape; @@ -37,8 +13,7 @@ typedef struct void *pVirAddr; } ax_runner_tensor_t; -class ax_runner_base -{ +class ax_runner_base { protected: std::vector moutput_tensors; std::vector minput_tensors; @@ -52,106 +27,124 @@ protected: std::map> map_group_output_tensors; std::map> map_group_input_tensors; + void build_tensor_maps() + { + map_input_tensors.clear(); + for (const auto &t : minput_tensors) map_input_tensors[t.sName] = t; + + map_output_tensors.clear(); + for (const auto &t : moutput_tensors) map_output_tensors[t.sName] = t; + + map_group_input_tensors.clear(); + for (const auto &grp : mgroup_input_tensors) { + for (const auto &t : grp) map_group_input_tensors[t.sName].push_back(t); + } + + map_group_output_tensors.clear(); + for (const auto &grp : mgroup_output_tensors) { + for (const auto &t : grp) map_group_output_tensors[t.sName].push_back(t); + } + } + public: + virtual ~ax_runner_base() + { + } + virtual int init(const char *model_file, bool use_mmap = false) = 0; - virtual int init(char *model_buffer, size_t model_size) = 0; + virtual int init(char *model_buffer, size_t model_size) = 0; + virtual void deinit() = 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) + int get_num_inputs() { - 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 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(); + }; - return map_input_tensors[name]; + 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(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) + const ax_runner_tensor_t &get_input(const 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]; + auto it = map_input_tensors.find(name); + if (it == map_input_tensors.end()) throw std::runtime_error("input tensor not found: " + name); + return it->second; } - 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) + const ax_runner_tensor_t &get_input(int grpid, int idx) { - 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]; + 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_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) + const ax_runner_tensor_t &get_input(int grpid, const 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]; + auto it = map_group_input_tensors.find(name); + if (it == map_group_input_tensors.end()) throw std::runtime_error("input tensor not found: " + name); + if (grpid < 0 || grpid >= (int)it->second.size()) + throw std::runtime_error("group id out of range for: " + name); + return it->second[grpid]; } - virtual int inference() = 0; + 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(const std::string &name) + { + auto it = map_output_tensors.find(name); + if (it == map_output_tensors.end()) throw std::runtime_error("output tensor not found: " + name); + return it->second; + } + + 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, const std::string &name) + { + auto it = map_group_output_tensors.find(name); + if (it == map_group_output_tensors.end()) throw std::runtime_error("output tensor not found: " + name); + if (grpid < 0 || grpid >= (int)it->second.size()) + throw std::runtime_error("group id out of range for: " + name); + return it->second[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); +}; \ No newline at end of file diff --git a/projects/llm_framework/main_llm/src/runner/ax_model_runner/ax_model_runner_ax650.cpp b/projects/llm_framework/main_llm/src/runner/ax_model_runner/ax_model_runner_ax650.cpp index 374261b..363c3d5 100644 --- a/projects/llm_framework/main_llm/src/runner/ax_model_runner/ax_model_runner_ax650.cpp +++ b/projects/llm_framework/main_llm/src/runner/ax_model_runner/ax_model_runner_ax650.cpp @@ -1,8 +1,9 @@ #include "ax_model_runner_ax650.hpp" -#include "string.h" -#include "fstream" -#include "memory" -// #include "utilities/file.hpp" +#include +#include +#include +#include +#include #include #include #include @@ -11,332 +12,211 @@ #include "sample_log.h" #define AX_CMM_ALIGN_SIZE 128 - const char *AX_CMM_SESSION_NAME = "npu"; -typedef enum -{ +typedef enum { AX_ENGINE_ABST_DEFAULT = 0, - AX_ENGINE_ABST_CACHED = 1, + AX_ENGINE_ABST_CACHED = 1, } AX_ENGINE_ALLOC_BUFFER_STRATEGY_T; -typedef std::pair INPUT_OUTPUT_ALLOC_STRATEGY; +struct ax_runner_ax650_handle_t { + AX_ENGINE_HANDLE handle = nullptr; + AX_ENGINE_CONTEXT_T context = 0; + std::vector io_info; + std::vector io_data; +}; -static void print_io_info(AX_ENGINE_IO_INFO_T *io_info) -{ - static std::map 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 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) +static int prepare_io_struct_only(AX_ENGINE_IO_INFO_T *info, AX_ENGINE_IO_T *io_data) { memset(io_data, 0, sizeof(*io_data)); - io_data->pInputs = new AX_ENGINE_IO_BUFFER_T[info->nInputSize]; + io_data->pInputs = new AX_ENGINE_IO_BUFFER_T[info->nInputSize]; io_data->nInputSize = info->nInputSize; + memset(io_data->pInputs, 0, sizeof(AX_ENGINE_IO_BUFFER_T) * 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->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); - } + memset(io_data->pOutputs, 0, sizeof(AX_ENGINE_IO_BUFFER_T) * info->nOutputSize); + + for (uint i = 0; i < info->nInputSize; ++i) io_data->pInputs[i].nSize = info->pInputs[i].nSize; + for (uint i = 0; i < info->nOutputSize; ++i) io_data->pOutputs[i].nSize = info->pOutputs[i].nSize; return 0; } -struct ax_joint_runner_ax650_handle_t +static int prepare_io_with_alloc( + AX_ENGINE_IO_INFO_T *info, AX_ENGINE_IO_T *io_data, + std::pair strategy, + std::vector skip_alloc_names = {}) { - AX_ENGINE_HANDLE handle; - AX_ENGINE_CONTEXT_T context; - std::vector io_info; - std::vector io_data; + int ret = prepare_io_struct_only(info, io_data); + if (ret != 0) return ret; - // int algo_width, algo_height; - // int algo_colorformat; -}; + for (uint i = 0; i < info->nInputSize; ++i) { + auto &buffer = io_data->pInputs[i]; + if (std::find(skip_alloc_names.begin(), skip_alloc_names.end(), info->pInputs[i].pName) != + skip_alloc_names.end()) { + continue; + } + if (strategy.first == AX_ENGINE_ABST_CACHED) { + ret = AX_SYS_MemAllocCached((AX_U64 *)(&buffer.phyAddr), &buffer.pVirAddr, buffer.nSize, AX_CMM_ALIGN_SIZE, + (const AX_S8 *)(AX_CMM_SESSION_NAME)); + } else { + ret = AX_SYS_MemAlloc((AX_U64 *)(&buffer.phyAddr), &buffer.pVirAddr, buffer.nSize, AX_CMM_ALIGN_SIZE, + (const AX_S8 *)(AX_CMM_SESSION_NAME)); + } + if (ret != 0) { + ALOGE("Alloc input[%d] failed", i); + return ret; + } + memset(buffer.pVirAddr, 0, buffer.nSize); + } + + for (uint i = 0; i < info->nOutputSize; ++i) { + auto &buffer = io_data->pOutputs[i]; + if (std::find(skip_alloc_names.begin(), skip_alloc_names.end(), info->pOutputs[i].pName) != + skip_alloc_names.end()) { + continue; + } + if (strategy.second == AX_ENGINE_ABST_CACHED) { + ret = AX_SYS_MemAllocCached((AX_U64 *)(&buffer.phyAddr), &buffer.pVirAddr, buffer.nSize, AX_CMM_ALIGN_SIZE, + (const AX_S8 *)(AX_CMM_SESSION_NAME)); + } else { + ret = AX_SYS_MemAlloc((AX_U64 *)(&buffer.phyAddr), &buffer.pVirAddr, buffer.nSize, AX_CMM_ALIGN_SIZE, + (const AX_S8 *)(AX_CMM_SESSION_NAME)); + } + if (ret != 0) { + ALOGE("Alloc output[%d] failed", i); + return ret; + } + memset(buffer.pVirAddr, 0, buffer.nSize); + } + return 0; +} int ax_runner_ax650::sub_init() { - // 4. create context + if (!m_handle) return -1; + int ret = AX_ENGINE_CreateContext(m_handle->handle); - if (0 != ret) - { - ALOGE("AX_ENGINE_CreateContext"); - return ret; - } + if (ret != 0) 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"); + if (ret != 0) return ret; - // 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); + ret = AX_ENGINE_GetGroupIOInfoCount(m_handle->handle, &io_count); + if (ret != 0) return ret; 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"); + std::vector skip_alloc_names = {"K_cache", "V_cache"}; + 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 (ret != 0) return ret; + m_handle->io_info[grpid] = io_info; - // 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); + if (grpid == 0) { + ret = prepare_io_with_alloc(io_info, &m_handle->io_data[grpid], + {AX_ENGINE_ABST_DEFAULT, AX_ENGINE_ABST_CACHED}); + } else if (grpid == io_count - 1) { + ret = prepare_io_with_alloc(io_info, &m_handle->io_data[grpid], + {AX_ENGINE_ABST_DEFAULT, AX_ENGINE_ABST_CACHED}, skip_alloc_names); + } else { + ret = prepare_io_struct_only(io_info, &m_handle->io_data[grpid]); + } + if (ret != 0) return ret; + } - 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; + if (io_count > 2) { + auto &first_io_data = m_handle->io_data[0]; + auto &first_io_info = m_handle->io_info[0]; + auto &last_io_data = m_handle->io_data[io_count - 1]; + auto &last_io_info = m_handle->io_info[io_count - 1]; + for (uint i = 0; i < last_io_data.nInputSize; ++i) { + if (std::find(skip_alloc_names.begin(), skip_alloc_names.end(), last_io_info->pInputs[i].pName) != + skip_alloc_names.end()) { + for (uint j = 0; j < first_io_data.nInputSize; ++j) { + if (first_io_info->pInputs[j].pName == last_io_info->pInputs[i].pName) { + last_io_data.pInputs[i].phyAddr = first_io_data.pInputs[j].phyAddr; + last_io_data.pInputs[i].pVirAddr = first_io_data.pInputs[j].pVirAddr; + } + } } } - for (size_t grpid = 0; grpid < io_count; grpid++) - { + for (size_t grpid = 1; grpid < io_count - 1; 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); + + size_t min_inputs = std::min(io_info->nInputSize, last_io_data.nInputSize); + for (size_t i = 0; i < min_inputs; i++) { + io_data.pInputs[i].phyAddr = last_io_data.pInputs[i].phyAddr; + io_data.pInputs[i].pVirAddr = last_io_data.pInputs[i].pVirAddr; } - 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); + size_t min_outputs = std::min(io_info->nOutputSize, last_io_data.nOutputSize); + for (size_t i = 0; i < min_outputs; i++) { + io_data.pOutputs[i].phyAddr = last_io_data.pOutputs[i].phyAddr; + io_data.pOutputs[i].pVirAddr = last_io_data.pOutputs[i].pVirAddr; } } - - moutput_tensors = mgroup_output_tensors[0]; - minput_tensors = mgroup_input_tensors[0]; - - _parepare_io = true; - } - else - { } - 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 = io_info->pOutputs[i].pName ? std::string(io_info->pOutputs[i].pName) : ""; + tensor.nSize = io_info->pOutputs[i].nSize; + tensor.phyAddr = io_data.pOutputs[i].phyAddr; + tensor.pVirAddr = io_data.pOutputs[i].pVirAddr; + for (size_t j = 0; j < io_info->pOutputs[i].nShapeSize; j++) { + tensor.vShape.push_back(io_info->pOutputs[i].pShape[j]); + } + 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 = io_info->pInputs[i].pName ? std::string(io_info->pInputs[i].pName) : ""; + tensor.nSize = io_info->pInputs[i].nSize; + tensor.phyAddr = io_data.pInputs[i].phyAddr; + tensor.pVirAddr = io_data.pInputs[i].pVirAddr; + for (size_t j = 0; j < io_info->pInputs[i].nShapeSize; j++) { + tensor.vShape.push_back(io_info->pInputs[i].pShape[j]); + } + mgroup_input_tensors[grpid].push_back(tensor); + } + } + + if (!mgroup_output_tensors.empty()) moutput_tensors = mgroup_output_tensors[0]; + if (!mgroup_input_tensors.empty()) minput_tensors = mgroup_input_tensors[0]; + + build_tensor_maps(); + + return 0; } int ax_runner_ax650::init(const char *model_file, bool use_mmap) { - if (use_mmap) - { + if (use_mmap) { MMap model_buffer(model_file); - if (!model_buffer.data()) - { - ALOGE("mmap"); - return -1; - } + if (!model_buffer.data()) 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; - } + } else { + char *model_buffer = nullptr; + size_t len = 0; + if (!read_file(model_file, &model_buffer, &len)) return -1; auto ret = init(model_buffer, len); delete[] model_buffer; return ret; @@ -345,96 +225,84 @@ int ax_runner_ax650::init(const char *model_file, bool use_mmap) 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"); + if (m_handle) deinit(); + m_handle = new ax_runner_ax650_handle_t; + int ret = AX_ENGINE_CreateHandle(&m_handle->handle, model_buffer, model_size); + if (0 != ret) { + ALOGE("AX_ENGINE_CreateHandle failed: 0x%x", ret); + delete m_handle; + m_handle = nullptr; return ret; } - // fprintf(stdout, "Engine creating handle is done.\n"); - return sub_init(); } -void ax_runner_ax650::release() +void ax_runner_ax650::deinit() { - 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]); + if (!m_handle) return; + + std::unordered_set freed_phy_addrs; + + for (size_t g = 0; g < m_handle->io_data.size(); ++g) { + auto &io = m_handle->io_data[g]; + + if (io.pInputs) { + for (size_t j = 0; j < io.nInputSize; ++j) { + AX_ENGINE_IO_BUFFER_T *pBuf = io.pInputs + j; + if (pBuf->phyAddr != 0) { + if (freed_phy_addrs.find(pBuf->phyAddr) == freed_phy_addrs.end()) { + AX_SYS_MemFree(pBuf->phyAddr, pBuf->pVirAddr); + freed_phy_addrs.insert(pBuf->phyAddr); + } + } + } + delete[] io.pInputs; + io.pInputs = nullptr; } - AX_ENGINE_DestroyHandle(m_handle->handle); - m_handle->handle = nullptr; + if (io.pOutputs) { + for (size_t j = 0; j < io.nOutputSize; ++j) { + AX_ENGINE_IO_BUFFER_T *pBuf = io.pOutputs + j; + if (pBuf->phyAddr != 0) { + if (freed_phy_addrs.find(pBuf->phyAddr) == freed_phy_addrs.end()) { + AX_SYS_MemFree(pBuf->phyAddr, pBuf->pVirAddr); + freed_phy_addrs.insert(pBuf->phyAddr); + } + } + } + delete[] io.pOutputs; + io.pOutputs = nullptr; + } } - if (m_handle) - { - delete m_handle; - m_handle = nullptr; + if (m_handle->handle) { + AX_ENGINE_DestroyHandle(m_handle->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() { + if (!m_handle) return -1; + for (size_t i = 0; i < get_num_inputs(); i++) { + auto &tensor = get_input(i); + AX_SYS_MflushCache(tensor.phyAddr, tensor.pVirAddr, tensor.nSize); + } + int ret = AX_ENGINE_RunSync(m_handle->handle, &m_handle->io_data[0]); - for (size_t i = 0; i < get_num_outputs(); i++) - { + + for (size_t i = 0; i < get_num_outputs(); i++) { auto &tensor = get_output(i); AX_SYS_MinvalidateCache(tensor.phyAddr, tensor.pVirAddr, tensor.nSize); } @@ -443,12 +311,19 @@ int ax_runner_ax650::inference() int ax_runner_ax650::inference(int grpid) { + if (!m_handle) return -1; + if (grpid < 0 || grpid >= (int)m_handle->io_data.size()) return -1; + + for (size_t i = 0; i < mgroup_input_tensors[grpid].size(); i++) { + auto &tensor = mgroup_input_tensors[grpid][i]; + AX_SYS_MflushCache(tensor.phyAddr, tensor.pVirAddr, tensor.nSize); + } + 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); + for (size_t i = 0; i < mgroup_output_tensors[grpid].size(); i++) { + auto &tensor = mgroup_output_tensors[grpid][i]; AX_SYS_MinvalidateCache(tensor.phyAddr, tensor.pVirAddr, tensor.nSize); } return ret; -} +} \ No newline at end of file diff --git a/projects/llm_framework/main_llm/src/runner/ax_model_runner/ax_model_runner_ax650.hpp b/projects/llm_framework/main_llm/src/runner/ax_model_runner/ax_model_runner_ax650.hpp index 351fb04..62b6538 100644 --- a/projects/llm_framework/main_llm/src/runner/ax_model_runner/ax_model_runner_ax650.hpp +++ b/projects/llm_framework/main_llm/src/runner/ax_model_runner/ax_model_runner_ax650.hpp @@ -1,20 +1,23 @@ #pragma once #include "ax_model_runner.hpp" -class ax_runner_ax650 : public ax_runner_base -{ +struct ax_runner_ax650_handle_t; + +class ax_runner_ax650 : public ax_runner_base { protected: - struct ax_joint_runner_ax650_handle_t *m_handle = nullptr; - - bool _parepare_io = false; - + struct ax_runner_ax650_handle_t *m_handle = nullptr; int sub_init(); public: + ax_runner_ax650() = default; + virtual ~ax_runner_ax650() + { + deinit(); + } + 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; diff --git a/projects/llm_framework/main_vlm/src/runner/LLM.hpp b/projects/llm_framework/main_vlm/src/runner/LLM.hpp index 59dd87e..3d6ad98 100644 --- a/projects/llm_framework/main_vlm/src/runner/LLM.hpp +++ b/projects/llm_framework/main_vlm/src/runner/LLM.hpp @@ -285,11 +285,11 @@ public: void Deinit() { for (int i = 0; i < _attr.axmodel_num; i++) { - llama_layers[i].layer.release(); + llama_layers[i].layer.deinit(); } - llama_post.release(); - vpm_encoder.release(); - vpm_resampler.release(); + llama_post.deinit(); + vpm_encoder.deinit(); + vpm_resampler.deinit(); embed_selector.Deinit(); } @@ -856,10 +856,10 @@ public: void Deinit() { for (int i = 0; i < _attr.axmodel_num; i++) { - llama_layers[i].layer.release(); + llama_layers[i].layer.deinit(); } - llama_post.release(); - image_encoder.release(); + llama_post.deinit(); + image_encoder.deinit(); embed_selector.Deinit(); } @@ -1958,10 +1958,10 @@ public: void Deinit() { for (int i = 0; i < _attr.axmodel_num; i++) { - llama_layers[i].layer.release(); + llama_layers[i].layer.deinit(); } - llama_post.release(); - image_encoder.release(); + llama_post.deinit(); + image_encoder.deinit(); embed_selector.Deinit(); } diff --git a/projects/llm_framework/main_vlm/src/runner/ax_model_runner/ax_model_runner.hpp b/projects/llm_framework/main_vlm/src/runner/ax_model_runner/ax_model_runner.hpp index d8c78d8..c4bcb2a 100644 --- a/projects/llm_framework/main_vlm/src/runner/ax_model_runner/ax_model_runner.hpp +++ b/projects/llm_framework/main_vlm/src/runner/ax_model_runner/ax_model_runner.hpp @@ -4,31 +4,7 @@ #include #include -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 -{ +typedef struct { std::string sName; unsigned int nIdx; std::vector vShape; @@ -37,8 +13,7 @@ typedef struct void *pVirAddr; } ax_runner_tensor_t; -class ax_runner_base -{ +class ax_runner_base { protected: std::vector moutput_tensors; std::vector minput_tensors; @@ -52,106 +27,124 @@ protected: std::map> map_group_output_tensors; std::map> map_group_input_tensors; + void build_tensor_maps() + { + map_input_tensors.clear(); + for (const auto &t : minput_tensors) map_input_tensors[t.sName] = t; + + map_output_tensors.clear(); + for (const auto &t : moutput_tensors) map_output_tensors[t.sName] = t; + + map_group_input_tensors.clear(); + for (const auto &grp : mgroup_input_tensors) { + for (const auto &t : grp) map_group_input_tensors[t.sName].push_back(t); + } + + map_group_output_tensors.clear(); + for (const auto &grp : mgroup_output_tensors) { + for (const auto &t : grp) map_group_output_tensors[t.sName].push_back(t); + } + } + public: + virtual ~ax_runner_base() + { + } + virtual int init(const char *model_file, bool use_mmap = false) = 0; - virtual int init(char *model_buffer, size_t model_size) = 0; + virtual int init(char *model_buffer, size_t model_size) = 0; + virtual void deinit() = 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) + int get_num_inputs() { - 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 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(); + }; - return map_input_tensors[name]; + 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(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) + const ax_runner_tensor_t &get_input(const 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]; + auto it = map_input_tensors.find(name); + if (it == map_input_tensors.end()) throw std::runtime_error("input tensor not found: " + name); + return it->second; } - 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) + const ax_runner_tensor_t &get_input(int grpid, int idx) { - 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]; + 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_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) + const ax_runner_tensor_t &get_input(int grpid, const 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]; + auto it = map_group_input_tensors.find(name); + if (it == map_group_input_tensors.end()) throw std::runtime_error("input tensor not found: " + name); + if (grpid < 0 || grpid >= (int)it->second.size()) + throw std::runtime_error("group id out of range for: " + name); + return it->second[grpid]; } - virtual int inference() = 0; + 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(const std::string &name) + { + auto it = map_output_tensors.find(name); + if (it == map_output_tensors.end()) throw std::runtime_error("output tensor not found: " + name); + return it->second; + } + + 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, const std::string &name) + { + auto it = map_group_output_tensors.find(name); + if (it == map_group_output_tensors.end()) throw std::runtime_error("output tensor not found: " + name); + if (grpid < 0 || grpid >= (int)it->second.size()) + throw std::runtime_error("group id out of range for: " + name); + return it->second[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); \ No newline at end of file +}; \ No newline at end of file diff --git a/projects/llm_framework/main_vlm/src/runner/ax_model_runner/ax_model_runner_ax650.cpp b/projects/llm_framework/main_vlm/src/runner/ax_model_runner/ax_model_runner_ax650.cpp index 0454229..363c3d5 100644 --- a/projects/llm_framework/main_vlm/src/runner/ax_model_runner/ax_model_runner_ax650.cpp +++ b/projects/llm_framework/main_vlm/src/runner/ax_model_runner/ax_model_runner_ax650.cpp @@ -1,8 +1,9 @@ #include "ax_model_runner_ax650.hpp" -#include "string.h" -#include "fstream" -#include "memory" -// #include "utilities/file.hpp" +#include +#include +#include +#include +#include #include #include #include @@ -11,332 +12,211 @@ #include "sample_log.h" #define AX_CMM_ALIGN_SIZE 128 - const char *AX_CMM_SESSION_NAME = "npu"; -typedef enum -{ +typedef enum { AX_ENGINE_ABST_DEFAULT = 0, - AX_ENGINE_ABST_CACHED = 1, + AX_ENGINE_ABST_CACHED = 1, } AX_ENGINE_ALLOC_BUFFER_STRATEGY_T; -typedef std::pair INPUT_OUTPUT_ALLOC_STRATEGY; +struct ax_runner_ax650_handle_t { + AX_ENGINE_HANDLE handle = nullptr; + AX_ENGINE_CONTEXT_T context = 0; + std::vector io_info; + std::vector io_data; +}; -static void print_io_info(AX_ENGINE_IO_INFO_T *io_info) -{ - static std::map 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 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) +static int prepare_io_struct_only(AX_ENGINE_IO_INFO_T *info, AX_ENGINE_IO_T *io_data) { memset(io_data, 0, sizeof(*io_data)); - io_data->pInputs = new AX_ENGINE_IO_BUFFER_T[info->nInputSize]; + io_data->pInputs = new AX_ENGINE_IO_BUFFER_T[info->nInputSize]; io_data->nInputSize = info->nInputSize; + memset(io_data->pInputs, 0, sizeof(AX_ENGINE_IO_BUFFER_T) * 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->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); - } + memset(io_data->pOutputs, 0, sizeof(AX_ENGINE_IO_BUFFER_T) * info->nOutputSize); + + for (uint i = 0; i < info->nInputSize; ++i) io_data->pInputs[i].nSize = info->pInputs[i].nSize; + for (uint i = 0; i < info->nOutputSize; ++i) io_data->pOutputs[i].nSize = info->pOutputs[i].nSize; return 0; } -struct ax_joint_runner_ax650_handle_t +static int prepare_io_with_alloc( + AX_ENGINE_IO_INFO_T *info, AX_ENGINE_IO_T *io_data, + std::pair strategy, + std::vector skip_alloc_names = {}) { - AX_ENGINE_HANDLE handle; - AX_ENGINE_CONTEXT_T context; - std::vector io_info; - std::vector io_data; + int ret = prepare_io_struct_only(info, io_data); + if (ret != 0) return ret; - // int algo_width, algo_height; - // int algo_colorformat; -}; + for (uint i = 0; i < info->nInputSize; ++i) { + auto &buffer = io_data->pInputs[i]; + if (std::find(skip_alloc_names.begin(), skip_alloc_names.end(), info->pInputs[i].pName) != + skip_alloc_names.end()) { + continue; + } + if (strategy.first == AX_ENGINE_ABST_CACHED) { + ret = AX_SYS_MemAllocCached((AX_U64 *)(&buffer.phyAddr), &buffer.pVirAddr, buffer.nSize, AX_CMM_ALIGN_SIZE, + (const AX_S8 *)(AX_CMM_SESSION_NAME)); + } else { + ret = AX_SYS_MemAlloc((AX_U64 *)(&buffer.phyAddr), &buffer.pVirAddr, buffer.nSize, AX_CMM_ALIGN_SIZE, + (const AX_S8 *)(AX_CMM_SESSION_NAME)); + } + if (ret != 0) { + ALOGE("Alloc input[%d] failed", i); + return ret; + } + memset(buffer.pVirAddr, 0, buffer.nSize); + } + + for (uint i = 0; i < info->nOutputSize; ++i) { + auto &buffer = io_data->pOutputs[i]; + if (std::find(skip_alloc_names.begin(), skip_alloc_names.end(), info->pOutputs[i].pName) != + skip_alloc_names.end()) { + continue; + } + if (strategy.second == AX_ENGINE_ABST_CACHED) { + ret = AX_SYS_MemAllocCached((AX_U64 *)(&buffer.phyAddr), &buffer.pVirAddr, buffer.nSize, AX_CMM_ALIGN_SIZE, + (const AX_S8 *)(AX_CMM_SESSION_NAME)); + } else { + ret = AX_SYS_MemAlloc((AX_U64 *)(&buffer.phyAddr), &buffer.pVirAddr, buffer.nSize, AX_CMM_ALIGN_SIZE, + (const AX_S8 *)(AX_CMM_SESSION_NAME)); + } + if (ret != 0) { + ALOGE("Alloc output[%d] failed", i); + return ret; + } + memset(buffer.pVirAddr, 0, buffer.nSize); + } + return 0; +} int ax_runner_ax650::sub_init() { - // 4. create context + if (!m_handle) return -1; + int ret = AX_ENGINE_CreateContext(m_handle->handle); - if (0 != ret) - { - ALOGE("AX_ENGINE_CreateContext"); - return ret; - } + if (ret != 0) 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"); + if (ret != 0) return ret; - // 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); + ret = AX_ENGINE_GetGroupIOInfoCount(m_handle->handle, &io_count); + if (ret != 0) return ret; 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"); + std::vector skip_alloc_names = {"K_cache", "V_cache"}; + 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 (ret != 0) return ret; + m_handle->io_info[grpid] = io_info; - // 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); + if (grpid == 0) { + ret = prepare_io_with_alloc(io_info, &m_handle->io_data[grpid], + {AX_ENGINE_ABST_DEFAULT, AX_ENGINE_ABST_CACHED}); + } else if (grpid == io_count - 1) { + ret = prepare_io_with_alloc(io_info, &m_handle->io_data[grpid], + {AX_ENGINE_ABST_DEFAULT, AX_ENGINE_ABST_CACHED}, skip_alloc_names); + } else { + ret = prepare_io_struct_only(io_info, &m_handle->io_data[grpid]); + } + if (ret != 0) return ret; + } - 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; + if (io_count > 2) { + auto &first_io_data = m_handle->io_data[0]; + auto &first_io_info = m_handle->io_info[0]; + auto &last_io_data = m_handle->io_data[io_count - 1]; + auto &last_io_info = m_handle->io_info[io_count - 1]; + for (uint i = 0; i < last_io_data.nInputSize; ++i) { + if (std::find(skip_alloc_names.begin(), skip_alloc_names.end(), last_io_info->pInputs[i].pName) != + skip_alloc_names.end()) { + for (uint j = 0; j < first_io_data.nInputSize; ++j) { + if (first_io_info->pInputs[j].pName == last_io_info->pInputs[i].pName) { + last_io_data.pInputs[i].phyAddr = first_io_data.pInputs[j].phyAddr; + last_io_data.pInputs[i].pVirAddr = first_io_data.pInputs[j].pVirAddr; + } + } } } - for (size_t grpid = 0; grpid < io_count; grpid++) - { + for (size_t grpid = 1; grpid < io_count - 1; 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); + + size_t min_inputs = std::min(io_info->nInputSize, last_io_data.nInputSize); + for (size_t i = 0; i < min_inputs; i++) { + io_data.pInputs[i].phyAddr = last_io_data.pInputs[i].phyAddr; + io_data.pInputs[i].pVirAddr = last_io_data.pInputs[i].pVirAddr; } - 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); + size_t min_outputs = std::min(io_info->nOutputSize, last_io_data.nOutputSize); + for (size_t i = 0; i < min_outputs; i++) { + io_data.pOutputs[i].phyAddr = last_io_data.pOutputs[i].phyAddr; + io_data.pOutputs[i].pVirAddr = last_io_data.pOutputs[i].pVirAddr; } } - - moutput_tensors = mgroup_output_tensors[0]; - minput_tensors = mgroup_input_tensors[0]; - - _parepare_io = true; - } - else - { } - 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 = io_info->pOutputs[i].pName ? std::string(io_info->pOutputs[i].pName) : ""; + tensor.nSize = io_info->pOutputs[i].nSize; + tensor.phyAddr = io_data.pOutputs[i].phyAddr; + tensor.pVirAddr = io_data.pOutputs[i].pVirAddr; + for (size_t j = 0; j < io_info->pOutputs[i].nShapeSize; j++) { + tensor.vShape.push_back(io_info->pOutputs[i].pShape[j]); + } + 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 = io_info->pInputs[i].pName ? std::string(io_info->pInputs[i].pName) : ""; + tensor.nSize = io_info->pInputs[i].nSize; + tensor.phyAddr = io_data.pInputs[i].phyAddr; + tensor.pVirAddr = io_data.pInputs[i].pVirAddr; + for (size_t j = 0; j < io_info->pInputs[i].nShapeSize; j++) { + tensor.vShape.push_back(io_info->pInputs[i].pShape[j]); + } + mgroup_input_tensors[grpid].push_back(tensor); + } + } + + if (!mgroup_output_tensors.empty()) moutput_tensors = mgroup_output_tensors[0]; + if (!mgroup_input_tensors.empty()) minput_tensors = mgroup_input_tensors[0]; + + build_tensor_maps(); + + return 0; } int ax_runner_ax650::init(const char *model_file, bool use_mmap) { - if (use_mmap) - { + if (use_mmap) { MMap model_buffer(model_file); - if (!model_buffer.data()) - { - ALOGE("mmap"); - return -1; - } + if (!model_buffer.data()) 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; - } + } else { + char *model_buffer = nullptr; + size_t len = 0; + if (!read_file(model_file, &model_buffer, &len)) return -1; auto ret = init(model_buffer, len); delete[] model_buffer; return ret; @@ -345,97 +225,105 @@ int ax_runner_ax650::init(const char *model_file, bool use_mmap) 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"); + if (m_handle) deinit(); + m_handle = new ax_runner_ax650_handle_t; + int ret = AX_ENGINE_CreateHandle(&m_handle->handle, model_buffer, model_size); + if (0 != ret) { + ALOGE("AX_ENGINE_CreateHandle failed: 0x%x", ret); + delete m_handle; + m_handle = nullptr; return ret; } - // fprintf(stdout, "Engine creating handle is done.\n"); - return sub_init(); } -void ax_runner_ax650::release() +void ax_runner_ax650::deinit() { - 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]); + if (!m_handle) return; + + std::unordered_set freed_phy_addrs; + + for (size_t g = 0; g < m_handle->io_data.size(); ++g) { + auto &io = m_handle->io_data[g]; + + if (io.pInputs) { + for (size_t j = 0; j < io.nInputSize; ++j) { + AX_ENGINE_IO_BUFFER_T *pBuf = io.pInputs + j; + if (pBuf->phyAddr != 0) { + if (freed_phy_addrs.find(pBuf->phyAddr) == freed_phy_addrs.end()) { + AX_SYS_MemFree(pBuf->phyAddr, pBuf->pVirAddr); + freed_phy_addrs.insert(pBuf->phyAddr); + } + } + } + delete[] io.pInputs; + io.pInputs = nullptr; } - AX_ENGINE_DestroyHandle(m_handle->handle); - m_handle->handle = nullptr; + if (io.pOutputs) { + for (size_t j = 0; j < io.nOutputSize; ++j) { + AX_ENGINE_IO_BUFFER_T *pBuf = io.pOutputs + j; + if (pBuf->phyAddr != 0) { + if (freed_phy_addrs.find(pBuf->phyAddr) == freed_phy_addrs.end()) { + AX_SYS_MemFree(pBuf->phyAddr, pBuf->pVirAddr); + freed_phy_addrs.insert(pBuf->phyAddr); + } + } + } + delete[] io.pOutputs; + io.pOutputs = nullptr; + } } - if (m_handle) - { - delete m_handle; - m_handle = nullptr; + if (m_handle->handle) { + AX_ENGINE_DestroyHandle(m_handle->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() { - return AX_ENGINE_RunSync(m_handle->handle, &m_handle->io_data[0]); + if (!m_handle) return -1; + for (size_t i = 0; i < get_num_inputs(); i++) { + auto &tensor = get_input(i); + AX_SYS_MflushCache(tensor.phyAddr, tensor.pVirAddr, tensor.nSize); + } + + 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) { - return AX_ENGINE_RunGroupIOSync(m_handle->handle, m_handle->context, grpid, &m_handle->io_data[grpid]); + if (!m_handle) return -1; + if (grpid < 0 || grpid >= (int)m_handle->io_data.size()) return -1; + + for (size_t i = 0; i < mgroup_input_tensors[grpid].size(); i++) { + auto &tensor = mgroup_input_tensors[grpid][i]; + AX_SYS_MflushCache(tensor.phyAddr, tensor.pVirAddr, tensor.nSize); + } + + int ret = AX_ENGINE_RunGroupIOSync(m_handle->handle, m_handle->context, grpid, &m_handle->io_data[grpid]); + + for (size_t i = 0; i < mgroup_output_tensors[grpid].size(); i++) { + auto &tensor = mgroup_output_tensors[grpid][i]; + AX_SYS_MinvalidateCache(tensor.phyAddr, tensor.pVirAddr, tensor.nSize); + } + return ret; } \ No newline at end of file diff --git a/projects/llm_framework/main_vlm/src/runner/ax_model_runner/ax_model_runner_ax650.hpp b/projects/llm_framework/main_vlm/src/runner/ax_model_runner/ax_model_runner_ax650.hpp index 351fb04..62b6538 100644 --- a/projects/llm_framework/main_vlm/src/runner/ax_model_runner/ax_model_runner_ax650.hpp +++ b/projects/llm_framework/main_vlm/src/runner/ax_model_runner/ax_model_runner_ax650.hpp @@ -1,20 +1,23 @@ #pragma once #include "ax_model_runner.hpp" -class ax_runner_ax650 : public ax_runner_base -{ +struct ax_runner_ax650_handle_t; + +class ax_runner_ax650 : public ax_runner_base { protected: - struct ax_joint_runner_ax650_handle_t *m_handle = nullptr; - - bool _parepare_io = false; - + struct ax_runner_ax650_handle_t *m_handle = nullptr; int sub_init(); public: + ax_runner_ax650() = default; + virtual ~ax_runner_ax650() + { + deinit(); + } + 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;