并行编程实战——PTX和JIT

一、介绍

在前面学习编程语言和编译器相关的知识时，曾经学习过LLVM的前后端架构，特别是对IR中间表示的优越之处进行了分析和说明。包括后面出现的WSAM其实也是有着相似的设计思想。从现实的情况来看，软件工程发展到今天，各种编程语言和相关的脚本语言已经很多，如何能够将它们从某种角度统一起来，一直是一些大佬们追求的梦想。

二、中间指令集

其实如果从硬件角度看，对软件语言的底层统一已经完成。那就是二进制或汇编语言。但这里面又出来一个问题，对GPU来说，硬件的产品也有很多，即使NVIDIA公司一家，也出了很多的系列。这就导致了一个问题，它们之间的指令集也各有不同。甚至，可以畅想一下，NVIDIA未必没有想统一硬件应用的想法。
但是，对于自家公司的产品都无法统一相同的指令集，那么想兼容外部硬件，难度更是可想而知。所以第一步就是先统一自己的指令集。但事实的指令集已经存在，如何统一呢？
还是软件设计中那句话，搞不定就加一层。NVIDIA也是这么做的，在硬件上抽象出来一套虚拟机指令集。它类似于JAVA的虚拟机字节码，也可以理解为LLVM的IR等等。这样一来，就可以做到CUDA程序编译时的一次编译到处运行。

三、PTX

PTX，Parallel Thread Execution。并行线程执行。在NVIDIA公司提供的说明中，指出它是一种与具体硬件无关的虚拟指令集。在CUDA的底层库中，提供了cuda::ptx的名空间，可以通过其进行编程应用。在前面的分析过程中，可以看到类似下面的代码，其实就是这个机制：

#include <cuda/ptx>
#include <cooperative_groups.h>

__global__ void init_barrier()
{
  __shared__ uint64_t bar;
  auto block = cooperative_groups::this_thread_block();

  if (block.thread_rank() == 0)
  {
    // A single thread initializes the total expected arrival count.
    cuda::ptx::mbarrier_init(&bar, block.size());
  }
  block.sync();
}

这种虚拟指令集会在真正安装到GPU执行时，由驱动程序和PTX转换器共同翻译成具体的硬件指令集，这样就可以由硬件直接执行了。
PTX出现的主要目的是为了高效运行和支持NVIDIA Tesla架构相关计算我的GPU而设计的。目前来看，CUDA和C/C++程序都可以使用编译器生成PTX指令，而在最终则可以优化翻译为目标架构的原生指令。这一点非常重要，可以理解为对PTX的优化同样会体现在原生指令中。其具体包括以下几点：

1. 提供能够跨越多个GPU世代的稳定指令集架构（ISA）。
2. 在编译型应用程序中实现与原生GPU性能相当的性能
3. 为 C/C++ 和其他编译器提供与机器无关的指令集架构 (ISA)
4. 为应用程序和中间件开发人员提供代码分发指令集架构 (ISA)
5. 为优化代码生成器和转换器提供通用的源代码级指令集架构 (ISA)，将 PTX 映射到特定的目标机器
6. 方便手动编写库、性能内核和架构测试
7. 提供一种可扩展的编程模型，能够涵盖从单个 GPU 单元到多个并行单元的各种 GPU 规模。

提供一种可扩展的编程模型，能够涵盖从单个 GPU 单元到多个并行单元的各种 GPU 规模。
PTX的优势就在于解耦，而解耦后，就可以实现最大可能的兼容。这样开发者就不必为每个不同的版本的硬件进行相关代码的兼容性修改。另外，NVIDIA还通过PTX提供了对底层硬件直接操作的能力，这就意味着其可以对硬件直接进行优化，这一点非常重要。

四、JIT

PTX的编程模型和抽象指令集是一种兼容和跨代的静态体现。而JIT（Just-In-Time，即时编译）则实现了动态的兼容转换。有过JAVA编程经验的可能都知道，JAVA可以在热点应用时进行即时的编译，从而可以保证运行的性能，达到兼顾性能和速度的双重优势。但CUDA中的JIT只是一种即时编译（没有热点一说），它负责在运行时将PTX编译为目标GPU具备的相关机器码即SASS（Streaming Assembler）。
PTX的兼容必须要JIT的实现才能真正达到设计的目的。JIT可以根据实际的目标GPU对PTX进行翻译。从而真正实现了上层对硬件的抽象。同时，其还提供了缓存的机制，如果下次还是同样的环境启动，驱动会直接加载缓存而不再进行相关的编译转换步骤，提高了运行速度。
其一般的流程如下：

1. 在CUDA12.8上开发编译了一个程序，其目标架构为较老的Maxwell
2. 现在要将程序运行在新的架构Blackwell上。先启动CUDA程序，调用相关的API
3. GPU驱动发现没有预编译的SASS，但发现包含有PTX代码
4. 驱动调用JIT编译器，针对新架构平台的GPU对PTX代码进行优化（比如使用最新指令等）并将其编译为优化后的SASS机器码
5. 将编译优化后的代码加载到GPU运行并缓存在本地

五、如何协同

通过上述的分析可以发现，其实这两技术从本质上并不是一种全新的技术。可以认为是一种拿来主义的具体的实践。在CUDA开发的过程中，PTX可以通过在编译阶段利用编译器将相关的代码编译为PTX代码的可执行文件，这样就实现中间层的硬件隔离，达到了事实上的软硬件解耦。然后，在真正的代码安装到GPU运行时，由CUDA运行时库将相关的转换过程移交给驱动，驱动会利用JIT编译器将PTX代码转换为SASS。这样就完美的实现了从静态解耦到动态运行的全流程。
如果想查看PTX到JIT编译后的原生代码，可以通过相关的工具进行处理如cuobjdump或nvdisasm，可以使用下面的命令获取这些工具使用的中间文件：

nvcc -keep -arch=sm_80 my_ir.cu

另外在NVIDIA的Triton框架中，可以直接通过对缓存的操作查看相关的代码，如kernel.asm[‘ptx’]或kernel.asm[‘llir’]等。有兴趣可以可以查阅相关的资料。

六、总结

PTX和JIT是中间层抽象的又一次具体实践。它提供了稳定的编程模型，为硬件不同版本甚至跨代（甚至有可能是跨硬件平台）兼容提供了基础的保障。达到了可以随时针对不同平台进行性能优化的可能。通过对软件与硬件版本的解耦，实现了类似于JAVA的一次编译多处运行的能力。