MNN 介绍

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

source/backend/
├── cpu/          # CPU通用后端
|    └───x86_x64
|    		└─── avx/          # x86 AVX2优化
|    		└─── avx512/       # x86 AVX512优化
|    		└─── avxfma/       # x86 AVXfma 支持Fused Multiply-Add
|	 		└─── sse/          # x86 SSE优化
|    └─── arm
|    		└───arm32		   # 32为arm指令支持 如ARMv7-A后端
|    		└───arm64		   # 64为arm指令支持
├── arm82/        # ARMv8.2 支持
|    └─── asm
|    		└───arm32		   # 32为arm指令兼容
|    		└───arm64		   # 64为arm指令支持
├── nnapi/        # NNAPI后端
├── metal/        # Metal后端（iOS/macOS）
├── opencl/       # OpenCL后端
├── opengl/       # OpenGL后端
├── vulkan/       # Vulkan后端
├── cuda/         # CUDA后端
├── tensorrt
├── coreml/       # CoreML后端
├── hiai/         # HIAI后端
└── qnn/          # QNN后端
└── neuropilot/   # neuropilot后端

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

// 1. 创建 下面的关系只表示时间上的顺序依赖关系 如Backend需要调用Runtime创建
Runtime (运行时, 各类后端继承Runtime, 如class CPURuntime : public Runtime)
  ↓ 创建
Backend (后端, 各类后端继承Backend, 如class CPUBackend : public Backend)
  ↓ 创建
Execution (执行器, 大部分算子实现是通过继承Execution实现,如:class CPUAttention : public Execution;
			小部分为汇编优化过的代码, 如source/backend/cpu/x86_x64/avx512/_AVX512_MNNGemmFloatUnit16x8.S)
			
// 2. 重新计算输出形状 在/workspace/code/MNN/source/shape/SizeComputer.hpp内完成 

// 3. 根据输入长度变化 执行的resize操作 下面的关系只表示时间上的顺序依赖关系
Backend::onResizeBegin (resize的准备阶段)
  ↓
Execution::onResize (算子的resize)
  ↓
Backend::onResizeEnd (resize的收尾阶段)

// 4. 执行算子  下面的关系只表示时间上的顺序依赖关系
Backend::onExecuteBegin (execute的准备阶段)
  ↓
Execution::onExecute (算子的execute)
  ↓
Backend::onExecuteEnd (execute的收尾阶段)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

class Runtime : public NonCopyable {
public:
    // 创建 Backend
    virtual Backend* onCreate(const BackendConfig* config, Backend* origin) const = 0;
    // 重置运行时配置
    virtual void onReset(int numberThread, const BackendConfig* config, bool full);
    
    // 垃圾回收
    virtual void onGabageCollect(int level) = 0;
    // 获取内存使用量
    virtual float onGetMemoryInMB();
	
    // 主要是cpuruntime使用 cpu后端实现了线程池(source/backend/cpu/ThreadPool.hpp)
    virtual void onConcurrencyBegin() const = 0;
    virtual void onConcurrencyEnd() const = 0;
	
    // 执行优化, mnn以算子图形式运行 下面是不同的运行方式
    enum CompilerType {
        Compiler_Geometry = 0, // 部分执行几何计算，分解形变算子，但不分解 BatchMatMul / Gather 等算子
        Compiler_Origin = 1, // 直接使用原始算子，不进行分解
        Compiler_Loop = 2, // 完全执行几何计算，仅此模式下，可以在算子不支持时自动回退到CPU计算
    };
private:    
    // 记录运行时信息 如kvcacheSizeLimit cpuIds等
    RuntimeHint mHint;
};

1
2
3
4
5
6

enum StorageType {
    STATIC,           // 不可重用内存，分配后立即释放
    DYNAMIC,          // 可重用内存，优先重用已有内存
    DYNAMIC_SEPERATE, // 不可重用内存，但延迟释放
    DYNAMIC_IN_EXECUTION // 执行时动态分配
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

class Backend : public NonCopyable {
public:
    // 创建 Execution 
    virtual Execution* onCreate(const std::vector<Tensor*>& inputs,
                                const std::vector<Tensor*>& outputs,
                                const MNN::Op* op) = 0;

    // 执行算子Resize 的开始/结束的一些准备/收尾工作
    virtual void onResizeBegin();
    virtual ErrorCode onResizeEnd() = 0;
    // 执行算子execute 的开始/结束的一些准备/收尾工作
    // 例如OpenCLBackend中 
    virtual void onExecuteBegin() const = 0;
    virtual void onExecuteEnd() const = 0;

    // 内存管理
    virtual MemObj* onAcquire(const Tensor* tensor, StorageType storageType) = 0;
    virtual bool onClearBuffer() = 0;
    virtual void onCopyBuffer(const Tensor* srcTensor, const Tensor* dstTensor) const = 0;
    
    // 把数据映射到指定后端 / 从指定后端映射回数据
    virtual void* onMapTensor(Tensor::MapType mtype, Tensor::DimensionType dtype, const Tensor* srcTensor);
    virtual bool onUnmapTensor(Tensor::MapType mtype, Tensor::DimensionType dtype, const Tensor* dstTensor, void* mapPtr);
    
    // 利用backend后端传递指针, llm中主要传递kvmeta(kv cache有关信息)
    void setMetaPtr(void* ptr);

private:
    // 枚举类 记录执行后端 如cpu opencl为: MNN_FORWARD_CPU MNN_FORWARD_OPENCL
    const MNNForwardType mType;
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

class Execution : public NonCopyable {
public:
    // 根据输入seq_len长度变化 调整执行时用到的值
    virtual ErrorCode onResize(const std::vector<Tensor*>& inputs,
                               const std::vector<Tensor*>& outputs);

    // 执行计算
    virtual ErrorCode onExecute(const std::vector<Tensor*>& inputs,
                                const std::vector<Tensor*>& outputs) = 0;

    // 克隆执行器 这在llm推理中很有用 例如prefill和decode之间的seq_len输入长度不一样 就会触发克隆 mnn默认保存不同输入长度的计算图
    // 这里通常不克隆权重, 权重在Op* op中 这里会通过指针引用
    virtual bool onClone(Backend* bn, const Op* op, Execution** dst);

    // 获取后端 例如在CPUAttention中利用这里获取kvmeta信息(mMeta = (KVMeta*)(backend->getMetaPtr());)
    Backend* backend() const;
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

struct CoreFunctions {
    // CPU 特性标志
    bool supportFp16arith = false;  // 支持 FP16 算术运算
    bool supportSDot = false;       // 支持 ARM S-Dot 指令
    bool supportI8mm = false;       // 支持 ARM I8MM 指令
    bool supportSME2 = false;       // 支持 ARM SME2 指令
    int  smeCoreNumber = 0;         // SME2 核心数量

    // 矩阵乘法相关函数
    void(*MNNGetMatMulPackMode)(int* eP, int *lP, int* hP);
    void(*MNNPackC4ForMatMul_A)(float* destOrigin, float const** sourceGroup,
                                 const int32_t* info, const int32_t* el);
    void(*MNNPackForMatMul_B)(float* dest, const float* source,
                              size_t h, size_t kernelsize, size_t ic, bool transpose);
    void(*MNNPackedMatMul)(float* C, const float* A, const float* B,
                           const size_t* parameter, const float* postParameters,
                           const float* bias, const float* k, const float* b);
    void(*MNNPackedMatMulRemain)(float* C, const float* A, const float* B,
                                 size_t eSize, const size_t* parameter,
                                 const float* postParameters, const float* bias,
                                 const float* k, const float* b);
    // 激活函数
    void(*MNNReluInt8)(int8_t* dst, const int8_t* src, size_t size, ssize_t zeroPoint);
    void(*MNNHardSwish)(float* dst, const float* src, size_t size);
    void(*MNNGelu)(float* dst, const float* src, size_t size, float* parameters);
    const float *beta, float epsilon, size_t size, bool RMSNorm);

    // 其他函数...
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36

// 在source/backend/cpu/CPUBackend.cpp 中创建runtime时 会进入到不同指令集的初始化注册过程
{
#ifdef MNN_SUPPORT_BF16
extern void registerBF16Backend();
#endif
#ifdef ENABLE_ARMV82
extern void registerArm82RuntimeCreator();
#endif
void registerCPURuntimeCreator() {
    MNNCoreFunctionInit();
    CPUBackend::initCreatorMap();
    registerCPUOps();
#ifdef MNN_SUPPORT_BF16
    registerBF16Backend();
#endif
#ifdef MNN_USE_ARMV82
    registerArm82RuntimeCreator();
#endif
    // TODO: Merge _initCoreFunction MNNFunctionInit and cpuinfo_arm_init
    MNNInsertExtraRuntimeCreator(MNN_FORWARD_CPU, new CPURuntimeCreator);
}

// 下面以armv82的初始化为例 在Arm82Functions::init()中
Arm82Functions::init(){
    
    // 部分特殊优化的底层代码
    FUNC_PTR_ASSIGN(gInstance->MNNPackC4ForMatMul_A, Arm82MNNPackForMatMul_A);
    /* 在声明中 这个代码是外部导入的, 路径位置source/backend/arm82/asm/arm64/Arm82MNNPackForMatMul_A.S
    extern "C" {
    // (UP_DIV(l,8), e, 8) -> (UP_DIV(e,eP), l, eP)
        void Arm82MNNPackForMatMul_A(float* destOrigin, float const** sourceGroup, const int32_t* info, const int32_t* el);
    }
    */
    
    // 其它代码...
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

// source/backend/cpu/CPUAttention.cpp
ErrorCode CPUAttention::onExecute(const std::vector<Tensor*>& inputs, const std::vector<Tensor*>& outputs) {
    auto gcore  = static_cast<CPUBackend *>(backend())->functions();
    auto core   = static_cast<CPUBackend*>(backend())->int8Functions();
    // 其它代码...
 	
    // 调用后端的core functions终端矩阵乘法
    gcore->MNNPackedMatMul(...);
    
    // 其它代码...
}