cuda-oxide:hello-constant 拆解 10——host binary 启动 + CUDA driver 加载 PTX

上一篇(llvm-export + llc)把 device 编译走完,.ptx 文件已经写盘。device 编译完成后,rustc 自带 LLVM 后端把 host 代码编成 x86_64 binary(main、DeviceBuffer、println! 等),cargo run 自动 exec 出来——进入运行时。本文拆 host 端从启动到 launch kernel 之间的全部 CUDA driver 调用。

0. 几个名词先说清楚

缩写 / 术语	英文全称	中文	含义
CUDA driver API	CUDA Driver API	CUDA 驱动 API	NVIDIA 提供的底层 CUDA API,以 `cu*` 开头(如 `cuCtxCreate`)
CUDA runtime API	CUDA Runtime API	CUDA 运行时 API	更高层的封装,以 `cuda*` 开头(如 `cudaMalloc`)——本文不用
context	context(上下文)	CUDA 上下文	一个线程绑定的 device + stream + 内存映射等状态集合
stream	stream(流)	CUDA 流	GPU 的命令队列,所有异步操作按 stream 排队
device pointer	device pointer / `CUdeviceptr`	设备指针	64-bit 整数,指向显存里的地址,host 不能 deref
JIT	Just-In-Time compilation	即时编译	CUDA driver 在加载 PTX 时把它编译成当前 GPU 的 SASS
RAII	Resource Acquisition Is Initialization	资源获取即初始化	Rust / C++ 风格:对象 drop 时自动释放资源

1. host binary 启动

device 编译完之后,cargo run --release 自动 exec 出 target/release/hello_constant 二进制。main 函数入口:

fn main() {
    init_tracing();
    tracing::info!("[PHASE 9/9 a] host binary started");
    println!("=== Unified Hello Constant Example ===\n");
    // ...
}

实际运行时你会看到:

INFO hello_constant: [PHASE 9/9 a] host binary started

注意 PHASE 9/9 a-e 这五条是新进程(host binary)输出的,跟前面 PHASE 1/9 - 8/9 完全不在同一个进程里——前面是 cargo-oxide / rustc / mir-importer,这里是编出来的 host 程序。

2. `CudaContext::new(0)` — 连 GPU

let ctx = CudaContext::new(0).expect("Failed to create CUDA context");
let stream = ctx.default_stream();

cuda-core 这个 crate 用 RAII 把 CUDA driver API 包成 Rust 友好的类型。CudaContext::new(0) 背后做三件 driver 调用:

cuInit(0);                          // 初始化 CUDA driver 子系统
cuDeviceGet(&dev, 0);               // 拿 device 0 的句柄
cuCtxCreate_v2(&ctx, 0, dev);       // 创建 device context

CUDA context 是个状态机——保存当前线程绑定的 device、stream、内存映射等。所有后续 cu* 调用都需要”当前线程有一个 active context”。

cuda-core 用 RAII 包装:CudaContext 这个 Rust struct drop 时自动调 cuCtxDestroy,不会泄漏。

let stream = ctx.default_stream() 拿默认流(stream 0)——stream 就是 GPU 的命令队列,所有 launch、memcpy 都按 stream 排队执行,默认是异步的。

实际运行时:

INFO hello_constant: [PHASE 9/9 b] initializing CUDA context (device 0)

3. `DeviceBuffer::<i32>::zeroed(&stream, 1)` — cudaMalloc

let out_dev = DeviceBuffer::<i32>::zeroed(&stream, 1).expect("Failed to allocate");

分配 1 个 i32 的 device 内存,初始化为 0。底层:

cuMemAlloc_v2(&dptr, sizeof(i32) * 1);      // 分配 4 字节 device 内存
cuMemsetD8Async(dptr, 0, 4, stream);        // 异步清零

dptr 是 device pointer——一个 64-bit 整数,指向 GPU 显存里的地址,不是 host 进程地址空间的指针。host 代码不能 deref 它,只能通过 cu* API 操作。

DeviceBuffer 同样 RAII:drop 时自动 cuMemFree_v2。

4. `load_module_from_file` — 最”魔法”的一步:PTX → SASS JIT

let cuda_module = ctx
    .load_module_from_file("hello_constant.ptx")
    .expect("Failed to load PTX module");

底层调 cuModuleLoad:

cuModuleLoad(&module, "hello_constant.ptx");

这一步是整条链路里最魔法的一步。cuModuleLoad 看到输入是 PTX 文本(不是 cubin 二进制),就触发 CUDA driver 内置的 JIT 编译器:

hello_constant.ptx (文本 PTX,虚拟 ISA)
        │
        ▼ CUDA driver JIT(ptxas 库,内嵌在 driver 里)
hello_constant SASS  (sm_61 特定的真实机器码)
        │
        ▼ 拷到 GPU 显存
GPU 准备执行

关键事实:

事实	解释
PTX 是虚拟 ISA	像 Java 字节码——由 driver 在 load 时编译成当前 GPU 的真实机器码 SASS
不同 GPU 编出来的 SASS 不一样	sm_61(Pascal)和 sm_90(Hopper)的 SASS 截然不同
JIT 结果会缓存	CUDA driver 把编出来的 SASS 缓存到 `~/.nv/ComputeCache/`,下次 load 同一 PTX 直接复用
PTX 是向前兼容的	sm_61 写的 PTX 可以在 sm_80 上 load(driver 帮你升级 SASS),反向不行

可以用 cuobjdump 看实际 SASS:

cuobjdump --dump-sass hello_constant.ptx
# 或从 ~/.nv/ComputeCache/ 里把 cubin 翻出来再 dump

实际运行时:

INFO hello_constant: [PHASE 9/9 c] loading PTX module from hello_constant.ptx

JIT 编译可能花几十 ms(首次 load),缓存命中后就是几 ms。

5. `kernels::from_module(...)` — 类型化封装

let module = kernels::from_module(cuda_module.clone())
    .expect("Failed to initialize typed CUDA module");

#[cuda_module] 宏在 host 端展开会生成一个 from_module 函数,给每个 #[kernel] 函数生成一个带类型签名的 launcher:

// 由 #[cuda_module] 宏自动生成
impl kernels::LoadedModule {
    pub unsafe fn hello_constant(
        &self,
        stream: &CudaStream,
        config: LaunchConfig,
        out: *mut i32,
    ) -> Result<(), DriverError> {
        // ... 把参数打包成 void**,调 cuLaunchKernel ...
    }
}

好处:调用 kernel 像调普通 Rust 函数,类型安全(编译期检查参数类型)。不需要手写 void* args[] 那套不安全的 C 风格参数打包。

6. `module.hello_constant(...)` — cuLaunchKernel

unsafe {
    module.hello_constant(
        stream.as_ref(),
        LaunchConfig::for_num_elems(1),       // grid=(1,1,1), block=(256,1,1)
        out_dev.cu_deviceptr() as *mut i32,
    )
}.expect("Kernel launch failed");

底层调 cuLaunchKernel:

void* args[] = { &out_ptr };
cuLaunchKernel(
    cu_function,              // kernel 句柄(之前从 module 里查到的)
    1, 1, 1,                  // gridDim (x, y, z)
    256, 1, 1,                // blockDim (x, y, z) ← 256 个线程!
    0,                        // sharedMem
    stream,                   // stream
    args,                     // args[]
    nullptr,                  // extra
);

cuLaunchKernel 是异步的——它只是把”启动这个 kernel”这条命令塞进 stream 队列,立刻返回。GPU 真正执行可能要等几微秒到几毫秒。

陷阱:for_num_elems(1) 启动的是 256 个线程,不是 1 个。block 固定 256,grid 按需要的元素数取上整。如果想真的单线程跑:

LaunchConfig { grid_dim: (1,1,1), block_dim: (1,1,1), shared_mem_bytes: 0 }

实际运行时:

INFO hello_constant: [PHASE 9/9 d] launching kernel kernel="hello_constant"

7. host 端时间线总结

[PHASE 9/9 a] host binary 启动
              ↓
              init_tracing()
              ↓
[PHASE 9/9 b] CudaContext::new(0)
              ├── cuInit(0)
              ├── cuDeviceGet(&dev, 0)
              └── cuCtxCreate_v2(&ctx, 0, dev)
              ↓
              DeviceBuffer::<i32>::zeroed(&stream, 1)
              ├── cuMemAlloc_v2(&dptr, 4)
              └── cuMemsetD8Async(dptr, 0, 4, stream)
              ↓
[PHASE 9/9 c] load_module_from_file("hello_constant.ptx")
              ├── 读文件
              ├── cuModuleLoad(...)
              │     ↓
              │     driver JIT:PTX → SASS(可能几十 ms,首次)
              │     ↓
              │     缓存到 ~/.nv/ComputeCache/
              ↓
              kernels::from_module(...) ← 类型化封装,无 driver 调用
              ↓
[PHASE 9/9 d] module.hello_constant(...)
              └── cuLaunchKernel(...)  ← 异步,立刻返回
              ↓
              (GPU 后台开始跑,host 继续往下)
              ↓
[PHASE 9/9 e] to_host_vec(&stream)  ← 下一篇讲

8. 关键观察

(1) 三层进程:cargo-oxide → cargo → rustc 编完之后 cargo 又 exec 出 hello_constant 二进制——你看到的 Output: 42 其实是孙进程输出的。

(2) PTX 在 host load 时才变成 SASS:这是为什么同一份 .ptx 可以跑在不同 GPU 上(向前兼容)。代价是首次 load 有几十 ms JIT 开销——生产环境通常会预先把 PTX 编成 cubin 缓存。

(3) device pointer 不能 deref:out_dev.cu_deviceptr() 返回的 CUdeviceptr 是 64-bit 整数,是显存地址。host 不能直接读写,必须通过 cuMem* API。

(4) cuLaunchKernel 是异步的:host 调完立刻返回,GPU 后台跑。如果之后立刻 host 端访问输出 buffer,会拿到旧值——必须先 cuStreamSynchronize(下一篇讲)。

9. 一句话总结

host binary 启动后,通过 cuda-core 这个 RAII 包装一步步调 CUDA driver API:cuInit + cuDeviceGet + cuCtxCreate 连 GPU,cuMemAlloc 分配显存,cuModuleLoad 触发 driver 内置 JIT 把 PTX 编成当前 GPU 的 SASS 并缓存,cuLaunchKernel 异步发射 grid × block × thread。整个流程类型安全(#[cuda_module] 宏帮你生成 launcher)、资源安全(RAII drop 自动释放)。最魔法的一步是 cuModuleLoad——PTX 这个虚拟 ISA 在这一刻才变成真正能跑在 GPU 上的机器码。

系列上一篇: cuda-oxide:hello-constant 拆解 09——llvm-export + llc 生成 PTX

评论区

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

cuda-oxide:hello-constant 拆解 10——host binary 启动 + CUDA driver 加载 PTX

0. 几个名词先说清楚

1. host binary 启动

2. CudaContext::new(0) — 连 GPU

3. DeviceBuffer::<i32>::zeroed(&stream, 1) — cudaMalloc

4. load_module_from_file — 最”魔法”的一步:PTX → SASS JIT

5. kernels::from_module(...) — 类型化封装

6. module.hello_constant(...) — cuLaunchKernel

7. host 端时间线总结

8. 关键观察

9. 一句话总结

2. `CudaContext::new(0)` — 连 GPU

3. `DeviceBuffer::<i32>::zeroed(&stream, 1)` — cudaMalloc

4. `load_module_from_file` — 最”魔法”的一步:PTX → SASS JIT

5. `kernels::from_module(...)` — 类型化封装

6. `module.hello_constant(...)` — cuLaunchKernel