UnTool (Union Tool) 说明文档

项目概述

UnTool 是一个用于 BM1684X 芯片的模型推理工具包，提供高效的模型部署和执行功能。该框架支持通用模型推理以及大语言模型(LLM)的推理，具有以下特点：

最小化依赖，重写bmruntime;
优化内存管理，统一将host的mem和device的mem管理为tensor，快速实现两端的数据同步；支持numpy，共享numpy数据的内存地址，零copy；
提供C++和Python双语言接口；使用ctypes链接c++和python，脱离版本依赖；
支持多nets多stages以及dynamic subnets；支持通用模型和LLM模型
融合tpukernel，将算子和tensor绑定，更方便debug和调用算子。（todo）

UnTool 框架允许用户从底层的C/C++接口到高层的Python API进行选择，支持从简单到复杂的多种使用场景。

项目结构

UnTool/
├── include/                    # C++ 头文件目录
│   ├── kernel/                 # 内核操作相关
│   │   └── tpu_kernel_runtime.h # TPU 内核运行时
│   ├── libsophon/              # Sophon芯片基础库
│   ├── llm/                    # LLM模型相关接口
│   │   └── llmruntime.h        # LLM运行时接口
│   ├── runtime/                # 运行时核心组件
│   │   ├── minibmrt.h          # 底层推理引擎接口
│   │   ├── miniruntime.h       # 通用模型推理接口
│   │   └── minitensor.h        # 张量操作接口
│   └── utils/                  # 实用工具
│       ├── datagen.h           # 数据生成工具
│       ├── device.h            # 设备管理工具
│       └── logs.h              # 日志工具
│
├── lib/                        # 库文件目录
│   └── libbm1684x_kernel_module.so # BM1684X 内核模块
│
├── python/                     # Python 接口目录
│   ├── cfun.txt                # C函数定义
│   ├── engineov.py             # 通用模型推理引擎
│   ├── enginellm.py            # LLM模型推理引擎(C封装，扩展性差，性能好)
│   ├── LLMBaseModel.py         # LLM模型基类(python封装，扩展性好，性能稍差)
│   ├── LLMBasePipeline.py      # LLM模型推理引擎(基于LLMBaseModel.py)
│   ├── help.py                 # 生成Python绑定的辅助脚本
│   └── untool.py               # 底层库绑定接口(由help.py生成)
│
├── release/                    # submodule: 用于发布 untool pypi包
│
├── CMakeLists.txt              # CMake 构建配置
├── envsetup.sh                 # 环境设置脚本
├── externc.cpp                 # 对外暴露的C接口实现
└── README.md                   # 项目说明文档

使用方法

1. 需要修改 C++源码，编译安装

source envsetup.sh

编译后的文件位置为：

install(TARGETS untool DESTINATION /opt/untool/lib)
install(DIRECTORY ${PROJECT_SOURCE_DIR}/include/ DESTINATION /opt/untool/include)
install(DIRECTORY ${PROJECT_SOURCE_DIR}/python/ DESTINATION /opt/untool/python)

.so文件在 /opt/untool/lib/libuntpu.so，.h文件在 /opt/untool/include/ 里面。如果需要使用这个库，可以通过这种方法引用：

set(UNTPU /opt/untool)
include_directories(${UNTPU}/include)
link_directories(${UNTPU}/lib)

Python 接口使用方法：

import sys
sys.path.append("/data/untool2/python")

2. 只使用 Python，pip安装

pip install untooll

import untool

框架架构

UnTool 框架采用分层设计，总体架构如下：

┌───────────────────────────────────┐
│       Python 应用层接口             │
│   (engineov.py / enginellm.py)    │
├───────────────────────────────────┤
│       Python 底层库绑定 (ctypes)    │ 
│    (untool.py <--> libuntool.so)  │
├───────────────────────────────────┤
│                                   │
│        C++  externc.cpp           │
│   ┌────────────┐  ┌─────────────┐ │
│   │            │  │             │ │
│   │ llmruntime │  │ miniruntime │ │
│   │            │  │             │ │
│   └────────────┘  └─────────────┘ │
│            │            │         │
│            ▼            ▼         │
│ ┌───────────────────────────────┐ │
│ │    张量操作     &  底层模型操作  │ │
│ │  (minitensor)    (minibmrt)   │ │
│ └───────────────────────────────┘ │
└───────────────────────────────────┘

源码说明

推理引擎 (minibmrt.h)

这是模型运行时的核心模块，负责bmodel模型加载、编译和运行。

主要功能：

bmodel模型加载与解析
内存管理和资源分配
命令编译与设备同步
网络执行与资源回收

核心数据结构：

ModelInfo结构体

struct ModelInfo {
    int                          device_id;     // 设备ID，用于多设备环境
    bm_handle_t                  bm_handle;     // 设备句柄，用于管理设备资源
    std::string                  chip_name;     // 芯片名称，如BM1684X
    std::vector<NetInfo>         nets;          // 网络列表，包含模型中所有网络
    size_t                       net_num;       // 网络数量
    uint64_t                     neuron_size;   // 神经元内存大小
    bm_device_mem_t              neuron_device; // 神经元设备内存
    uint64_t                     neuron_global_addr; // 神经元全局地址
    std::unordered_map<std::string, size_t> net_name_to_index; // 网络名称到索引的映射
};

NetInfo结构体

struct NetInfo {
    std::string                  name;       // 网络名称
    std::vector<StageInfo>       stages;     // 网络阶段列表
    size_t                       stage_num;  // 阶段数量
    int                          addr_mode;  // 地址模式(0:neuron共享，1:io独立)
};

StageInfo结构体

struct StageInfo {
    std::vector<TensorInfo>      input_tensor;   // 输入张量信息
    std::vector<TensorInfo>      output_tensor;  // 输出张量信息
    std::vector<uint64_t>        input_tensor_global_addr;  // 输入张量全局地址
    std::vector<uint64_t>        output_tensor_global_addr; // 输出张量全局地址

    // IO相关
    bool                         io_alone;       // IO是否独立
    uint64_t                     io_addr;        // IO地址
    uint64_t                     io_size;        // IO大小
    uint64_t                     io_offset;      // IO偏移
    uint64_t                     io_global_addr; // IO全局地址
    bm_device_mem_t              io_device;      // IO设备内存

    // 子网络相关
    std::vector<SubNetInfo>      subnets;        // 子网络列表
    size_t                       subnet_num;     // 子网络数量
    int32_t                      is_dynamic;     // 是否为动态网络
};

核心API函数：

// 加载bmodel模型
struct ModelInfo read_bmodel(const char* filename, int device_id = 0);

// 运行模型指定的网络和阶段
void run_model(struct ModelInfo* model_info_p, size_t net_idx, size_t stage_idx);

// 释放模型资源
void free_model(struct ModelInfo* model_info_p);

// 编译IO地址
void compile_io_addr(struct ModelInfo* model_info_p);

// 将模型数据移动到设备
void move_to_device(struct ModelInfo* model_info_p);

// 填充API信息
void fill_api_info(struct ModelInfo* model_info_p);

张量操作 (minitensor.h)

管理张量数据，处理主机与设备之间的数据传输，提供张量的创建、操作和销毁功能。

主要功能：

张量创建与内存管理
数据在主机和设备间的同步
数据类型转换与显示

核心数据结构：

UnTensor结构体

struct UnTensor {
    // 基本信息
    char         name[MAX_CHAR_NUM]; // 张量名称，用于标识不同张量
    int          dtype;              // 数据类型(BM_FLOAT32, BM_INT8等)
    size_t       size;               // 数据大小(字节)
    size_t       dims;               // 维度数(1D, 2D, 3D, 4D等)
    uint64_t     shape[MAX_SHAPE_NUM]; // 形状数组，存储每个维度的大小

    // 主机内存相关
    void*        data;               // 主机端数据指针，指向CPU内存中的数据
    bool         is_malloc_host;     // 是否在主机端分配了内存
    bool         is_have_data;       // 是否有有效数据

    // 设备内存相关
    int          device_id;          // 设备ID，用于多设备环境
    bm_handle_t  bm_handle;          // 设备句柄，用于管理设备资源
    bool         is_in_device;       // 数据是否在设备中
    u64          addr;               // 设备内存地址，指向TPU内存中的数据位置
    u64          device_start;       // 设备内存起始地址
    u64          device_size;        // 设备内存大小
    u64          offset;             // 偏移量，相对于device_start的偏移
    int          dmabuf_fd;          // DMA缓冲区文件描述符
    uint32_t     reserved;           // 保留字段
    uint32_t     rawflags;           // 原始标志位
};

核心API函数：

// 创建一个新的空张量
untensor untensor_create();

// 初始化张量
void untensor_init(untensor tensor);

// 设置张量数据
void untensor_set_data(untensor tensor, void* data, size_t size, bool copy = true);

// 在主机和设备间同步数据
void untensor_sync(untensor tensor, bool to_device = true, bool force = true);

// 显示张量数据
void untensor_show(untensor tensor, int start = 0, int len = -1, char location = 'd');

// 释放张量资源
void untensor_destroy(untensor tensor);

通用模型推理 (miniruntime.h)

提供通用模型的推理接口，包括模型加载、设置输入、推理执行以及获取输出结果。

主要功能：

模型加载与网络阶段设置
输入数据传输(主机到设备)
模型推理执行
输出结果获取(设备到主机)
资源管理与回收

核心数据结构：

UnRuntime结构体

struct UnRuntime {
    bm_handle_t           bm_handle;       // 模型句柄
    int                   device_id;       // 设备ID
    ModelInfo             model_info;      // 模型信息
    size_t                cur_net_idx;     // 当前网络索引
    size_t                cur_stage_idx;   // 当前阶段索引
    std::vector<untensor> input_tensors;   // 输入张量
    std::vector<untensor> output_tensors;  // 输出张量
};

核心API函数：

// 加载bmodel模型并创建runtime对象
unruntime unruntime_init(const char* bmodel_path, int device_id);

// 获取网络名称对应的 net_id
size_t unruntime_get_net_num(unruntime runtime, const char* net_name);

// 重设运行时的网络和阶段（初始化时，默认设为 net0 stage0）
void unruntime_set_net_stage(unruntime runtime, size_t net_idx, size_t stage_idx);

// 获取io untensor（参数 io_type：'i'为输入，'o'为输出）
untensor unruntime_get_io_tensor(unruntime runtime, char io_type, size_t tensor_idx);

// 获取io 数量（参数 io_type：'i'为输入，'o'为输出）
size_t unruntime_get_io_count(unruntime runtime, char io_type);

// 
void      unruntime_set_input_s2d(unruntime runtime, size_t tensor_idx, void* data, size_t size);
void      unruntime_set_input_d2d(unruntime runtime, size_t tensor_idx, untensor src_tensor);
void      unruntime_run(unruntime runtime, bool output_to_host);
void      unruntime_free(unruntime runtime);

LLM模型推理 (llmruntime.h)

专门为大语言模型优化的推理接口。

主要功能：

LLM模型加载与初始化
Token生成与处理
KV-Cache管理
自回归生成与推理

核心数据结构：

LLMBase结构体

struct LLMBase {
    bm_handle_t bm_handle;                             // 设备句柄，用于管理设备资源
    uint16_t mask;                                     // 注意力掩码值(BF16或F16格式)
    ModelInfo model_info;                              // 模型信息，包含网络结构等详细信息
    size_t seq_len;                                    // 最大序列长度，模型支持的最大token数
    size_t token_len;                                  // 当前token长度，已处理的token数量
    size_t num_layers;                                 // transformer层数量
    size_t embedding_idx;                              // embedding网络索引
    size_t embedding_cache_idx;                        // embedding_cache网络索引
    size_t lm_head_idx;                                // lm_head网络索引，用于生成下一个token
    size_t greedy_idx;                                 // greedy索引
    std::vector<size_t> block_ids;                     // transformer block网络索引
    std::vector<size_t> block_cache_ids;               // block cache网络索引，用于保存KV缓存
    std::vector<std::vector<untensor>> input_tensors;  // 所有网络的输入张量
    std::vector<std::vector<untensor>> output_tensors; // 所有网络的输出张量
};

核心API函数：

// 初始化LLM模型
llmbase llm_init(const char* model_path, int device_id = 0);

// 首次前向推理，处理输入提示并生成第一个token
int llm_forward_first(llmbase llm, int* tokens, size_t token_len);

// 后续推理，基于上一个token生成下一个token
int llm_forward_next(llmbase llm);

// 获取序列最大长度
size_t llm_get_seq_len(llmbase llm);

// 获取当前token长度
size_t llm_get_token_len(llmbase llm);

// 释放LLM资源
void llm_free(llmbase llm);

Python接口

底层库绑定 (untool.py)

这是Python与C/C++库之间的桥梁，负责类型转换和函数调用。untool.py由help.py自动生成，封装了底层C++库的所有功能，提供给Python应用使用。

主要功能：

C/C++库函数的Python封装
数据类型转换（Python类型 <--> C类型）
内存管理和指针操作
张量操作的Python接口

核心接口：

__all__ = [
    # C API
    "convert_model_info",
    "get_model_info_p",
    "find_net_num",
    "get_first_mask_ptr",
    "get_next_mask_ptr",
    "free_mask_ptr",
    "move_to_device",
    "compile_io_addr",
    "fill_api_info",
    "run_model",
    "free_model",
    "untensor_create",
    "untensor_init",
    "untensor_destroy",
    "untensor_set_data",
    "untensor_sync",
    "untensor_show",
    "untensor_s2d_bytes",
    "untensor_d2d_bytes_offset",
    "unruntime_init",
    "unruntime_get_net_num",
    "unruntime_set_net_stage",
    "unruntime_get_io_tensor",
    "unruntime_get_io_count",
    "unruntime_set_input_s2d",
    "unruntime_set_input_d2d",
    "unruntime_run",
    "unruntime_free",
    # Python API
    "EngineLLM",
    "EngineOV",
    "LLMBaseModel",
    "LLMBasePipline",
    # utils
    "get_lib_path",
    "type_map",
    "convert_to_ctypes",
    "convert_to_python",
    "char_point_2_str",
]

通用推理类：EngineOV

class EngineOV:
    def __init__(self, model_path="", device_id=0):
        """
        初始化推理引擎
        
        参数:
            model_path: 模型路径，bmodel文件的位置
            device_id: 设备ID
        """
        
    def __call__(self, args, to_host=True):
        """
        运行推理（可直接调用对象）
        
        参数:
            args: 输入数据列表，每个元素对应一个输入，支持 numpy和 torch.tensor
            to_host: 是否将输出同步到host
            
        返回:
            模型的输出结果列表
        """

使用示例：

# 加载模型
engine = EngineOV("/path/to/model.bmodel", 0)

# 准备输入数据
input_data = np.array([...], dtype=np.float32)

# 运行推理
results = engine([input_data])

# 使用结果
output = results[0]

LLM推理类

EngineLLM

基于llmruntime.h提供的高度封装的接口，性能好。需要调整时，应对底层cpp代码进行修改并重新编译。

class EngineLLM():
    def __init__(self, model_path, config_path, device_id):
        """
        参数:
            model_path: 模型路径，LLM的bmodel文件位置
            config_path: tokenizer配置路径
            device_id: 设备ID，用于多设备环境
        """
        
    def encode(self):
        """
        编码输入文本为tokens
        
        使用tokenizer将input_str编码为token序列，用于模型输入
        
        返回:
            编码后的token数量
        """
        
    def chat(self, enable_history=False):
        """
        交互式聊天接口
        """

使用示例：

import argparse
args = argparse.Namespace(
            model_path="/path/to/model.bmodel",
            tokenizer_path="/path/to/tokenizer",
            device_id=0,
        )
# 启动交互式聊天
engine = EngineLLM(args)
engine.chat()

LLMBaseModel & LLMBasePipeline

LLMBaseModel是基于minibmrt.h提供的底层接口编写的模型基类，可以基于该类进行修改，使用方便，扩展性强。

LLMBasePipeline根据LLMBaseModel编写的推理基类。

使用示例：

python llmbasepipeline.py -m path/to/model.bmodel -t path/to/tokennizer -d 0

UnTool

Model card

UnTool (Union Tool) 说明文档

目录

项目概述

项目结构

使用方法

1. 需要修改 C++源码，编译安装

2. 只使用 Python，pip安装

框架架构

源码说明

推理引擎 (minibmrt.h)

主要功能：

核心数据结构：

ModelInfo结构体

NetInfo结构体

StageInfo结构体

核心API函数：

张量操作 (minitensor.h)

主要功能：

核心数据结构：

UnTensor结构体

核心API函数：

通用模型推理 (miniruntime.h)

主要功能：

核心数据结构：

UnRuntime结构体

核心API函数：

LLM模型推理 (llmruntime.h)

主要功能：

核心数据结构：

LLMBase结构体

核心API函数：

Python接口

底层库绑定 (untool.py)

主要功能：

核心接口：

通用推理类：EngineOV

使用示例：

LLM推理类

EngineLLM

LLMBaseModel & LLMBasePipeline