4、MLU内存管理与数据布局：全局内存、共享内存、寄存器文件、数据对齐与内存访问优化

寒武纪MLU（Machine Learning Unit）作为面向深度学习训练与推理的专用加速器，其存储层次结构直接决定了算力能否被有效释放。与GPU类似但又不尽相同，MLU采用“全局内存（Global Memory）—共享内存（Shared Memory）—寄存器文件（Register File）”三级架构，配合片上SRAM与片外HBM/DDR的协同工作。本章将深入剖析MLU的内存模型、数据布局规则以及访存优化方法论，帮助开发者在编写MLU Kernel时，从“能跑”走向“跑得快”。

上图展示了MLU存储层次的核心数据流。全局内存是唯一可被主机端（Host）直接访问的存储空间，容量大但延迟高；共享内存是片上高速暂存区，由软件显式管理，用于数据复用与线程间通信；寄存器文件则完全私有，访问延迟仅1个时钟周期。理解这三者的带宽与延迟差异，是进行内存优化设计的前提。

4.1 核心原理：MLU存储层次与访存路径

MLU的存储架构设计遵循“算力越强，数据供给越要跟上”的原则。以寒武纪MLU370系列为例，其全局内存采用HBM2e，理论带宽可达1.2TB/s，但单个线程的随机访问延迟高达数百周期。若所有数据都直接从全局内存读取，计算单元将大量时间处于等待状态。因此，MLU提供了两级片上缓存/暂存机制：共享内存（SMEM）与寄存器文件（RF）。

共享内存本质上是一块可由开发者完全控制的SRAM，通常被划分为多个Bank（例如32个），每个Bank在同一时钟周期内只能响应一次访问。当多个线程同时访问同一Bank的不同地址时，就会发生Bank冲突，导致访问串行化，性能下降数倍。寄存器文件则与线程一一绑定，编译器负责分配与生命周期管理，开发者通常通过控制线程内局部变量的使用方式来间接影响寄存器利用率。

从编程模型看，MLU的线程层次为“Grid→Block→Thread”，其中Block内的线程共享同一块共享内存。这种设计使得Block内协作成为可能，例如矩阵分块计算中的子块缓存、卷积滑动窗口的数据复用等。理解这一原理后，开发者应建立“数据先入共享内存，再入寄存器，最后参与计算”的思维定式。

4.2 判断方法与操作流程：数据布局与对齐策略

数据布局的优劣直接影响内存带宽利用率。MLU硬件要求全局内存访问必须满足128字节对齐，且连续访问长度建议为32字节的整数倍。若数据布局不满足对齐条件，硬件将自动拆分访问事务，导致有效带宽下降50%以上。

以下是标准的数据布局优化流程：

1. 分析访存模式：检查Kernel中每个循环的索引表达式，判断是否存在跨步（Stride）访问或随机访问。
2. 选择布局格式：对于图像或特征图数据，推荐使用NC1HWC0格式（C0=32），将通道维拆分为C1和C0，使得最内层C0维度连续存放，正好匹配MLU的128字节对齐要求。
3. 共享内存Padding：在分配共享内存时，每行末尾额外增加若干字节（例如16字节），使每行的起始地址对齐到不同Bank，避免Bank冲突。
4. 向量化加载：使用MLU提供的向量加载指令（如__memcpy_async或vload），一次加载128位（16字节）数据到寄存器，减少指令数。
5. 双缓冲：在共享内存中分配两个缓冲区，一个用于当前计算，另一个用于DMA预取下一块数据，隐藏全局内存延迟。

4.3 案例拆解：共享内存Bank冲突的定位与消除

我们以一个实际的反向传播（Backward）Kernel为例。初始实现中，每个线程负责计算输出梯度矩阵中的一个元素，梯度矩阵大小为32×32，按行主序存储在共享内存中。线程索引与矩阵索引的映射关系为：row = threadIdx.x / 32; col = threadIdx.x % 32;。

初始现象：Kernel运行时间比预期慢3.2倍，通过性能分析工具（如cnprof）发现共享内存访问的Bank冲突率为87%。

分析过程：由于矩阵行长为32个float（128字节），而共享内存Bank数为32，每行恰好占用32个Bank。当线程束（32个线程）同时访问第0行时，所有线程访问的地址落在不同的Bank上，无冲突。但当线程束访问第1行时，线程0访问地址row=1, col=0，该地址对应的Bank为(1*32+0) % 32 = 0，线程1访问row=1, col=1，Bank为(1*32+1)%32=1……实际上，由于行长为32的整数倍，每一行的起始Bank都相同，导致同一列的不同行映射到同一Bank。当线程束内线程访问不同行但相同列时（例如线程0访问第0行第0列，线程1访问第1行第0列），就会发生32路冲突。

解决方案：在共享内存分配时进行Padding，将矩阵声明为float smem[32][33]，即每行增加1个float的偏移。这样，第1行的起始Bank变为(1*33)%32=1，第2行为(2*33)%32=2，以此类推。修改后，同一线程束内任意两个线程访问的Bank均不相同，冲突率降为0%。

结果对比：Kernel运行时间从12.4ms降至3.8ms，加速比3.26倍，与理论预期一致。

复盘经验：Bank冲突是共享内存优化的首要排查点。在编写Kernel前，应先用公式Bank = (row * stride + col) % 32手工推演一遍线程束的访问模式。若stride是32的倍数，则必然发生冲突。

4.4 实战场景：全局内存的合并访问与DMA异步搬运

在MLU上，全局内存的访问效率高度依赖“合并访问”模式。所谓合并访问，是指线程束内32个线程访问的地址落在同一段连续内存中（通常为128字节或256字节）。若线程束访问的地址分散在不同内存页，硬件将产生多个内存事务，带宽利用率骤降。

以矩阵乘法C[M,N] = A[M,K] * B[K,N]为例，朴素实现中，每个线程计算C的一个元素，需要读取A的一行和B的一列。A的行连续，合并访问良好；但B的列访问跨步为N个float，完全无法合并。优化方案是使用共享内存分块：将A和B分别切成大小为TILE×TILE的子块，先由线程束协作将子块从全局内存加载到共享内存（此时全局内存访问是合并的），再从共享内存中读取数据进行计算。

进一步地，MLU提供了异步DMA指令（如__mlu_async_memcpy），允许全局内存到共享内存的拷贝与计算重叠。操作流程如下：

// 伪代码示意
__shared__ float As[2][TILE][TILE]; // 双缓冲
for (int k0 = 0; k0 < K; k0 += TILE) {
    // 预取下一块到As[1]
    __mlu_async_memcpy(As[1], &A[blockIdx.y*TILE][k0], sizeof(float)*TILE*TILE);
    __mlu_async_wait(); // 等待当前块就绪
    // 计算当前块As[0]
    compute(As[0], Bs[0]);
    // 交换缓冲区
    swap(As[0], As[1]);
}

该模式将全局内存延迟完全隐藏在计算过程中，对于大矩阵乘法，性能可提升40%以上。但需注意，异步拷贝需要显式管理同步点，且缓冲区大小受共享内存容量限制。

4.5 常见误区与总结复盘

在MLU内存优化实践中，开发者常陷入以下误区：

• 误区一：盲目增加共享内存使用量。共享内存是稀缺资源，过度使用会降低Block并行度（Occupancy），导致调度器无法隐藏延迟。应通过cnDeviceGetAttribute查询硬件限制，并计算理论占用率。
• 误区二：忽略寄存器溢出。当线程内局部变量过多时，编译器会将寄存器变量溢出到局部内存（Local Memory，实际位于全局内存），导致性能断崖式下跌。应通过编译选项（如--maxrregcount）控制寄存器数量，或重构代码减少变量生命周期。
• 误区三：认为所有数据都适合放入共享内存。对于一次性读取且无复用的数据（如权重矩阵的偏置项），直接从全局内存读取反而更快，因为省去了DMA搬运和同步开销。
• 误区四：忽视数据对齐的“最后一公里”。即使整体布局正确，若Kernel内部使用reinterpret_cast将字节数组转为float数组，也可能破坏对齐假设。应使用__builtin_assume_aligned或MLU提供的对齐分配函数。

总结复盘：MLU内存优化的核心是“让数据在正确的时间出现在正确的位置”。具体而言，应遵循以下优先级：

1. 优先保证全局内存合并访问（128B对齐、连续）。
2. 其次利用共享内存实现数据复用，并消除Bank冲突。
3. 最后通过寄存器缓存高频访问数据，减少共享内存压力。
4. 使用异步DMA和双缓冲隐藏全局内存延迟。

通过本章的学习，开发者应能够熟练诊断内存瓶颈，并针对不同计算模式选择合适的数据布局与访存策略。在实际项目中，建议使用性能分析工具（如cnmonitor、cnprof）量化各层内存的利用率，以数据驱动优化决策。