一、GEMM 算子概述

GEMM（General Matrix Multiplication，通用矩阵乘法）是深度学习、大模型、科学计算的核心算子，公式为 C = αA×B + βC，占 Transformer、CNN 模型计算量 60%~70%。昇腾达芬奇架构以Cube 单元为核心，通过硬件加速、多级缓存、数据分块、双缓冲、流水线并行，实现 GEMM 极致性能。

二、昇腾硬件架构与 GEMM 适配

2.1 达芬奇 AI Core 核心单元

• Cube 单元：16×16×16 脉动阵列，单周期完成 4096 次乘加（MAC），FP16 算力 256 GFLOPS/core
• Vector 单元：处理向量运算（激活、归一化）
• Scalar 单元：控制流、循环、地址计算
• 存储层级：GM（全局内存）→ L2 Cache → L1 Buffer → UB（统一缓存）→ L0A/L0B（Cube 专用缓存）

2.2 核心约束

• Cube 对齐：输入矩阵需16 字节对齐（M/N/K 维度为 16 倍数）
• UB 容量：单 AI Core UB 约 2MB，需分块（Tiling） 适配
• 内存墙：GM→UB 带宽远低于 Cube 算力，需数据复用、双缓冲、预取

三、GEMM 算子完整执行流程（6 阶段）

3.1 Host 侧：Tiling 与参数准备

核心：将大矩阵切分为适配 UB 的 Tile 块，确定并行策略

1. 参数解析：读取 A/B/C 维度、α/β、数据类型（FP16/FP32/INT8）
2. Tiling 计算：
   • M 方向：TileM=64（A 行）
   • N 方向：TileN=64（B 列）
   • K 方向：TileK=16（公共维度）
   • 每个 AI Core 负责 C 的 TileM×TileN 子块
3. 内存分配：Host/Device 内存、Stream、事件、同步信号
4. 下发任务：将 Tiling 参数、内存地址、Kernel 函数下发至 Device

3.2 Device 侧：数据搬运（GM→UB）

核心：双缓冲、异步预取、数据重排

1. 初始化：L0A/L0B、UB、累加器清零
2. 双缓冲：UB 分 2 个缓冲区（buf0/buf1），计算当前块时预取下一块
3. 数据搬运：
   • GM→L2→L1→UB：异步 DMA（不阻塞 Cube）
   • 重排：A/B 转置、对齐、填充（Padding）
4. 同步：数据搬运完成后触发 Barrier

3.3 核心计算：Cube 矩阵乘（AIC）

核心：16×16×16 Cube 脉动计算、累加、流水线并行

1. Cube 指令：调用 mmad（矩阵乘累加），执行 A×B 分块
2. 计算流程：
   • 加载 A [TileM, TileK] 到 L0A
   • 加载 B [TileK, TileN] 到 L0B
   • Cube 并行计算：C_tile = A_tile × B_tile
   • 累加：C = α×C_tile + β×C
3. 双缓冲流水：
   • 计算 buf0 → 预取 buf1
   • 计算 buf1 → 预取 buf0
4. 同步：K 维度遍历完成，所有 Tile 计算结束

3.4 向量 / 标量处理（AIV）

核心：Padding、偏置、激活、精度转换

• Padding：非 16 倍数维度补零
• 偏置加法：C = C + bias
• 激活函数：ReLU、GELU、Sigmoid
• 精度转换：FP16→FP32、INT8→FP16

3.5 结果写回（UB→GM）

核心：结果合并、同步、写回全局内存

1. 结果合并：将各 AI Core 的 Tile 结果拼接为完整 C 矩阵
2. 同步：所有计算完成、写回完成
3. 释放资源：释放 UB、L0、L1 缓存

3.6 Host 侧：结果获取与后处理

1. 数据拷贝：Device→Host 内存
2. 结果校验：维度、数值、精度检查
3. 资源释放：Stream、事件、内存

四、Ascend C 代码实现

4.1 Tiling 与 Kernel 入口

// GEMM Tiling定义
BEGIN_TILING_DATA_DEF(GemmTiling)
  TILING_DATA_FIELD_DEF(uint32_t, tileM);
  TILING_DATA_FIELD_DEF(uint32_t, tileN);
  TILING_DATA_FIELD_DEF(uint32_t, tileK);
  TILING_DATA_FIELD_DEF(uint32_t, numK);
END_TILING_DATA_DEF;

// Kernel入口（Device侧）
__global__ __aicore__ void gemm_kernel(
  GM_ADDR gmA, GM_ADDR gmB, GM_ADDR gmC,
  const GemmTiling tiling, float alpha, float beta
) {
  // 1. 初始化UB、L0、累加器
  LocalTensor<half> localA, localB, localC;
  localA.SetBuffer((half*)UB_BASE, tiling.tileM, tiling.tileK);
  localB.SetBuffer((half*)UB_BASE + 1024, tiling.tileK, tiling.tileN);
  localC.SetBuffer((half*)UB_BASE + 2048, tiling.tileM, tiling.tileN);

  // 2. 双缓冲预取
  uint32_t bufIdx = 0;
  for (uint32_t k = 0; k < tiling.numK; k++) {
    // 3. 异步搬运A/B到UB
    DataCopy(localA, gmA + k*tiling.tileK, tiling.tileM, tiling.tileK);
    DataCopy(localB, gmB + k*tiling.tileK*tiling.tileN, tiling.tileK, tiling.tileN);
    Sync();

    // 4. Cube矩阵乘（核心）
    mmad(localC, localA, localB, alpha, beta);

    // 5. 双缓冲切换
    bufIdx = 1 - bufIdx;
  }

  // 6. 结果写回GM
  DataCopy(gmC, localC, tiling.tileM, tiling.tileN);
}

4.2 双缓冲与流水线优化

// 双缓冲（预取+计算并行）
uint32_t bufIdx = 0;
for (uint32_t k = 0; k < tiling.numK; k++) {
  // 异步预取下一块（不阻塞Cube）
  if (k + 1 < tiling.numK) {
    DataCopyAsync(localA[1-bufIdx], gmA + (k+1)*tiling.tileK);
    DataCopyAsync(localB[1-bufIdx], gmB + (k+1)*tiling.tileK*tiling.tileN);
  }

  // 计算当前块
  mmad(localC, localA[bufIdx], localB[bufIdx]);

  // 切换缓冲区
  bufIdx = 1 - bufIdx;
}

五、性能优化关键技术

1. Tiling 策略：16×16×16 对齐、2D 分块、UB 利用率最大化
2. 双缓冲：计算与数据搬运并行，Cube 利用率 100%
3. 数据重排：转置、对齐、填充，减少非对齐访问
4. 算子融合：GEMM + 偏置 + 激活，减少内存访问
5. 多核并行：多 AI Core 并行计算，加速比接近核数

六、总结

昇腾 GEMM 算子通过达芬奇 Cube 单元、多级缓存、分块 Tiling、双缓冲、流水线并行，实现从 Host 到 Device、从数据搬运到核心计算、从结果写回的全链路优化。其执行流程严格遵循硬件特性、内存层级、并行计算三大原则，性能可达理论峰值 90% 以上，是大模型、AI 训练 / 推理的核心算力支撑。