XiaoMi/mace项目中的内存布局设计解析

XiaoMi/mace项目中的内存布局设计解析

mace MACE is a deep learning inference framework optimized for mobile heterogeneous computing platforms. 项目地址: https://gitcode.com/gh_mirrors/ma/mace

概述

在深度学习推理框架中，内存布局设计是影响性能的关键因素之一。XiaoMi/mace项目针对CPU和GPU运行时采用了不同的内存布局策略，以充分利用硬件特性，提高计算效率。本文将详细解析mace框架中的内存布局设计原理。

CPU运行时内存布局

CPU运行时采用传统的张量缓冲区布局方式，主要特点如下：

1. 中间输入/输出张量：采用NCHW布局

   • N：批次大小
   • C：通道数
   • H：高度
   • W：宽度

2. 卷积滤波器：采用OIHW布局

   • O：输出通道数
   • I：输入通道数
   • H：高度
   • W：宽度

3. 深度可分离卷积滤波器：采用MIHW布局

   • M：乘数（通常为1）
   • I：输入通道数
   • H：高度
   • W：宽度

4. 一维参数：简单的线性布局（W）

这种布局方式与主流深度学习框架一致，便于CPU进行连续内存访问，提高缓存命中率。

GPU运行时内存布局

GPU运行时基于OpenCL实现，采用2D图像(CL_RGBA)作为张量存储，要求OpenCL 1.2及以上版本。这种设计充分利用了GPU的纹理内存和并行处理能力。

输入/输出张量布局

输入/输出张量采用NHWC格式存储，但根据不同的使用场景有三种映射方式：

1. 通道优先(Channel-Major)：

   • 默认输入/输出格式
   • 图像尺寸：[W × (C+3)/4, N × H]
   • 每个像素包含4个连续通道的数据

2. 高度优先(Height-Major)：

   • 主要用于Winograd变换和矩阵乘法输出
   • 图像尺寸：[W × C, N × (H+3)/4]
   • 每个像素包含4个连续高度的数据

3. 宽度优先(Width-Major)：

   • 当前未使用
   • 图像尺寸：[(W+3)/4 × C, N × H]
   • 每个像素包含4个连续宽度的数据

滤波器张量布局

1. 常规卷积滤波器：

   • 缓冲区格式：OIHW
   • 图像尺寸：[I, (O+3)/4 × W × H]
   • 每个像素包含4个输出通道的数据

2. 深度可分离卷积滤波器：

   • 缓冲区格式：MIHW
   • 图像尺寸：[H × W × M, (I+3)/4]
   • 每个像素包含4个输入通道的数据
   • 当前仅支持乘数为1的情况

一维参数张量布局

主要用于偏置等一维参数：

• 缓冲区格式：W
• 图像尺寸：[(W+3)/4, 1]
• 每个像素包含4个连续的数据元素

性能优化考量

mace的内存布局设计体现了以下优化思想：

1. 数据局部性：将经常一起访问的数据放在相邻位置，提高缓存利用率
2. 向量化处理：利用RGBA通道同时处理4个数据元素，提高并行度
3. 内存对齐：通过(+3)/4确保数据对齐，提高访问效率
4. 硬件适配：针对CPU和GPU的不同特性采用不同的布局策略

实际应用建议

1. 在自定义算子开发时，应遵循框架的内存布局规范
2. 数据预处理阶段应考虑目标设备的内存布局特点
3. 性能调优时可尝试不同的内存布局方式
4. 深度可分离卷积的实现需要注意乘数限制

理解这些内存布局设计原理，有助于开发者更好地使用mace框架，并在必要时进行定制化优化。

mace MACE is a deep learning inference framework optimized for mobile heterogeneous computing platforms. 项目地址: https://gitcode.com/gh_mirrors/ma/mace