cuda-oxide:hello-constant 拆解 09——llvm-export + llc 生成 PTX

上一篇(dialect conversion driver)把 IR 降到全 dialect-llvm。这一篇看 device 编译的最后两步:llvm-export 把 in-memory IR 序列化成文本 .ll,然后调外部 llc 工具用 NVPTX backend 编成文本 .ptx。

0. 几个名词先说清楚

缩写 / 术语	英文全称	中文	含义
LLVM IR	LLVM Intermediate Representation	LLVM 中间表示	LLVM 的核心 IR,文本格式 `.ll`、二进制格式 `.bc`
NVPTX	NVIDIA PTX (target name in LLVM)	NVIDIA PTX 后端	LLVM 的一个 target backend,把 LLVM IR 编成 PTX
target triple	target triple(目标三元组)	目标三元组	LLVM 描述编译目标的字符串,形如 `arch-vendor-os`
datalayout	data layout(数据布局)	数据布局	描述各类型在该 target 上的大小、对齐方式、字节序的字符串
PTX	Parallel Thread Execution	并行线程执行	NVIDIA GPU 的虚拟 ISA,文本格式,跨代兼容
SASS	Streaming ASSembler	流式汇编	NVIDIA GPU 的真实机器码,每代 GPU 不同
llc	LLVM static compiler	LLVM 静态编译器	LLVM 自带工具,把 `.ll` 编译成目标平台汇编
!nvvm.annotations	named metadata `!nvvm.annotations`	NVVM 注解元数据	告诉 NVPTX backend”哪个函数是 GPU kernel”

1. 全景

dialect-llvm module (in-memory)
        │
        ▼ llvm-export
hello_constant.ll  (文本 LLVM IR)
        │
        ▼ Command::new("llc-22") -march=nvptx64 -mcpu=sm_61 ...
hello_constant.ptx (文本 PTX)
        │
        ▼ 下一站:CUDA driver JIT
SASS (GPU 真实机器码)

2. llvm-export 做的事

pliron 没有内置 LLVM IR 序列化(不像 MLIR 标准库自带),所以 cuda-oxide 在 dialect-llvm 这个 crate 里手写了一个 exporter。

输入是 dialect-llvm 组成的 module,输出是文本 LLVM IR 字符串。主流程:

1. 写文件头:datalayout + target triple
2. 写外部声明(declare,如 vprintf)
3. 写每个函数(define):
     a. signature
     b. 每个 basic block
        i. 每条 op 序列化(涉及 llvm_nvvm_* → llvm.nvvm.* 反向转换)
4. 写 attribute groups(#0、#1 这种)
5. 写 metadata(!nvvm.annotations)

3. hello_constant.ll 真实输出

跑完编译,hello_constant.ll 文件内容真实长这样:

3.1 文件头

; ModuleID = 'builtin.module'
source_filename = "hello_constant"
target datalayout = "e-i64:64-i128:128-v16:16-v32:32-n16:32:64"
target triple = "nvptx64-nvidia-cuda"

字段	含义
`target triple = "nvptx64-nvidia-cuda"`	告诉 llc”我是 NVPTX 目标,请用 NVPTX backend”(三段:arch / vendor / os)
`target datalayout = "..."`	描述类型布局(`e` = 小端、`i64:64` = i64 对齐 64 位等)

3.2 外部声明

declare i32 @vprintf(ptr, ptr)

vprintf 是 CUDA runtime 提供的 device 函数,只声明不定义——后面 driver JIT 时会把它链接进来。

3.3 kernel 函数

define ptx_kernel void @hello_constant(ptr %v0) {
entry:
  br label %bb0
bb0:
  %v1 = phi ptr [ %v0, %entry ]
  %v2 = call i32 asm sideeffect "mov.u32 $0, 42;", "=r"() #0
  br label %bb1
bb1:
  %v3 = zext i32 %v2 to i64
  %v4 = insertvalue { i64 } undef, i64 %v3, 0
  ...构造 "%lld <> \0" 字符串...
  %v14 = alloca [9 x i8]
  store [9 x i8] %v13, ptr %v14
  ...
  %v20 = call i32 @vprintf(ptr %v17, ptr %v19)
  br label %bb2
bb2:
  store i32 42, ptr %v1
  ret void
}

注意几个 LLVM IR 风格细节:

ptx_kernel 是 calling convention(调用约定)关键字,告诉 NVPTX backend 这是 GPU 入口
%v2 = call i32 asm sideeffect "mov.u32 $0, 42;", "=r"() #0 —— 你之前自己写的 inline asm 最终就是这一行,asm sideeffect 告诉优化器有副作用
#0 是 attribute group(属性组)引用,后面 attributes #0 = { ... convergent ... } 定义具体属性
store i32 42, ptr %v1 —— *out = 42 在 LLVM IR 里干净到不能再干净

4. 关键 trick:`llvm_nvvm_` → `llvm.nvvm.`

dialect-llvm 里 op 名字用下划线(llvm_nvvm_read_ptx_sreg_tid_x),因为 pliron 的 identifier 不允许点号。但真正的 LLVM IR intrinsic 名字必须带点(llvm.nvvm.read.ptx.sreg.tid.x)——LLVM 的命名约定。

llvm-export 在序列化函数调用时做反向转换:

match callee {
    CallOpCallable::Direct(identifier) => {
        let name = identifier.to_string();
        let fixed_name = if name.starts_with("llvm_") {
            name.replace('_', ".")          // ← 全部 _ 改成 .
        } else {
            strip_device_prefix(&name)       // 用户函数名脱掉 cuda-oxide 加的前缀
        };
        // 写出 `call ... @<fixed_name>(...)`
    }
}

例如 llvm_nvvm_read_ptx_sreg_tid_x 在导出时变成 llvm.nvvm.read.ptx.sreg.tid.x——这才是 NVPTX backend 认识的标准名字。

5. NVPTX backend 选哪些函数当 kernel

.ll 文件里函数前面那个 ptx_kernel(calling convention)是一种方式。另一种更通用的方式是 !nvvm.annotations 元数据:

!nvvm.annotations = !{!0, !1}
!0 = !{ptr @hello_constant, !"kernel", i32 1}
!1 = !{ptr @hello_kernel, !"kernel", i32 1}

读法:“@hello_constant 这个函数,角色是 kernel,值为 1”。NVPTX backend 看到这个标记就知道输出 PTX .entry 而不是 .func。

gpu_kernel: "true" 这个 attribute 从 dialect-mir 一路保留到 dialect-llvm,最后被 exporter 翻译成 !nvvm.annotations 元数据——一条信息走完整个 IR 链路。

6. llc 调用

mir-importer 的 pipeline 里调外部 llc:

let llc_candidates = [("llc-22", target), ("llc-21", target)];
//                                        ↑
//                          target 是 sm_61 / sm_80 / sm_90 等

for (llc_cmd, llc_target) in llc_candidates {
    let result = std::process::Command::new(llc_cmd)
        .arg("-march=nvptx64")               // 64-bit NVPTX
        .arg(format!("-mcpu={}", llc_target)) // 例如 sm_61
        .arg(ll_path)                         // 输入 hello_constant.ll
        .arg("-o")
        .arg(ptx_path)                        // 输出 hello_constant.ptx
        .output();
    
    if let Ok(output) = result && output.status.success() {
        return Ok(llc_target.to_string());
    }
}

候选 llc-22 → llc-21 顺序尝试。LLVM 21 是 cuda-oxide 的最低版本要求——更老版本拒绝 cuda-oxide 发的 TMA / tcgen05 / WGMMA intrinsic 签名。

实际运行时你会看到:

Using llc: llc-22 (auto-detected)
✓ PTX written to .../hello_constant.ptx (target: sm_61)
INFO [PHASE 8/9] PTX file written (LLVM IR → PTX via llc) target=sm_61

sm_61 是 Pascal 架构(GTX 10 系),自动从 CUDA device 0 探测出来的。

可以用环境变量覆盖 llc 路径:

CUDA_OXIDE_LLC=/usr/local/llvm-22/bin/llc cargo oxide run hello_constant

7. NVPTX backend 做什么

llc 是 LLVM 的”静态编译器”工具,根据 target triple 选具体的 backend:

.ll 文件 → llc → 目标平台机器码 / 汇编

target triple = "nvptx64-nvidia-cuda"  →  NVPTX backend  →  PTX
target triple = "x86_64-pc-linux-gnu"  →  X86 backend     →  x86 ASM
target triple = "aarch64-apple-darwin" →  AArch64 backend →  ARM ASM

NVPTX backend 在 LLVM 源码里 llvm/lib/Target/NVPTX/。它做的事:

指令选择(instruction selection):每条 LLVM IR 指令选一条 PTX 指令
寄存器分配(register allocation):把 LLVM 虚拟 SSA value 分配到 PTX 虚拟寄存器(%r1、%rd1…)
inline asm 嵌入:你写的 PTX 字符串直接塞进输出(带 // begin inline asm 标记)
kernel 注解处理:根据 !nvvm.annotations 决定输出 .entry 还是 .func
state space 推断:根据指针的 addrspace(N) 决定 ld.global / ld.shared / ld.local

8. hello_constant.ptx 真实输出

最终 PTX 文件头部:

//
// Generated by LLVM NVPTX Back-End
//

.version 5.0
.target sm_61
.address_size 64

.extern .func  (.param .b32 func_retval0) vprintf
(
    .param .b64 vprintf_param_0,
    .param .b64 vprintf_param_1
);

.visible .entry hello_constant(
    .param .u64 .ptr .align 1 hello_constant_param_0
)
{
    .reg .b32   %r<2>;
    .reg .b64   %rd<8>;

    ld.param.b64    %rd2, [hello_constant_param_0];
    cvta.to.global.u64  %rd1, %rd2;
    
    // begin inline asm
    mov.u32 %r1, 42;
    // end inline asm
    
    cvt.u64.u32     %rd3, %r1;
    ... printf 调用相关指令 ...
    
    st.global.b32   [%rd1], 42;                       // *out = 42
    ret;
}

段	含义
`.version 5.0`	PTX ISA 版本
`.target sm_61`	目标 GPU 架构
`.extern .func ... vprintf`	外部函数声明,对应 LLVM IR 的 `declare`
`.visible .entry hello_constant`	kernel 入口(`.entry` 来自 `!nvvm.annotations`)
`.reg .b32 %r<2>`	声明用了 2 个 32-bit 寄存器
`// begin inline asm` / `// end inline asm`	包裹你写的 PTX 字符串

9. LLVM IR ↔ PTX 一一对应

LLVM IR 的每一行几乎对应 PTX 的一行(只多了 NVPTX backend 加的入口框架):

LLVM IR:                              PTX:
%v1 = phi ptr [ %v0, %entry ]         ld.param.b64 %rd2, ...           ← 读参数
                                       cvta.to.global.u64 %rd1, %rd2;   ← 转 global 地址
%v2 = call i32 asm "mov.u32 $0, 42;"  mov.u32 %r1, 42;                   ← inline asm
store i32 42, ptr %v1                  st.global.b32 [%rd1], 42;          ← *out = 42
ret void                               ret;                                ← 返回

LLVM IR 的每条指令都能在 PTX 里找到精确对应。这就是 NVPTX backend 的”指令选择”做的事——一对一映射,没有任何高层优化。

10. llc 是外部进程,不是库

cuda-oxide 不在自己的进程里直接调 LLVM C++ API,而是 std::process::Command spawn 出去的外部命令。

替代方案是用 LLVM C API 在进程里直接调 backend,但那样 cuda-oxide 会强依赖某个特定 LLVM 版本(C++ ABI 不稳定)。用 llc 命令行调用换来版本独立性,代价是多个进程启动成本(编译时几十 ms)。

跟整个 cuda-oxide 的设计哲学一致:关注点分离 + 用外部成熟工具。

11. 一句话总结

llvm-export 手写一个文本 LLVM IR 序列化器(因为 pliron 没自带),把 dialect-llvm module 写成 .ll 文件,过程中做 llvm_nvvm_* → llvm.nvvm.* 的反向转换(下划线还原为点号),还把 gpu_kernel: true attribute 翻译成 !nvvm.annotations 元数据。然后 Command::new("llc-22") spawn 外部 LLVM 工具,用 NVPTX backend 把 .ll 编成 .ptx 文本。LLVM IR 每条指令在 PTX 里都有精确对应,因为 NVPTX backend 主要做的是指令选择(一对一映射),不做高层优化。device 编译到此结束,接下来是 host runtime + GPU 执行。

系列上一篇: cuda-oxide:hello-constant 拆解 08——mir-lower:dialect conversion driver

评论区

评论功能即将上线, 敬请期待。