下冰雹 · 2022年10月10日 · 广东

CuAssembler 开源 | 探求 NVIDIA GPU 极限性能的利器

在通常的 CUDA 编程中,用户主要通过 CUDA C/C++ 或 python 语言实现 CUDA 功能的调用。在 NVIDIA 对 CUDA C 的官方支持工具链中,CUDA 设备端代码从 CUDA C 转换为 NVVM IR(类似 LLVM IR),再生成对应架构的 ptx,最后由 ptxas 编译生成最终设备端的机器码。这其中CUDA C、NVVM IR 和 PTX 都有相对完整的工具和文档支持。但对于最终生成的设备上运行的汇编(称为Streaming ASSembly,SASS)及其机器码,却没有提供类似 gcc 的 gas 和 llvm 的 llvm-mc 这种汇编工具。这就意味着最终代码生成质量将更多取决于编译器,而用户对底层汇编缺乏直接控制手段,从而限制功能使用和性能发挥。

CuAssembler 是个 CUDA SASS 汇编器,主要是把 nvdisasm 生成的反汇编输出汇编成可加载执行的 cubin 文件,实现对最终汇编机器码的绝对控制,弥补当前 NV 没有官方汇编器支持的不足。

同类工具里有asfermi,KeplerAs,maxas,turingas等针对特定 NV 架构的汇编工具。与这些工具相比,CuAssembler 对多架构的支持更完整(当前主要有 SM60/SM61/SM70/SM75/SM80/SM86),指令编码支持和扩展也更容易。同时,CuAssembler 能更好地嵌入现有 NVIDIA 官方工具链,可以实现大部分代码的功能继承自原有编译流程,从而支持更多 CUDA feature。

CuAssembler 的主要使用场景有两类。

一是性能敏感算子(如 conv、gemm 等)的汇编级的极致优化。由于汇编是设备运行的最终代码,使用汇编可以支持最细粒度的性能调优。

第二个场景是通过一些特定汇编代码实现 microbenchmarking,测试指令集及微架构的性能信息,从而为程序性能优化提供参考。

CuAssembler 之前版本的简要介绍可参考:https://zhuanlan.zhihu.com/p/...

更早的专栏文章有介绍实现原理:

这里发布的版本在原版基础上有多个更新,主要有:

1. 补充完善了对 SM60/61/70/75/80/86 的支持,特别是解决了 Ampere 架构中的一些隐藏问题。

2. 补充了多个脚本文件,日常使用更方便快捷。

3. 内部实现做了大量优化,大大提升了可靠性和易用性。

4. 改进了编码的实现方式,能覆盖更多隐藏编码和提高可靠性。

安装、配置与使用

一、安装说明

新版 CuAssembler 仓库地址为:https://github.com/openppl-public/CuAssembler ,原版仓库后面也会尽量同步更新。用户可以自由选择 git clone或是下载 zip 文件包等方式。

CuAssembler 是个 python 包,所以依赖 python(推荐 python3.8+)。除 python 自带的标准库外,还使用了两个第三方库:sympy, pyelftools,都可以通过 pip 或 conda 下载安装。另外,因为需要用到 cuda 的反汇编和 dump 功能,需要 CUDA SDK 中的 cuobjdump 和 nvdisasm 两个可执行文件(不需要完整的 CUDA SDK,只需这两个可执行文件即可,建议加到系统 PATH)。由于 cuobjdump 和 nvdisasm 可能有一些隐藏 bug 导致部分功能受限(主要是一些指令编码处理问题),建议使用最新版本的 CUDA SDK。

CuAssembler 的目录结构如下:

image.png

CuAssembler 目录结构

•  bin : bin 目录主要是一些常用功能的脚本,方便用户更快捷的调用。建议加入系统 PATH。

  • • CuAsm :CuAsm 目录包括所有的模块文件,想在 python 中 import CuAsm(也包括运行 bin 中的 python 脚本),需要将其拷贝到系统 PATH 或是已知的 PYTHONPATH 下,也可以把其父目录添加到 PYTHONPATH中(推荐)。
  • • TestData : 这个目录主要是一些测试 cu 或 cubin 文件,由于版本变化,部分格式可能有出入,仅供参考。
  • • Tests : 一些调用 CuAsm 模块功能的简单脚本,仅供开发者参考。
  • • Tools : Tools 目录下包括一个 cuasm 文件的 vscode 语法高亮插件,也可以用来高亮 cuobjdump 得到的 sass。它很简陋,也没有加入 vscode 的官方市场,也没有安装文件,可以复制到 vscode 的插件目录直接使用。建议配色使用 vscode 的 Dark+(default Dark)。

image.png

  • VSCode 语法高亮效果

二、配置说明

Linux 用户可以直接修改.bashrc, 设置相应的环境变量(注意修改为对应路径):

 

 export PATH=${PATH}:~/works/CuAssembler/bin  
  export PYTHONPATH=${PYTHOPATH}:~/works/CuAssembler/

bin 目录的脚本基本都有详细的介绍文档。Linux 下可以直接运行./cuasm.py 、./dsass.py等 python 文件(前提是能被/usr/bin/env找到 python 路径,否则需要用python *.py方式运行), 但是我通常会创建相应的符号链接并设为可执行:

ln -s cuasm.py cuasm  
chmod a+x cuasm

这样就可以直接用cuasm/dsass/hcubin/hnvcc作为命令了(不用加后缀也不用python cuasm.py这种繁琐的方式)。

windows 下虽然可以配置 py 运行方式,但并不方便。所以我为这几个 py 脚本都创建了 bat 批处理文件,方便运行:

@echo off  
python "%~dp0\\cuasm.py" %\*

其中的%~dp0 表示当前 bat 所在目录,如果需要把 bat 复制到其他目录,把它替换为绝对或相对路径即可。配置完成后,cuasm t.cubin就可以将t.cubin反汇编为t.cuasmcuasm t.cuasm -o t2.cubin 就可以将其汇编为相应的二进制文件。

脚本都有相应的介绍和命令行参数文档,可以自行参考。这里简单介绍一下各个脚本的功能:

• cuasm : 将二进制 cubin 文件反汇编为用户可编辑的 cuasm 文本文件,或是将用户修改后的 cuasm 文件汇编为 cubin 文件,从而在CUDA 程序中调用运行。这是 CuAssembler 提供的主要功能。

• hcubin : hcubin 主要是解决 Ampere 架构的 SM80、SM86 等反汇编文本不显示默认Cache-policy descriptor的问题。它会修改相应 bit 强制所有 desc 都必须显示。具体介绍见后文。

• dsass : NV 提供的官方反汇编工具中,cuobjdump/nvdisasm 都不会显示 control code 相关信息,dsass 可以提供一个方便的接口把原汇编和 control code 都完整的显示出来。dsass既可以接受cubin文件和包含cubin的二进制文件输入,也可以输入从 cuobjdump -sass *.exe 得到的sass文本文件。dsass输入cubin文件时会自动检测是否需要hack desc,其他输入格式则会保留原编码。

• hnvcc : hnvcc是一个hack版的nvcc,它可以在环境变量(HNVCC_OP)的控制下打断或更改nvcc中涉及cubin的编译流程。不设置HNVCC_OP时,hnvcc将直接调用nvcc;在HNVCC_OP=dump时,它会把ptxas的输出cubin保存下来,供用户反汇编和修改;而HNVCC_OP=hack时,它会用对应的cubin文件替换ptxas原来的输出。这样即使是使用runtime api,也可以直接替换cubin,无需再专门写一个driver api的版本去运行cubin文件。

• makefile : 这不是脚本,而是个makefile模板。它提供了一个简单的cu文件编译、dump以及hack的逻辑。make d2h就可以dump相应cubin并反汇编为cuasm,make hack就能用修改的cuasm汇编后替换对应的cubin,并生成最终的可执行文件。

三、使用说明

这里介绍 linux 下最方便快捷地使用 CuAssembler 的方法:

  1. 1. 在工程目录下创建一个cu文件,例如cudatest.cu
  2. 2. 将CuAssembler/bin/makefile复制到工程目录,将其中的BASENAME修改为前面的文件名(这里是BASENAME=cudatest)。将ARCH变量设置为要编译的SM版本,这里假设是ARCH=sm_75。如果需要处理额外的include或link文件,修改对应的$INCLUDE 和 $LINK 变量即可。
  3. 3. 运行make d2h,其中 d2h 表示 "dump to hack"。这会生成3个新文件:
  • • dump.cudatest.sm_75.cubin : 由原始输入 cudatest.cu 编译生成的cubin文件.
  • • dump.cudatest.sm_75.cuasm : 原始cubin的反汇编文件.
  • • hack.cudatest.sm_75.cuasm : 原始反汇编dump.cudatest.sm_75.cuasm的副本,方便用户修改.
  1. 1. 根据用户需求修改hack.cudatest.sm_75.cuasm
  2. 2. 运行make hack。这会将用户修改后的 hack.cudatest.sm_75.cuasm 汇编为 hack.cudatest.sm_75.cubin,用它替换原编译链中的cudatest.sm_75.cubin,最终生成可执行文件cudatest
  3. 3. 得到最终可执行文件cudatest后,就可以按原来的方式运行,例如./cudatest

整个过程与通常开发方式比较接近,也不用专门写driver api相关接口去调用 cubin。

一些提示:

  1. 1. 由于 CuAssembler 当前的功能限制,cuasm所有的接口需要继承自原cubin文件,所以需要修改kernel参数、全局变量、constant memory配置时,需要重新make d2h生成新的cuasm并在此基础上重新修改。不过由于cuasm是文本格式,用户可以使用一些版本控制软件尽量减少重复修改的部分。

    2. 这个方式比较适合独立的CUDA程序,如果需要修改大型程序(比如多个cu文件)的一部分,nvcc的参数和编译链条会较为复杂,存在较多风险;

功能介绍更新

一、Ampere Cache Policy UR 问题

新版发布的一个主要更新是改善了对 Ampere 架构(SM80/SM86 等)的支持。Ampere 除了新加一些指令外,总体与 Turing 非常接近,但存在一些隐藏问题,就是 Cache Policy UR 不显示的问题。

对于 Ampere 架构,用 NVCC 编译的 cubin 中,如果有 load、store 类指令(比如 LDG/STG/ATOMG/RED/LDGSTS 等等),通常会使用default cache policy(通常位于c[0x0][0x118])。它会在程序早期被载入某对UR(64bit,比如UR[4:5]),其他非默认cache policy需要相应的显式初始化。例如,我们使用 ptx 的createpolicy指令创建 policy 并用于load指令中:

.reg .b64 %cp<3>;  
  
createpolicy.fractional.L2::evict\_last.L2::evict\_unchanged.b64 %cp0, 0.75;  
createpolicy.fractional.L2::evict\_last.L2::evict\_unchanged.b64 %cp1, 0.5;  
createpolicy.fractional.L2::evict\_last.L2::evict\_unchanged.b64 %cp2, 0.25;  
  
ld.global.b32 %r0, \[%rd0\];  
ld.global.L2::cache\_hint.b32 %r1, \[%rd0+1024\], %cp0;  
ld.global.L2::cache\_hint.b32 %r2, \[%rd0+2048\], %cp1;  
ld.global.L2::cache\_hint.b32 %r3, \[%rd0+3072\], %cp2;  
  
st.global.v4.b32  \[%rd1\], {%r0, %r1, %r2, %r3};

用ptxas编译后,汇编中UR[4:5]是默认cache policy,UR[6:7]/UR[8:9]/UR[10:11]是自定义的非默认cache policy:

code for sm_86
        Function : _Z5ktestPmPi
.headerflags    @"EF_CUDA_SM86 EF_CUDA_PTX_SM(EF_CUDA_SM86)"
/*0000*/                   IMAD.MOV.U32 R1, RZ, RZ, c[0x0][0x28] ;  /* 0x00000a00ff017624 */
                                                                    /* 0x000fc400078e00ff */
/*0010*/                   S2R R2, SR_TID.X ;                       /* 0x0000000000027919 */
                                                                    /* 0x000e220000002100 */
/*0020*/                   ULDC.64 UR4, c[0x0][0x118] ;             /* 0x0000460000047ab9 */
                                                                    /* 0x000fe40000000a00 */
/*0030*/                   UMOV UR6, 0x0 ;                          /* 0x0000000000067882 */
                                                                    /* 0x000fe40000000000 */
/*0040*/                   UMOV UR7, 0x14b00000 ;                   /* 0x14b0000000077882 */
                                                                    /* 0x000fe40000000000 */
/*0050*/                   UMOV UR8, 0x0 ;                          /* 0x0000000000087882 */
                                                                    /* 0x000fe40000000000 */
/*0060*/                   UMOV UR9, 0x14700000 ;                   /* 0x1470000000097882 */
                                                                    /* 0x000fe40000000000 */
/*0070*/                   UMOV UR10, 0x0 ;                         /* 0x00000000000a7882 */
                                                                    /* 0x000fc40000000000 */
/*0080*/                   UMOV UR11, 0x14300000 ;                  /* 0x14300000000b7882 */
                                                                    /* 0x000fe20000000000 */
/*0090*/                   LEA R2, P0, R2, c[0x0][0x160], 0x2 ;     /* 0x0000580002027a11 */
                                                                    /* 0x001fca00078010ff */
/*00a0*/                   IMAD.X R3, RZ, RZ, c[0x0][0x164], P0 ;   /* 0x00005900ff037624 */
                                                                    /* 0x000fca00000e06ff */
/*00b0*/                   LDG.E R8, [R2.64] ;                      /* 0x0000000402087981 */
                                                                    /* 0x000ea8000c1e1900 */
/*00c0*/                   LDG.E R9, desc[UR6][R2.64+0x400] ;       /* 0x0004000602097981 */
                                                                    /* 0x000ea8200c1e1900 */
/*00d0*/                   LDG.E R10, desc[UR8][R2.64+0x800] ;      /* 0x00080008020a7981 */
                                                                    /* 0x000ea8200c1e1900 */
/*00e0*/                   LDG.E R11, desc[UR10][R2.64+0xc00] ;  
...

这里可以发现/*00b0*/LDG.E R8, [R2.64] ; 其实应该是LDG.E R8, desc[UR4][R2.64] ;,由于UR[4:5]是默认policy,没有显示在汇编文本里。但是,默认policy的存放位置是不固定的,可能是UR[4:5], 也可能是其他。这导致显示的汇编文本无法完全反映原始语义信息,相当于经过反汇编、汇编后,部分编码信息丢失了。这种编码信息的丢失也会导致无法正确修改NVCC输出的cubin文件。

经过观察,我发现设置102 bit(也就是LDG第二行的0x000ea8200...中的2)可以强制让nvdisasm显示所有的desc。其他指令看起来不用这个bit,所以也无害。因此,cubin转为cuasm时,CuAssembler会自动把所有指令的这个bit置1,从而使编码信息完备。

修改所有指令编码也会带来新的问题。因为CuAssembler所有的编码规则不是手工填的,而是根据输入解方程算到的。额外的无用bit会导致每个操作数或modifier的权重计算错误。因此,CuAssembler/bin中的脚本hcubin可以专门用来修复Ampere架构中使用了desc的指令编码。dsass工具在dump sass时,如果输入是cubin文件,也会自动完成这个编码hack。

当然,对于多数用户来说,编码变化是汇编器内部机制问题,不用过多关心。只要汇编文本能显示完备的汇编语义信息就足够了。

二、Reuse bit 编码问题

在Scott Gray最开始介绍Maxwell Control Codes的文章中,reuse bit是4bit,每个对应一个操作数slot。后来的多数文章也继承了这个观点,CuAssembler也用了显式的[R-R-]这种方式来标记对应的slot。但是根据对Turing及Ampere大量反汇编及编码对应关系的观察,这里面存在两个问题:

有些指令没有显示 Reuse,但仍可设置 Reuse bit:典型的比如LDG,BAR,TEX等指令。它会在原Reuse bit中写入一些额外信息,具体信息含义不明。

第四个操作数的 Reuse,不对应 Reuse 的第 4 位:在Ampere之前,没有发现有4个GPR操作数且能reuse的指令,所有的reuse都只用了3个bit。但到Ampere中,HMMA.SP支持4个GPR输入,且都可以reuse。而第4个操作数的reuse并没有设置第4个bit,而是在原编码的中间位置。因此,原先的reuse应该就只有3个bit,第4位其实一直未使用。

因此,在新版的CuAssembler中,reuse不再作为单独的control code域显示,而是作为通常的modifier处理,由程序自动确认编码。当然,这也导致前述LDG、BAR、TEX等指令在通常reuse bit上写的信息会丢失。这需要在反汇编时就提取相应信息,并将其写入反汇编文本中,这样才能在汇编时恢复。由于这些bit含义暂不明,这个暂不支持。

三、改进内部构建机制和错误检查

由于缺乏完备的编码规则,CuAssembler不可避免的会遇到很多编码问题,如无法识别的modifier,不能编码的组合等。我们改进了相关描述文件,可以更容易发现编码问题。如对指令BRA_II, 它的内部信息如下(参考CuAssembler/CuAsm/InsAsmRepos/DefaultInsAsmRepos.sm_75.txt):

'BRA_II':CuInsAssembler("", {"InsKey" : 'BRA_II', 
  "InsRepos" : [([7, -16], ['0_BRA', '0_NegAddrOffset'], 1088033237653166188516112711),
([8, 48], ['0_BRA'], 1083197534374707946695199047),
([8, 368], ['0_BRA'], 1083197534374709321084733767),
([0, 1296], ['0_BRA'], 1083197534374713306814351687),
([7, 10960], ['0_BRA', '0_DIV'], 1083197534374754821968263495),
([7, 48], ['0_BRA', '0_CONV'], 1083197534374707959580096839),
([9, 1008], ['0_BRA', '0_U'], 1083197534374712074158774599)], 
  "InsModiSet" : {'0_BRA': 0, '0_NegAddrOffset': 1, '0_DIV': 2, '0_CONV': 3, '0_U': 4}, 
  "ValMatrix" : Matrix([
[1, 1, 0, 0, 0, 7,   -16],
[1, 0, 0, 0, 0, 8,    48],
[1, 0, 0, 0, 0, 8,   368],
[1, 0, 0, 0, 0, 0,  1296],
[1, 0, 1, 0, 0, 7, 10960],
[1, 0, 0, 1, 0, 7,    48],
[1, 0, 0, 0, 1, 9,  1008]]), 
  "PSol" : Matrix([
[ 0x38000000000000000000947], # 0_BRA
[   0x400000000000000000000], # 0_NegAddrOffset
[               0x200000000], # 0_DIV
[               0x300000000], # 0_CONV
[               0x100000000], # 0_U
[                    0x1000], # Pred
[               0x100000000], # V1
]), 
  "PSolFac" : 1, 
  "ValNullMat" : None, 
  "InsRecords" : [(0x0000e0, 0x000000000383fffffffffff000007947, "BRA 0xe0;"),
(0x000150, 0x00000000038000000000003000008947, "@!P0 BRA 0x190 ;"),
(0x000270, 0x00000000038000000000017000008947, "@!P0 BRA 0x3f0 ;"),
(0x0002b0, 0x00000000038000000000051000000947, "@P0 BRA 0x7d0 ;"),
(0x000e60, 0x000000000380000000002ad200007947, "BRA.DIV 0x3940 ;"),
(0x0016f0, 0x00000000038000000000003300007947, "BRA.CONV 0x1730 ;"),
(0x005f00, 0x0000000003800000000003f100009947, "@!P1 BRA.U 0x6300 ;"),
], 
  "ErrRecords" : {-0x14000000000000000000000000000000 : (0x003770, 0x14000000038000000000009000002947, "@P2 BRA 0x3810 ;"),
},   "Rhs" : Matrix([
[0x383fffffffffff000007947],
[0x38000000000003000008947],
[0x38000000000017000008947],
[0x38000000000051000000947],
[0x380000000002ad200007947],
[0x38000000000003300007947],
[0x3800000000003f100009947],
]), 
  "Arch" : CuSMVersion(75) }

其中ValNullMat为空,表明这个是完备的系数矩阵(可能有未见过的modifier,但所有操作数组合都可编码)。ErrRecords里面记录了所有已知编码与预期不一致的指令,例如这里BRA就在通常reuse bit的位置中写入了0x14,这是普通指令的编码系统没有的。如果发现一些指令的编码存疑,检查这些信息有利于快速判断是否是规则问题。

另外,对于用户手写汇编时常见的typo问题,增加了更友好的出错提示,例如SM75中错把IADD3写成了IADD,则会提示:

File c.cuasm:3126 :
    [B0-----:R-:W-:Y:S04]         /*0d10*/                   IADD R1, R1, 0x8, RZ ;
Error when assembling instruction "[B0-----:R-:W-:Y:S04] IADD R1, R1, 0x8, RZ ;":
    Unknown InsKey(IADD_R_R_II_R) in Repos!
Available InsKeys:
    IADD3_R_R_II_R
    IMAD_R_R_II_R
    IADD3_R_P_R_II_R
    IMAD_R_R_II_R_P
    IMAD_R_P_R_II_R

这可以更好的帮助用户快速找到错误并修正。输入了未知modifier或不可编码组合时,也会列出当前所有已知组合,方便用户查找原因,等等。

结语

面对 CUDA 底层这种封闭生态,做汇编器其实是有一些事倍功半的。毕竟 CUDA 的中上层工具和生态已经相当完备,功能和性能已经能满足多数需求。

**所以很多时候我们也没太追求完备性,只是希望能多一个入口,多一种选择。或者说,做汇编器不是目的,而是手段。

后面我们会基于汇编器开发一些更有实际意义的功能,希望各位同仁继续关注~**

欢迎 Star ✨ 

github.com/openppl-public/CuAssembler

作者: 杨仍才
文章来源:OpenPPL

推荐阅读

更多IC设计技术干货请关注IC设计技术专栏。
迎添加极术小姐姐微信(id:aijishu20)加入技术交流群,请备注研究方向。
推荐阅读
关注数
16350
内容数
1226
嵌入式端AI,包括AI算法在推理框架Tengine,MNN,NCNN,PaddlePaddle及相关芯片上的实现。欢迎加入微信交流群,微信号:aijishu20(备注:嵌入式)
目录
极术微信服务号
关注极术微信号
实时接收点赞提醒和评论通知
安谋科技学堂公众号
关注安谋科技学堂
实时获取安谋科技及 Arm 教学资源
安谋科技招聘公众号
关注安谋科技招聘
实时获取安谋科技中国职位信息