MNN 源码阅读

邢越峰
396 阅读5分钟

MNN是开源边缘AI计算框架的一种，内部有很多值得学习的设计理念和代码实现。这个章节是我从工程实现角度走读源码的笔记。

文章部分内容利用GPT做过整理润色

能力架构图

image.png

备注：

  • Pre-Inference:
    • Geometry Compute: 处理几何计算，可能涉及到输入数据的预处理。
    • Semi-Auto Search: 半自动搜索，可能用于优化模型结构或参数。
    • Resource Manage: 资源管理，确保计算资源的有效利用。
  • Frontend （Algorithm Level Optimize）
    • Strassen Matrix Multiplication: 一种高效的矩阵乘法算法，减少计算复杂度。
    • Winograd Convolution/Deconvolution: 优化卷积和反卷积操作，减少计算量。
    • Piecewise Linearization: 将非线性函数分段线性化，简化计算。
    • Low precision (FP16/BF16/Int8): 使用低精度数据类型，减少内存占用和计算时间。
  • Backend (Hardware Level Optimize)
    • CPU:
      • SSE, AVX, NEON, AVX512: 这些是不同的SIMD（单指令多数据）指令集，用于加速计算。
    • NPU:
      • CoreML, HIAI: 针对特定硬件（如苹果的NPU或华为的AI处理器）进行优化。
    • GPU:
      • OpenCL, Vulkan, CUDA, Metal: 这些是用于GPU编程的框架和API，支持并行计算。

运行时架构图

这个架构图主要描述 MNN 应用层提供的功能和接口，不属于推理框架部分，属于应用层拓展。

框架

  • 执行路径：模型转换 -> 模型加载 -> 预推理 -> 计算核选择 -> 后端执行 -> 输出结果
  • 功能架构概览：
    • 功能模块分布：关键技术点见上文框图，但其中没有描述控制流，更像是技术力总结
    • 功能的层次和实现概览：
      • 配置+接口 ：提供统一的配置接口，允许用户根据具体需求配置不同的推理参数和后端选项。
      • frontend（构建、转换、模型压缩）支持：
        • 离线转换工具：模型压缩、模型转换是独立工具，代码位置 /tools。模型转换结果是 MNN 自有格式 .mnn。其中 PTQ 由 CPP 代码支持，剪枝由 python+tensorflow1.x 实现
        • 推理调度：计算核运行调度、混合精度支持、多 backend 选择。
      • backend（硬件适配和计算加速）支持：
        • 模型加载：通过预推理选择计算核优化方案（依赖具体 backend）、内存（包括 backend 上） 预分配和复用逻辑生成。
        • 计算优化：计算核在 backend 上的优化。
    • 应用组件：
      • 多语言的 API 调用
      • 应用层功能拓展，如 CV 类型前处理、后处理
  • 软件结构的关键点：
    • 控制流：
      • 调度相关：
        • Interpreter：持有模型的计算图结构和相关信息。
        • Session：持有运行时数据，如输入输出张量。支持同步和异步两种接口。调用异构设备同步接口可能需要再调用 waitdone 类型接口。
          • ScheduleConfig：配置Session的前端优化力度和后端选择。默认情况下，使用CPU进行运行。
      • Backend 相关：
        • BackendConfig：配置内存管理（memory）、功耗管理（power）和计算精度（precision）。
          • Power 配置：支持不同的内存分配和绑核等策略，实现功耗配置。
          • Precision 配置：支持不同的计算精度，如FP32和INT8，通过配置可以选择最适合的精度模式。
        • Runtime：后端的抽象层，持有平台相关的资源和配置，可以在多个Session之间共享。
      • 内存相关：
        • 数据结构：Tensor 是数据传输抽象，具备多 backend 数据类型持有的能力。CV::ImageProcess 支持对 Matrix 的 CV 操作。
        • 内存复用：非激活值的张量（tensor）被复用，除非设置了saveTensors标志，以防止不必要的内存分配和释放。
        • inptut zerocopy：支持 inputtensor 的内存虚拟地址指针获取
        • 输入尺寸和 batch （影响 session 的内存分配）：对 session 执行 resizeSession 和 resizeTensor，batch 数由模型转换期确定。
      • 运行调度可配置：
        • 单 net 多 session 实现模型参数共享的多路运行；（存疑，是否存在类似 trt 的多 profile 的 context）
        • 多 Runtime 多 net/session 实现 backend 在 session 间共享；
        • 单 session 多 scheduler 实现推理过程多 backend 或 backendconfig；
    • 并发与并行：
      • Backend：
        • CPUThreadPool source/backend/cpu/ThreadPool.hpp：具体见下子小节。
      • Frontend：多 session 串行流水；多线程 session 并行；待实验；
    • 数据流 & 数据结构
      • session 运行的 tensor：session 包含多个 pipeline（存疑，为什么多个 pipeline），pipeline 包含 PipelineInfo，pipelineinfo 持有 tensor；
      • net 权值的内存：一般采用构建期静态内存分配，也支持 mmap，具体见子章节。
      • 如何共享 tensor：通过多个 Tensor 与 MemNode 绑定实现，具体见子章节。
    • 内存管理与缓存
      • 内存分配（CPU、GPU）
        • 计算内存预分配：位于 Backend 和 Runtime 的逻辑。具体见子章节。
      • 内存池化策略
        • CPU 内存池化：支持用户态的内存预分配和碎片管理。见各种 Allocator 相关实现。具体见子章节。
        • GPU 内存池化：依赖于图形API（如OpenCL、Vulkan）进行内存管理，采用内存池和复用策略减少内存分配次数。
      • 优化手段
    • 图优化
      • 图融合、计算图线性拆分
      • 计算核的方案选择（模型加载期执行），方式是半自动搜索。如拆分计算利用多种 conv 计算优化方式（如 winograd、滑窗） 实现
    • 计算优化
      • 汇编实现计算核
      • 权值内存重排
      • 矩阵计算算法优化，在通用的计算核优化算法上对更细粒度计算过程进行优化

实战操作

开发环境：使用 githubcodespace，2-x86cpu 8G-RAM

代码编译：cmake .. -DMNN_BUILD_TEST=ON -DMNN_BUILD_TOOLS=ON -DMNN_BUILD_QUANTOOLS=ON -DMNN_BUILD_DEMO=ON -DMNN_BUILD_CONVERTER=ON -DMNN_BUILD_BENCHMARK=ON -DMNN_AVX2=OFF

子页面链接

MNN 工程实现代码阅读 - 掘金

avatar
文章
阅读
粉丝