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/...

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

需求及指导思想: https://zhuanlan.zhihu.com/p/279341172
自动指令编码: https://zhuanlan.zhihu.com/p/296908633
指令编码的特殊处理: https://zhuanlan.zhihu.com/p/314410214

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

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 的目录结构如下：

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)。

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.cuasm, cuasm 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. 在工程目录下创建一个cu文件，例如cudatest.cu。
2. 将CuAssembler/bin/makefile复制到工程目录，将其中的BASENAME修改为前面的文件名（这里是BASENAME=cudatest）。将ARCH变量设置为要编译的SM版本，这里假设是ARCH=sm_75。如果需要处理额外的include或link文件，修改对应的$INCLUDE 和 $LINK 变量即可。
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. 根据用户需求修改hack.cudatest.sm_75.cuasm。
2. 运行make hack。这会将用户修改后的 hack.cudatest.sm_75.cuasm 汇编为 hack.cudatest.sm_75.cubin，用它替换原编译链中的cudatest.sm_75.cubin，最终生成可执行文件cudatest。
3. 得到最终可执行文件cudatest后，就可以按原来的方式运行，例如./cudatest。

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

一些提示：

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或不可编码组合时，也会列出当前所有已知组合，方便用户查找原因，等等。