编写你的第一个内核

本节将引导你完成 cuda-oxide 的安装、创建项目、编写 GPU 内核并运行——全部使用纯 Rust 完成。

---

安装 cargo-oxide

如果你还没有安装构建工具，请执行：

cargo install --git https://github.com/NVlabs/cuda-oxide.git cargo-oxide

验证你的环境是否配置正确：

cargo oxide doctor

这会检查兼容的 GPU、CUDA toolkit、LLVM 以及代码生成后端。在继续之前请修复它报告的所有问题（详情参见安装）。

---

创建项目

使用 cargo oxide new 脚手架创建一个新项目：

cargo oxide new my_first_kernel
cd my_first_kernel

这会生成一个可以直接运行的项目：

my_first_kernel/
├── Cargo.toml          # 依赖于 cuda-device、cuda-host、cuda-core
├── rust-toolchain.toml # 固定所需的 nightly 工具链
└── src/
    └── main.rs          # 内核 + 主机代码在同一文件中

编译并运行：

cargo oxide run

你应该会看到 PASSED: all 1024 elements correct。生成的模板是一个向量加法内核——这是一个不错的起点，但让我们看看更有趣的内容。

---

内核剖析

下面是一个带点变化的向量加法：逐元素加法被提取到了一个普通的辅助函数中。内核和辅助函数与主机代码共存于同一个文件中：

use cuda_device::{cuda_module, kernel, thread, DisjointSlice};
use cuda_core::{CudaContext, DeviceBuffer, LaunchConfig};

/// 普通辅助函数——无需注解。
/// 编译器会自动发现它，因为 `vecadd` 调用了它。
fn add(a: f32, b: f32) -> f32 {
    a + b
}

#[cuda_module]
mod kernels {
    use super::*;

    #[kernel]
    pub fn vecadd(a: &[f32], b: &[f32], mut c: DisjointSlice<f32>) {
        let idx = thread::index_1d();
        let i = idx.get();
        if let Some(c_elem) = c.get_mut(idx) {
            *c_elem = add(a[i], b[i]);
        }
    }
}

fn main() {
    let ctx = CudaContext::new(0).unwrap();
    let stream = ctx.default_stream();

    const N: usize = 1024;
    let a_host: Vec<f32> = (0..N).map(|i| i as f32).collect();
    let b_host: Vec<f32> = (0..N).map(|i| (i * 2) as f32).collect();

    let a_dev = DeviceBuffer::from_host(&stream, &a_host).unwrap();
    let b_dev = DeviceBuffer::from_host(&stream, &b_host).unwrap();
    let mut c_dev = DeviceBuffer::<f32>::zeroed(&stream, N).unwrap();

    let module = kernels::load(&ctx).expect("Failed to load embedded module");
    module
        .vecadd(
            &stream,
            LaunchConfig::for_num_elems(N as u32),
            &a_dev,
            &b_dev,
            &mut c_dev,
        )
        .unwrap();

    let c_host = c_dev.to_host_vec(&stream).unwrap();
    let errors = (0..N)
        .filter(|&i| (c_host[i] - (a_host[i] + b_host[i])).abs() > 1e-5)
        .count();

    if errors == 0 {
        println!("PASSED: all {} elements correct", N);
    } else {
        eprintln!("FAILED: {} errors", errors);
        std::process::exit(1);
    }
}

这里发生了很多事情。让我们逐一拆解关键部分。

单源编译

内核和主机代码位于同一个文件中，通过一次 cargo 命令调用即可编译。代码生成后端会拦截编译过程，将 #[kernel] 函数路由到 MIR→PTX 管线，而其他所有内容则委托给标准 LLVM 处理。最终生成的二进制文件包含原生主机代码以及内嵌的设备构件。

#[kernel]

将函数标记为可启动的内核入口点——相当于 GPU 的 main。该函数通过以下管线编译为 PTX：

Rust 源码 → MIR → Pliron IR → LLVM IR → PTX

同一个函数对主机编译器也是可见的（用于类型检查），但其函数体永远不会在 CPU 上被调用。

设备函数（自动发现）

上面的 add 辅助函数没有任何注解。当编译器处理 #[kernel] 时，它会遍历调用图并自动发现内核调用的每个函数。这些函数会被编译为 PTX 设备函数并由后端内联——你不需要手动标记它们。

备注

#[device] 属性确实存在，但服务于不同的目的：它将函数标记为独立的设备编译根（用于构建供 C++ 消费的 Rust 设备库），或在 #[device] extern "C" { ... } 块中用于声明外部设备函数，以便通过 CUDA C++ LTOIR 进行 FFI 互操作。对于从内核调用的私有辅助函数，你不需要 #[device]。

#[cuda_module]

#[cuda_module] 包裹内联的内核模块并生成带类型的主机 API：

let module = kernels::load(&ctx)?;
module.vecadd(&stream, LaunchConfig::for_num_elems(N as u32), &a_dev, &b_dev, &mut c_dev)?;

加载器从主机二进制文件中读取内嵌的设备构件，缓存内核函数句柄，并将每个 #[kernel] 暴露为一个 Rust 方法。方法签名镜像了内核签名，设备切片映射为 DeviceBuffer 借用。

load_kernel_module 和 cuda_launch! 仍然作为更低层的 API 可用，用于手动加载旁载构件和自定义启动代码。

参数标量化

切片以 (ptr, len) 组件的形式跨越主机/设备 ABI——主机将它们作为两个内核参数传递，设备编译器在入口块中重新组装切片。按值传递的结构体和闭包则作为一个 byval .param 传递，因此主机端将整个聚合体作为单个槽位推送（这与启动器实际所做的相匹配，避免了按字段逐个声明时的不一致）。所有这些完全是透明的——内核签名看起来仍然像普通的 Rust：

Host:   module.vecadd(..., &data, ...)
          → 提取切片的 (ptr, len)，传递两个参数

PTX:    .entry kernel(.param .u64 ptr, .param .u64 len, ...)
          → 接收扁平的切片参数

Device: 内核体看到统一的 &[T] 切片
          → 编译器在入口处重建

动态结构体布局

当你向 GPU 传递结构体时，cuda-oxide 会向 rustc 查询每个字段的精确字节偏移量，并在设备端使用显式填充重建布局。这意味着 #[repr(C)] 不是必需的——普通的 Rust 结构体可以直接使用，即使在 HMM（GPU 直接访问主机内存）下也是如此。

---

异步编程

对于多内核管线或并发工作负载，cuda-oxide 提供了基于 Tokio 构建的异步执行模型。让我们脚手架一个异步项目并逐步了解其中的差异。

创建异步项目

cargo oxide new my_async_kernel --async
cd my_async_kernel
cargo oxide run

--async 标志会生成一个预配置了 tokio 和 cuda-async 依赖的项目。

完整示例

下面是生成的异步 vecadd 模板（为可读性做了少量格式调整）：

use cuda_device::{cuda_module, kernel, thread, DisjointSlice};
use cuda_async::device_context::init_device_contexts;
use cuda_async::device_operation::DeviceOperation;
use cuda_core::LaunchConfig;

#[cuda_module]
mod kernels {
    use super::*;

    #[kernel]
    pub fn vecadd(a: &[f32], b: &[f32], mut c: DisjointSlice<f32>) {
        let idx = thread::index_1d();
        let i = idx.get();
        if let Some(c_elem) = c.get_mut(idx) {
            *c_elem = a[i] + b[i];
        }
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    use cuda_async::device_box::DeviceBox;
    use cuda_core::memory::{malloc_async, memcpy_dtoh_async, memcpy_htod_async};
    use std::mem;

    // 1. 初始化设备上下文映射（默认设备 0，1 个设备）。
    //    轮询流池在首次使用时惰性创建。
    init_device_contexts(0, 1)?;

    // 2. 从异步设备上下文加载内嵌的内核模块。
    let module = kernels::load_async(0)?;

    const N: usize = 1024;
    let a_host: Vec<f32> = (0..N).map(|i| i as f32).collect();
    let b_host: Vec<f32> = (0..N).map(|i| (i * 2) as f32).collect();

    // 3. 分配设备内存并拷贝主机数据。
    let (a_dev, b_dev, mut c_dev) =
        cuda_async::device_context::with_cuda_context(0, |ctx| {
            let stream = ctx.default_stream();
            let bytes = N * mem::size_of::<f32>();
            unsafe {
                let a = malloc_async(stream.cu_stream(), bytes).unwrap();
                let b = malloc_async(stream.cu_stream(), bytes).unwrap();
                let c = malloc_async(stream.cu_stream(), bytes).unwrap();
                memcpy_htod_async(a, a_host.as_ptr(), bytes, stream.cu_stream()).unwrap();
                memcpy_htod_async(b, b_host.as_ptr(), bytes, stream.cu_stream()).unwrap();
                stream.synchronize().unwrap();
                (
                    DeviceBox::<[f32]>::from_raw_parts(a, N, 0),
                    DeviceBox::<[f32]>::from_raw_parts(b, N, 0),
                    DeviceBox::<[f32]>::from_raw_parts(c, N, 0),
                )
            }
        })?;

    // 4. 启动——返回一个惰性的 DeviceOperation，此时尚未进行 GPU 工作。
    module
        .vecadd_async(
            LaunchConfig::for_num_elems(N as u32),
            &a_dev,
            &b_dev,
            &mut c_dev,
        )?
        .sync()?;  // 阻塞直到 GPU 完成。

    // 5. 将结果拷贝回主机。
    let mut c_host = vec![0.0f32; N];
    cuda_async::device_context::with_cuda_context(0, |ctx| {
        let stream = ctx.default_stream();
        unsafe {
            memcpy_dtoh_async(
                c_host.as_mut_ptr(),
                c_dev.cu_deviceptr(),
                N * mem::size_of::<f32>(),
                stream.cu_stream(),
            )
            .unwrap();
            stream.synchronize().unwrap();
        }
    })?;

    // 6. 验证。
    let errors = (0..N)
        .filter(|&i| (c_host[i] - (a_host[i] + b_host[i])).abs() > 1e-5)
        .count();

    if errors == 0 {
        println!("PASSED: all {} elements correct", N);
    } else {
        eprintln!("FAILED: {} errors", errors);
        std::process::exit(1);
    }

    Ok(())
}

与同步版本的差异

内核本身是完全相同的——异步只改变了你在主机端启动和管理 GPU 工作的方式。

{kernel}_async 代替 {kernel}

返回一个惰性的 DeviceOperation，而不是立即启动。在显式调度之前不会执行任何 GPU 工作。这让你可以在提交资源之前先构建计算图。

init_device_contexts(default_device, num_devices)

初始化线程本地的设备上下文映射，设置默认 GPU 序号和多设备使用的容量。轮询流池在首次使用时惰性创建。操作随后以轮询方式分配到不同的流，无需手动管理流即可最大化 GPU 占用率。

DeviceBox 代替 DeviceBuffer

线程安全的设备内存包装器。与流池配合工作，通过 malloc_async 支持异步分配。

.sync() 与 .await

.sync() 会阻塞调用线程直到 GPU 完成——当你在主机端没有其他事情可做时使用它。.await 挂起当前的 Tokio 任务，让其他任务在等待期间继续执行——当你需要并发执行主机工作或有多个 GPU 管线并行时使用它。

and_then / zip!

使用 .and_then(|result| next_op) 链接依赖操作。使用 zip!(op_a, op_b) 并发运行独立操作——两者都会被提交到流池，合并的结果在两者都完成后可用。这些组合器让你以声明式的方式表达复杂的多内核管线。

小技巧

要查看更完整的异步示例，请参阅 async_mlp——一个包含 and_then 链接、zip! 并行分配以及跨并发批次共享 Arc 权重的多内核前向传播（GEMM、MatVec、ReLU）示例。在 cuda-oxide 工作空间中运行 cargo oxide run async_mlp。