image

国产GPGPU踩坑 | 曦云C500 HGEMM算子优化

该话题被推 实用技术 AI 其它
230
Aoikajitsu
Aoikajitsu

233

曦云 C500 是沐曦（MetaX）于 2022 年 发布的旗舰级通用计算 GPU。该产品基于沐曦自主研发的“曦云”架构设计，旨在为科学计算、人工智能训练及推理提供强大的国产算力解决方案。曦云 C500 采用自研 XCORE 1.0 架构及指令集，配备标量、矢量和张量计算单元，支持多种混合精度计算，搭载 64GB HBM2e 显存与 7 个高速 MetaXLink 互连接口，可实现 2 卡至 64 卡多种互连拓扑，具备国内稀缺的高带宽、超多卡互连能力；依托自研 MXMACA 软件栈，适配主流算法框架、运算库等工具，编程接口在 API 层面高度兼容 GPU 行业国际主流 CUDA 生态。

针对GEMM数据搬运的说法均默认为K major情况下，其他情况可以以此类推

机器整体上类似A卡，不过这是中途AI大人告诉我的，我对A卡并没有系统认知，可能也不太一样

编程难度： 国产灵车≫\ggCUDA

鲲 Galgame 表情包 \[1\] - 16

算力配置

由于优化的HGEMM,只关注FP16算力，实际上我没查到具体算力，这个机器分了PCIe版和OAM版，二者频率不一样，总之这个机器对标的是A100，大概是240TFLOPS的水平，后面的microbenchmark会对此有更深理解，由于是GEMM，带宽无所谓了，正常机器的配比根本不会bound到带宽上

每AP(SM) smem大小是64KB，register file是512KB，warp(wave) size为64，每线程最大寄存器数为255

NVIDIA Ampere风格GEMM

Ampere风格GEMM理想优化目标如下，图为CTA视角：

(图上的memory并不是真的表达mainloop没有memory操作，更多只是展示这个阶段无compute吞吐)

prologue负责启动软件流水线，epilogue负责写回数据到HBM，这两段不需要计算单元参与，剩下的全部时间都希望能够吃满计算单元以达到最佳性能

具体的流水线如下：

在这种设计下mainloop段保持计算单元持续在工作，能吃满计算吞吐，而这依赖对mma operand的prefetch，所以对于寄存器只需要2 buffer就足够了，一份在用，一份装下一次使用的；对于smem，stage数需要tuning，如图所示，如果wait的时间过长会导致mma工作暂停，因延迟导致没有最大化mma吞吐，运行时间变长，性能变差，更多的stage数只能让in flight的指令数更多，不能据此推测memory system的具体情况，只是希望让更多指令in flight去尽可能隐藏延迟，并非越多越好，不过1到2是质变，意义在于生产和消费overlap

吐槽一句：其实我个人感觉Ampere风格的GEMM比Hopper及之后的GEMM更难分析(对于mainloop的设计)，Ampere风格每个线程既负责搬运数据又负责发起计算还有繁琐的控制逻辑，这么多东西挤在同一条指令流上还要考虑吃满计算单元吞吐，分析起来很繁琐，而且还有warp切换去隐藏延迟提高吞吐，这个就更是分析不了的东西了，倒是Hopper那种异步设计，搬运和计算的逻辑由不同warp负责，两者是完全不同的指令流，谁都不会堵塞谁，在算力膨胀的今天，就算stall了1cycle也损失了不少flops

会用到的builtin函数

cpp
#pragma once

#include <maca_fp16.h>
#include <__clang_maca_vector_types.h>
#include <stddef.h>

// 资料来源：MetaX Developer Documentation
// https://developer.metax-tech.com/doc

// Thin C++ wrappers for the MXC builtins.
// The MACA toolchain provides __NATIVE_VECTOR__ through its vector-types
// header. Keep this wrapper on that official spelling so it works with mxcc.

using v1u32 = __NATIVE_VECTOR__(1, unsigned);
using v2u32 = __NATIVE_VECTOR__(2, unsigned);
using v4u32 = __NATIVE_VECTOR__(4, unsigned);
using v4f16 = __NATIVE_VECTOR__(4, __fp16);
using v4f32 = __NATIVE_VECTOR__(4, float);

namespace mxc {

// The immediate-control operands of MXC memory/barrier builtins must remain
// compile-time constants. Encode them as non-type template parameters rather
// than ordinary inline-function arguments.

// Load 128 bits from global memory into shared memory.
// The return value is a synchronization flag, not loaded data.
template <int Offset, size_t Mask = 0, bool Ret0En = true,
          bool SaddrFlag = true, bool PredNeg = false, bool IsAsync = true>
__device__ __forceinline__ v4u32 ldg_b128_bsm(void* shared_addr,
                                               void* global_addr) {
  return __builtin_mxc_ldg_b128_bsm(shared_addr, global_addr, Offset, Mask,
                                    Ret0En, SaddrFlag, PredNeg, IsAsync);
}

// F16 16x16x16 matrix multiply-accumulate: D = A * B + C.
__device__ __forceinline__ v4f32 mma_16x16x16f16(v4f16 a, v4f16 b, v4f32 c) {
  return __builtin_mxc_mma_16x16x16f16(a, b, c);
}

// Wait until the selected outstanding global/shared-memory operations reach
// their target counts. flag == 0 waits for all preceding operations.
template <unsigned Flag>
__device__ __forceinline__ void arrive() {
  __builtin_mxc_arrive(Flag);
}

// Global-memory queue fence. gvmcnt is valid in [0, 63].
template <unsigned Gvmcnt>
__device__ __forceinline__ void arrive_gvmcnt() {
  __builtin_mxc_arrive_gvmcnt(Gvmcnt);
}

// Shared-memory queue fence. bsmcnt is valid in [0, 15].
template <unsigned Bsmcnt>
__device__ __forceinline__ void arrive_bsmcnt() {
  __builtin_mxc_arrive_bsmcnt(Bsmcnt);
}

// Instruction-level barrier. Prevents instruction reordering and waits for
// outstanding stores and instruction fetches to return.
__device__ __forceinline__ void barrier_inst() {
  __builtin_mxc_barrier_inst();
}

}  // namespace mxc

其中核心的__builtin_mxc_mma_16x16x16f16的AB矩阵TV Layout如下，而acc则是下图把1x4的矩形换成4x1的矩形

__builtin_mxc_ldg_b128_bsm则是类似cp.async的gmem2smem搬运指令，不过寄存器，而__builtin_mxc_arrive_gvmcnt效果上类似cp.async.wait_group，不过不需要自己cp.async.commit_group，每条指令都是一个group，实际上C500上靠显式屏障指令完成对于可变延迟操作的等待，不自己写的情况下编译器会自行插入相关屏障指令

Microbenchmark

使用clock64测出来，__builtin_mxc_mma_16x16x16f16的latency为16cycle，且throughput和latency一致，不需要多组accumulator去吃满MMA单元的吞吐，这点和NVIDIA机器是不一样的，这个ILP特性后面会被使用到，一个AP(SM)有4个MMA单位，故整AP的吞吐为刚才的数据乘4，而A100的MMA单元吞吐是8cycle，不过shape是16x8x16，2条16x8x16就是16x16x16,就是16cycle，两者的MMA单元性能一致，且一个SM有4个MMA单元，C500有104AP，A100有108SM，两者FP16算力几乎只有频率差异

硬件特点(坑)

一个很大的坑点是——机器的LSU throttle很低，无法支持较多指令in flight，动不动就阻塞指令发射端，带来很大影响的latency，解决方法是减小每次发出请求的粒度，需要更细粒度的流水线设计以overlap并避免指令被阻止发射

鲲 Galgame 表情包 \[1\] - 23

最终流水线设计

每线程最大寄存器是255，根据经验设置accumulator为max register的1/2，也就是128，每个MMA Atom的accumulator为4个寄存器，总共32个MMA Atom，最终设置warp tile为64x128，为了吃满512KB选择开8个warp，设置CTA tile为256x256x64，smem使用量为64KB

由于smem大小的限制，在大tile下甚至是开不出来2 buffer(BK=64)，如果使用register去staging那么就会遇到寄存器紧张的情况，与AMD不同，AMD每线程最大寄存器数量为512，且还有上述的坑点，缩小BK换取2 buffer是无意义的，实际上没办法等到数据加载满一个BMxBK的buffer后才开始计算，得另寻方法让MMA和LDG overlap

放弃NV的思维进一步思考后，发觉到其实消费完一份数据后(搬运到register)就安全的可以发起LDG进行原地更新了，这不需要2 buffer，但需要更细粒度的流水线管理，若反过来思考NV的最佳实践为什么不是消费完就原地更新而是把下一个k tile要使用的数据写到另一个buffer，可以发现其原因之一在于NV的mma指令的吞吐要求多组无关联的accumulator，在n卡上一个k block消费的smem数据形状为BMxMMA K，这个形状和CTA内线程一次搬运数据的形状是对不上的，没办法做到原地更新的要求，一条搬运指令下来搬了X * BK，却计算了BMxMMA K这样的区域，只有等多次计算后多个MMA K铺满BK才能安全进行原地更新，此时整个buffer都消费完了，而假设一次block消费是WarpNumM * MMA M * BK，这就和搬运时的形状匹配了，可以安全原地更新，C500的mma指令ILP特点支持去这样消费数据

其实还是multistage设计，不过每个stage从一个BMxBK的buffer变成了一个WarpNumM * MMA M x BK的区域，gmem2smem部分流水线设计是一致的，剩下的则是一个stage内的smem2rmem和mma之间overlap的设计了，在BK=64的情况下，K上其实2个LDS就读完了，不太需要更细粒度的划分了(也没寄存器压力)，直接全搬进寄存器，由于是优先消费K上的数据，而K上并没有数据复用的特点，启动阶段若只准备1stage能执行的mma数量就太少了，形成的时间窗口可能不太能盖住gmem和smem2rmem的latency，所以初始阶段等了2stage，理想情况下是执行一段mma后wait一下next stage然后load数据到register然后继续执行mma，不过实现上我也没这么表达，目前的写法还是单纯的按需load，符合直觉，按理说编译器应该能合理安排指令的位置，不过就后续我profile出来的结果编译器并没有对此进行优化，不过大概率是相关的屏障指令阻止了指令重排，知道有这么一回事就行了，最终mainloop代码变成了下面这坨，这种手动展开写法在HPC中挺常见的，主要还是不信任编译器

cpp
    XCORE1000_HGEMM_MMA(0, 0, 0, 0);
    XCORE1000_HGEMM_LDG_A(next_k_offset, 0);
    XCORE1000_HGEMM_MMA(0, 0, 0, 1);
    XCORE1000_HGEMM_MMA(0, 0, 1, 0);
    XCORE1000_HGEMM_MMA(0, 0, 1, 1);
    XCORE1000_HGEMM_MMA(1, 0, 0, 0);
    XCORE1000_HGEMM_MMA(1, 0, 0, 1);
    XCORE1000_HGEMM_MMA(1, 0, 1, 0);
    XCORE1000_HGEMM_MMA(1, 0, 1, 1);
    XCORE1000_HGEMM_MMA(0, 1, 0, 0);
    XCORE1000_HGEMM_MMA(0, 1, 0, 1);
    XCORE1000_HGEMM_MMA(0, 1, 1, 0);
    XCORE1000_HGEMM_MMA(0, 1, 1, 1);
    XCORE1000_HGEMM_MMA(1, 1, 0, 0);
    XCORE1000_HGEMM_MMA(1, 1, 0, 1);
    XCORE1000_HGEMM_MMA(1, 1, 1, 0);

    mxc::arrive_gvmcnt<3>();
    mxc::barrier_inst();
    XCORE1000_HGEMM_MMA(1, 1, 1, 1);
    XCORE1000_HGEMM_LDS_A_FRAGMENT(2, 0);
    XCORE1000_HGEMM_MMA(0, 2, 0, 0);
    XCORE1000_HGEMM_LDG_B(next_k_offset, 0);
    XCORE1000_HGEMM_MMA(0, 2, 0, 1);
    XCORE1000_HGEMM_LDS_A_FRAGMENT(2, 1);
    XCORE1000_HGEMM_MMA(0, 2, 1, 0);
    XCORE1000_HGEMM_MMA(0, 2, 1, 1);
    XCORE1000_HGEMM_LDS_B_FRAGMENT(2, 0, 0);
    XCORE1000_HGEMM_MMA(1, 2, 0, 0);
    XCORE1000_HGEMM_MMA(1, 2, 0, 1);
    XCORE1000_HGEMM_LDS_B_FRAGMENT(2, 0, 1);
    XCORE1000_HGEMM_MMA(1, 2, 1, 0);
    XCORE1000_HGEMM_MMA(1, 2, 1, 1);
    XCORE1000_HGEMM_LDS_B_FRAGMENT(2, 1, 0);
    XCORE1000_HGEMM_MMA(0, 3, 0, 0);
    XCORE1000_HGEMM_MMA(0, 3, 0, 1);
    XCORE1000_HGEMM_LDS_B_FRAGMENT(2, 1, 1);
    XCORE1000_HGEMM_MMA(0, 3, 1, 0);
    XCORE1000_HGEMM_MMA(0, 3, 1, 1);
    XCORE1000_HGEMM_MMA(1, 3, 0, 0);
    XCORE1000_HGEMM_MMA(1, 3, 0, 1);
    XCORE1000_HGEMM_MMA(1, 3, 1, 0);
    XCORE1000_HGEMM_MMA(1, 3, 1, 1);
    XCORE1000_HGEMM_MMA(2, 0, 0, 0);
    XCORE1000_HGEMM_LDG_A(next_k_offset, 1);
    XCORE1000_HGEMM_MMA(2, 0, 0, 1);
    XCORE1000_HGEMM_MMA(2, 0, 1, 0);
    XCORE1000_HGEMM_MMA(2, 0, 1, 1);
    XCORE1000_HGEMM_MMA(2, 1, 0, 0);
    XCORE1000_HGEMM_MMA(2, 1, 0, 1);
    XCORE1000_HGEMM_MMA(2, 1, 1, 0);
    XCORE1000_HGEMM_MMA(2, 1, 1, 1);
    XCORE1000_HGEMM_MMA(2, 4, 0, 0);
    XCORE1000_HGEMM_MMA(2, 4, 0, 1);
    XCORE1000_HGEMM_MMA(2, 4, 1, 0);
    XCORE1000_HGEMM_MMA(2, 4, 1, 1);
    XCORE1000_HGEMM_MMA(0, 4, 0, 0);
    XCORE1000_HGEMM_MMA(0, 4, 0, 1);
    XCORE1000_HGEMM_MMA(0, 4, 1, 0);

    mxc::arrive_gvmcnt<3>();
    mxc::barrier_inst();
    XCORE1000_HGEMM_MMA(0, 4, 1, 1);
    XCORE1000_HGEMM_LDS_A_FRAGMENT(3, 0);
    XCORE1000_HGEMM_MMA(2, 5, 0, 0);
    XCORE1000_HGEMM_LDG_B(next_k_offset, 1);
    XCORE1000_HGEMM_MMA(2, 5, 0, 1);
    XCORE1000_HGEMM_LDS_A_FRAGMENT(3, 1);
    XCORE1000_HGEMM_MMA(2, 5, 1, 0);
    XCORE1000_HGEMM_MMA(2, 5, 1, 1);
    XCORE1000_HGEMM_LDS_B_FRAGMENT(3, 0, 0);
    XCORE1000_HGEMM_MMA(0, 5, 0, 0);
    XCORE1000_HGEMM_MMA(0, 5, 0, 1);
    XCORE1000_HGEMM_LDS_B_FRAGMENT(3, 0, 1);
    XCORE1000_HGEMM_MMA(0, 5, 1, 0);
    XCORE1000_HGEMM_MMA(0, 5, 1, 1);
    XCORE1000_HGEMM_LDS_B_FRAGMENT(3, 1, 0);
    XCORE1000_HGEMM_MMA(1, 4, 0, 0);
    XCORE1000_HGEMM_MMA(1, 4, 0, 1);
    XCORE1000_HGEMM_LDS_B_FRAGMENT(3, 1, 1);
    XCORE1000_HGEMM_MMA(1, 4, 1, 0);
    XCORE1000_HGEMM_MMA(1, 4, 1, 1);
    XCORE1000_HGEMM_MMA(1, 5, 0, 0);
    XCORE1000_HGEMM_MMA(1, 5, 0, 1);
    XCORE1000_HGEMM_MMA(1, 5, 1, 0);
    XCORE1000_HGEMM_MMA(1, 5, 1, 1);
    XCORE1000_HGEMM_MMA(0, 6, 0, 0);
    XCORE1000_HGEMM_LDG_A(next_k_offset, 2);
    XCORE1000_HGEMM_MMA(0, 6, 0, 1);
    XCORE1000_HGEMM_MMA(0, 6, 1, 0);
    XCORE1000_HGEMM_MMA(0, 6, 1, 1);
    XCORE1000_HGEMM_MMA(0, 7, 0, 0);
    XCORE1000_HGEMM_MMA(0, 7, 0, 1);
    XCORE1000_HGEMM_MMA(0, 7, 1, 0);
    XCORE1000_HGEMM_MMA(0, 7, 1, 1);
    XCORE1000_HGEMM_MMA(3, 0, 0, 0);
    XCORE1000_HGEMM_MMA(3, 0, 0, 1);
    XCORE1000_HGEMM_MMA(3, 0, 1, 0);
    XCORE1000_HGEMM_MMA(3, 0, 1, 1);
    XCORE1000_HGEMM_MMA(3, 1, 0, 0);
    XCORE1000_HGEMM_MMA(3, 1, 0, 1);
    XCORE1000_HGEMM_MMA(3, 1, 1, 0);

    mxc::arrive_gvmcnt<3>();
    mxc::barrier_inst();
    XCORE1000_HGEMM_MMA(3, 1, 1, 1);
    XCORE1000_HGEMM_LDS_A_FRAGMENT(0, 0);
    XCORE1000_HGEMM_MMA(3, 4, 0, 0);
    XCORE1000_HGEMM_LDG_B(next_k_offset, 2);
    XCORE1000_HGEMM_MMA(3, 4, 0, 1);
    XCORE1000_HGEMM_LDS_A_FRAGMENT(0, 1);
    XCORE1000_HGEMM_MMA(3, 4, 1, 0);
    XCORE1000_HGEMM_MMA(3, 4, 1, 1);
    XCORE1000_HGEMM_LDS_B_FRAGMENT(0, 0, 0);
    XCORE1000_HGEMM_MMA(3, 5, 0, 0);
    XCORE1000_HGEMM_MMA(3, 5, 0, 1);
    XCORE1000_HGEMM_LDS_B_FRAGMENT(0, 0, 1);
    XCORE1000_HGEMM_MMA(3, 5, 1, 0);
    XCORE1000_HGEMM_MMA(3, 5, 1, 1);
    XCORE1000_HGEMM_LDS_B_FRAGMENT(0, 1, 0);
    XCORE1000_HGEMM_MMA(1, 6, 0, 0);
    XCORE1000_HGEMM_MMA(1, 6, 0, 1);
    XCORE1000_HGEMM_LDS_B_FRAGMENT(0, 1, 1);
    XCORE1000_HGEMM_MMA(1, 6, 1, 0);
    XCORE1000_HGEMM_MMA(1, 6, 1, 1);
    XCORE1000_HGEMM_MMA(1, 7, 0, 0);
    XCORE1000_HGEMM_MMA(1, 7, 0, 1);
    XCORE1000_HGEMM_MMA(1, 7, 1, 0);
    XCORE1000_HGEMM_MMA(1, 7, 1, 1);
    XCORE1000_HGEMM_MMA(2, 2, 0, 0);
    XCORE1000_HGEMM_LDG_A(next_k_offset, 3);
    XCORE1000_HGEMM_MMA(2, 2, 0, 1);
    XCORE1000_HGEMM_MMA(2, 2, 1, 0);
    XCORE1000_HGEMM_MMA(2, 2, 1, 1);
    XCORE1000_HGEMM_MMA(3, 2, 0, 0);
    XCORE1000_HGEMM_MMA(3, 2, 0, 1);
    XCORE1000_HGEMM_MMA(3, 2, 1, 0);
    XCORE1000_HGEMM_MMA(3, 2, 1, 1);
    XCORE1000_HGEMM_MMA(2, 3, 0, 0);
    XCORE1000_HGEMM_MMA(2, 3, 0, 1);
    XCORE1000_HGEMM_MMA(2, 3, 1, 0);
    XCORE1000_HGEMM_MMA(2, 3, 1, 1);
    XCORE1000_HGEMM_MMA(3, 3, 0, 0);
    XCORE1000_HGEMM_MMA(3, 3, 0, 1);
    XCORE1000_HGEMM_MMA(3, 3, 1, 0);

    mxc::arrive_gvmcnt<3>();
    mxc::barrier_inst();
    XCORE1000_HGEMM_MMA(3, 3, 1, 1);
    XCORE1000_HGEMM_LDS_A_FRAGMENT(1, 0);
    XCORE1000_HGEMM_MMA(2, 6, 0, 0);
    XCORE1000_HGEMM_LDG_B(next_k_offset, 3);
    XCORE1000_HGEMM_MMA(2, 6, 0, 1);
    XCORE1000_HGEMM_LDS_A_FRAGMENT(1, 1);
    XCORE1000_HGEMM_MMA(2, 6, 1, 0);
    XCORE1000_HGEMM_MMA(2, 6, 1, 1);
    XCORE1000_HGEMM_LDS_B_FRAGMENT(1, 0, 0);
    XCORE1000_HGEMM_MMA(3, 6, 0, 0);
    XCORE1000_HGEMM_MMA(3, 6, 0, 1);
    XCORE1000_HGEMM_LDS_B_FRAGMENT(1, 0, 1);
    XCORE1000_HGEMM_MMA(3, 6, 1, 0);
    XCORE1000_HGEMM_MMA(3, 6, 1, 1);
    XCORE1000_HGEMM_LDS_B_FRAGMENT(1, 1, 0);
    XCORE1000_HGEMM_MMA(2, 7, 0, 0);
    XCORE1000_HGEMM_MMA(2, 7, 0, 1);
    XCORE1000_HGEMM_LDS_B_FRAGMENT(1, 1, 1);
    XCORE1000_HGEMM_MMA(2, 7, 1, 0);
    XCORE1000_HGEMM_MMA(2, 7, 1, 1);
    XCORE1000_HGEMM_MMA(3, 7, 0, 0);
    XCORE1000_HGEMM_MMA(3, 7, 0, 1);
    XCORE1000_HGEMM_MMA(3, 7, 1, 0);
    XCORE1000_HGEMM_MMA(3, 7, 1, 1);

    next_k_tile = (next_k_tile + 1) % k_tile_num;

奇奇怪怪的技巧

设计出合理的流水线后发现性能依旧拉胯，profile出来发现是L1 Cache Hit Rate很低，由于不清楚硬件memory hierarchy的具体情况，只能瞎猜，最后CTA swizzle也没拯救命中率，放弃走Cache的这条路了，转而大方吃满带宽，应用上StaggerU技巧后最终MMA吞吐到了85%，达到了200T的性能

StaggerU: 主要解决 DRAM channel / cache / TLB 冲突。特别是 K 很大、stride 是 2 的幂时，不同 CTA 如果同时从 K=0 开始，它们的 global load 地址可能呈现非常规律的 power-of-two 间距，于是大量请求会在同一个时间阶段撞到相同 memory channel / cache set / TLB 区域。

可能的优化方向

就和上面提到的一样，编译器没有生成理想的代码，当然这锅不一定是编译器的，可能只是屏障指令不合理的使用阻止了编译器优化代码，目前的代码是在进行计算前才把数据从smem load到register，这部分等待产生的latency应该存在，不过似乎没有相关的计数器能显示这个，profiler倒是显示出有明显超量的对smem进行读请求由于时间上过于密度导致相关请求队列爆满而使得发射被阻塞的计数；顺带一提，一些常规的通用的优化，比如解决bank conflict什么的，都已经用上了，没什么和CUDA不一样的地方就不提了

总结

鲲 Galgame 表情包 \[6\] - 23

对我来说探索出高效的软件流水线设计已经心满意足，懒得继续凹了，本来就没有什么硬件资料，再进一步深入单纯是猜，不值得，代码开源在https://github.com/Aoi979/CudaOps/blob/main/src/cuda_ops_core/gemm/hgemm/kernels/maca/xcore1000_hgemm_f16_nt_m256n256k64_fp32acc.hpp，感兴趣的话可以自行查阅

本文版权遵循 CC BY-NC 协议 和 本站版权政策

1 条回复

Aoikajitsu
Aoikajitsu

233

#1
发布于 (编辑于 )

AI计算也是AI(

(｡>︿<｡) 已经一滴回复都不剩了哦~