MNN 工程实现代码阅读CPU Backend ThreadPool source/backend/cpu/Thread

父页面

概述

父页面已经从架构实现角度进行模块梳理。这个页面主要记录部分关键模块的代码精度笔记。

CPU Backend ThreadPool

source/backend/cpu/ThreadPool.hpp

一个功能模拟 CUDA/SIMT 的定制线程池，支持单线程线下无锁的任务下发和并发逻辑，通过绑定虚拟线程号实现某种任务的统一发射、运行、同步结束。

// 类型：
    // <任务，同时下发的任务数量>
    // 任务包含入参int，指的是线程数，类似cuda的threadIdx。与CUDA逻辑类似：理论上并发多个线程调用task，
    //。   task函数中根据threadIdx选择要执行的计算内容，然后同步返回。
    typedef std::pair<std::function<void(int)>, int> TASK;
// 接口：
    // 阻塞接口，tasks在全部并发线程中完成后返回。此线程会作为0号线程执行一次task。
    // 全程无锁，通过遍历对应mTask元素中的atomic判断是否全部运行完成
    // 通过yield实现自旋判断完成状态
    static void enqueue(TASK&& task, int index, int threadNumber);

    // 设置若干线程的启动和停止。需要成对调用。线程活跃计数为零时停止运行。
    // 也就是说不同线程的调用者可能使用同一个线程进行任务执行。
    static void active(int threadNumber);
    static void deactive(int threadNumber);

    // 每个线程调用enqueue前应该获取一个独立的mTask元素的index
    // 这个index能够让enqueue调用不需要去其他线程竞争生产消费队列（虽然这个队列长度是1）
    // index还能为每个线程单独维护一份线程任务执行状态标记，当这一任务在全部线程中完成后即刻返回
    static int acquireWorkIndex();
    static void releaseWorkIndex(int index);

    // 初始化、释放 ThreadPool 单例和线程池并发数量
    static int init(int number);
    static void destroy();

// 成员：

    // 私有单例
    static ThreadPool* gInstance;
    // 初始化变量，构造线程池中的SIMT逻辑
    // 线程池运行逻辑：在mActiveCount为true时通过遍历mTasks中的atomic_bool感知是否有任务到来并执行
    //。  执行完任务yield然后自旋等待，直到mActiveCount为false进入wait状态
    ThreadPool(int number = 0);
    ~ThreadPool();

    // 牛马数量
    int mNumberThread            = 0;
    // 干活牛马
    std::vector<std::thread> mWorkers;
    // 牛马是否干活的标志
    std::atomic<bool> mStop = {false};
    // 牛马是否等待干活的通知
    std::condition_variable mCondition;
    // 1. 牛马通知的互斥量；2. mTasksAvailable数组的保护
    std::mutex mQueueMutex;
        
    // 每个线程只能持有一个mTasks元素，atomic_bool是SIMT中并发任务的完成状态
    std::vector<std::pair<TASK, std::vector<std::atomic_bool*>>> mTasks;
    // mTasks可用状态，与mTasks分离，能够避免多线程获取此状态时锁住mTasks
    std::vector<bool> mTaskAvailable;

    // 活跃的SIMT调用者数量；atomic不可移动构造和复制构造，故用指针；
    std::vector<std::atomic_int*> mActiveCount;

CPUBackend

这部分通过是 backend 对 cpu 的实现。此模块主要包含三个类：

CPUBackend 对外封装推理调度接口，对内执行内存申请、输入尺寸修改、推理调用，持有资源（如 cpuruntime）。
CPUMemObj 自动释放内存的封装
CPURuntime 主要封装池化线程和池化内存两种重要的功能，还支持内核绑定、功耗调节、异步调度任务等待功能。

子模块功能：

**CPURuntime** 功能：
- 静态内存分配器 (**mStaticAllocator**)：用于管理模型的静态内存，包括 weight、feature、kvcache。可通过配置 hint 中的weightMemoryPath 实现 MmapAllocator 映射到文件上。默认使用 malloc 并进行内存地址对齐。
- 静态内存分配器 cache（mStaticAllocatorCache）：如果有映射mStaticAllocator 到文件，那么mStaticAllocatorCache 可用于后续分配EagerBufferAllocator。
- 动态内存分配器 (**mDynamic**)：用于管理推理过程中动态变化的内存需求，如中间计算结果。
- 动态内存映射 (**mDynamicMmap**)：从文件系统映射内存。可通过配置 hint 中的 midMemoryPath启用。
- 内存分配器的创建与管理（**createDynamicBufferAlloctor**）：动态内存分配优先使用 DeferBufferAllocator通过mDynamic或mDynamicMmap分配（关于 Defer 分配，见下文），其次使用EagerBufferAllocator通过mStaticAllocator或者mStaticAllocatorCache分配。
**CPUBackend** 功能：
- 具体模块
 - mDmaInfo 动态分配器（Dynamic Memory Allocator）：通过CPURuntime的createDynamicBufferAlloctor创建。
 - mStaticAllocator 静态内存分配器：与CPURuntime的分配器保持一致。
- 调用过程
 - 构造函数通过依赖注入的方式从CPURuntime中初始化内存分配器
 - allocBuffer接口支持对入参Tensor数据结构进行内存分配，支持 STATIC、DYNAMIC、DYNAMIC_SEPERATE 类型内存申请。Tensor是计算核直接使用的数据结构。
 - onClearBuffer等接口支持对其持有的 runtime 中全部内存做批量操作。往往在充值流程中调用。
 - onMapTensor``onCopyBuffer等函数支持对 host、device 内存的映射、深拷贝操作。外部可能不需要清楚底层数据结构，直接调用 backend 的接口能实现大多数内存操作。

功能总结：

CPURuntime和CPUBackend在基类Runtime``Backend基础上实现 CPU 上不同层级的 device 特性封装。runtime 层面直接配置和操作多线程并发、内存池化、核心绑定的功能。backend 层面封装对外接口，提供 pipeline 视角的静态动态内存的控制、pipeline 级别调度。通过这些封装，实现 pipeline 视角在业务逻辑上的 backend 抽象。

BufferAllocator 和 Tensor

这部分主要实现内存管理和数据结构的抽象，能够支持计算核、backend、runtime 的内存操作和优化策略，通过内部封装的池化机制优化内存碎片化、提升内存复用。

Allocator 用于申请和释放 MemChunk：

Allocator（抽象基类）：支持onAlloc和onRelease接口
MmapAllocator（派生类）：每次调用onAlloc都新创建名称不同的文件用于映射
RecurseAllocator（派生类）：持有另一个Allocator，辅助类型的实现
DefaultAllocator（派生类）：通过malloc 申请内存并强制自行 64 位对齐，内存封装在MemChunk中。

BufferAllocator 的类型和特性：

BufferAllocator（抽象基类） ：支持alloc和free接口实现内存申请释放，支持多线程共享的barrierBegin等接口。compute接口实现内存到Tensor的一次性重申请和重新映射，用于推理引擎 resize 输入内存，看起来和原本的架构体系不太搭，应该是后期为支持 resize 输入尺寸的同时保证池化内存的内存连续性而嫁接的特性。
EagerBufferAllocator（派生类）：通过FREELIST管理chunk，模式与ptmalloc相似且简化。alloc和free会立即触发内存分配或chunk分割或合并。
DeferBufferAllocator（派生类）：通过简化的FREELIST管理预期要持有的chunksize，在alloc和free时不会立即触发内存分配，但是会触发chunk切分或合并。在compute调用中一次性实现全部chunk的内存分配。

Tensor 的功能：（类似 CV::Mat）

描述 tensor 的属性，如 dim、batch、memtype、size 等，支持相关函数
持有内存的引用，包含异构设备内存的访问方式，具体是 host 指针和对应的 device 内存地址。