昇腾 GEMM 类算子执行流程
一、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 块，确定并行策略
- 参数解析：读取 A/B/C 维度、α/β、数据类型（FP16/FP32/INT8）
- Tiling 计算：
- M 方向：TileM=64（A 行）
- N 方向：TileN=64（B 列）
- K 方向：TileK=16（公共维度）
- 每个 AI Core 负责 C 的 TileM×TileN 子块
- 内存分配：Host/Device 内存、Stream、事件、同步信号
- 下发任务：将 Tiling 参数、内存地址、Kernel 函数下发至 Device
3.2 Device 侧：数据搬运（GM→UB）
核心：双缓冲、异步预取、数据重排
- 初始化：L0A/L0B、UB、累加器清零
- 双缓冲：UB 分 2 个缓冲区（buf0/buf1），计算当前块时预取下一块
- 数据搬运：
- GM→L2→L1→UB：异步 DMA（不阻塞 Cube）
- 重排：A/B 转置、对齐、填充（Padding）
- 同步：数据搬运完成后触发 Barrier
3.3 核心计算：Cube 矩阵乘（AIC）
核心：16×16×16 Cube 脉动计算、累加、流水线并行
- Cube 指令：调用 mmad（矩阵乘累加），执行 A×B 分块
- 计算流程：
- 加载 A [TileM, TileK] 到 L0A
- 加载 B [TileK, TileN] 到 L0B
- Cube 并行计算：C_tile = A_tile × B_tile
- 累加：C = α×C_tile + β×C
- 双缓冲流水：
- 计算 buf0 → 预取 buf1
- 计算 buf1 → 预取 buf0
- 同步：K 维度遍历完成，所有 Tile 计算结束
3.4 向量 / 标量处理（AIV）
核心：Padding、偏置、激活、精度转换
- Padding：非 16 倍数维度补零
- 偏置加法：C = C + bias
- 激活函数：ReLU、GELU、Sigmoid
- 精度转换：FP16→FP32、INT8→FP16
3.5 结果写回（UB→GM）
核心：结果合并、同步、写回全局内存
- 结果合并：将各 AI Core 的 Tile 结果拼接为完整 C 矩阵
- 同步：所有计算完成、写回完成
- 释放资源：释放 UB、L0、L1 缓存
3.6 Host 侧：结果获取与后处理
- 数据拷贝：Device→Host 内存
- 结果校验：维度、数值、精度检查
- 资源释放：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; }五、性能优化关键技术
- Tiling 策略：16×16×16 对齐、2D 分块、UB 利用率最大化
- 双缓冲：计算与数据搬运并行，Cube 利用率 100%
- 数据重排：转置、对齐、填充，减少非对齐访问
- 算子融合：GEMM + 偏置 + 激活，减少内存访问
- 多核并行：多 AI Core 并行计算，加速比接近核数
六、总结
昇腾 GEMM 算子通过达芬奇 Cube 单元、多级缓存、分块 Tiling、双缓冲、流水线并行，实现从 Host 到 Device、从数据搬运到核心计算、从结果写回的全链路优化。其执行流程严格遵循硬件特性、内存层级、并行计算三大原则，性能可达理论峰值 90% 以上，是大模型、AI 训练 / 推理的核心算力支撑。