Rust CUDA绑定生成库bindgen_cuda的使用:高效生成Rust与CUDA交互的FFI绑定代码

Rust CUDA绑定生成库bindgen_cuda的使用:高效生成Rust与CUDA交互的FFI绑定代码

简介

bindgen_cuda是一个类似于bindgen的库,它可以帮助自动生成CUDA内核源文件的绑定,使它们更容易直接从Rust中使用。

PTX包含示例

假设你有一个CUDA源文件src/cuda.cu

__global__ void cuda_hello(){
    printf("Hello World from GPU!\n");
}

你可以将bindgen_cuda添加为构建依赖:

cargo add --build bindgen_cuda

然后创建build.rs文件:

fn main() {
    let builder = bindgen_cuda::Builder::default();
    let bindings = builder.build_ptx().unwrap();
    bindings.write("src/lib.rs");
}

这将创建一个包含以下代码的源文件:

pub const CUDA: &str = include_str!(concat!(env!("OUT_DIR"), "/cuda.ptx"));

然后你可以使用cudarc这样的库直接在Rust代码中使用PTX。

原始CUDA调用示例

你也可以构建一个静态库,然后在build.rs中链接它,以便直接用C代码调用CUDA。

src/cuda.cu文件:

__global__ void cuda_hello(){
    printf("Hello World from GPU!\n");
}

int run() {
    cuda_hello<<<1,1>>>(); 
    return 0;
}

然后编写build.rs

fn main() {
    let builder = bindgen_cuda::Builder::default();
    builder.build_lib("libcuda.a");
    println!("cargo:rustc-link-lib=cuda");
}

然后你可以在src/lib.rs中通过FFI接口调用:

extern "C" {
    fn cuda_hello();
}
fn main(){
    unsafe{ cuda_hello();}
}

完整示例代码

使用PTX的完整示例

  1. 创建CUDA源文件src/cuda.cu
// CUDA内核函数
__global__ void cuda_hello() {
    printf("Hello World from GPU!\n");
}
  1. 创建build.rs
fn main() {
    // 创建默认构建器
    let builder = bindgen_cuda::Builder::default();
    // 构建PTX并写入src/lib.rs
    let bindings = builder.build_ptx().unwrap();
    bindings.write("src/lib.rs");
}
  1. 创建src/main.rs
use cudarc::driver::CudaDevice;

fn main() {
    // 创建CUDA设备
    let dev = CudaDevice::new(0).unwrap();
    // 加载PTX模块
    let ptx = include_str!("../src/lib.rs");
    let module = dev.load_ptx(ptx, "cuda").unwrap();
    // 获取CUDA函数
    let func = module.function("cuda_hello").unwrap();
    // 调用CUDA函数
    unsafe { func.call([1, 1, 1], 1).unwrap() };
}

使用原始CUDA调用的完整示例

  1. 创建CUDA源文件src/cuda.cu
__global__ void cuda_hello() {
    printf("Hello World from GPU!\n");
}

// 封装函数
int run() {
    cuda_hello<<<1,1>>>(); 
    return 0;
}
  1. 创建build.rs
fn main() {
    // 创建默认构建器
    let builder = bindgen_cuda::Builder::default();
    // 构建静态库
    builder.build_lib("libcuda.a");
    // 告诉cargo链接库
    println!("cargo:rustc-link-lib=cuda");
}
  1. 创建src/main.rs
// 声明外部函数
extern "C" {
    fn run() -> i32;
}

fn main() {
    // 安全调用CUDA函数
    unsafe {
        let result = run();
        println!("CUDA function returned: {}", result);
    }
}

许可证

MIT License


1 回复

Rust CUDA绑定生成库bindgen_cuda的使用指南

完整示例demo

以下是基于内容中提供的向量加法示例的完整实现,包含CUDA核函数和完整的Rust集成:

首先需要创建CUDA核函数文件vector_add.cu:

// vector_add.cu - CUDA核函数实现
__global__ void vector_add(float *a, float *b, float *c, int n) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) {
        c[i] = a[i] + b[i];
    }
}

然后编译为PTX格式:

nvcc -ptx vector_add.cu -o vector_add.ptx

接下来是完整的Rust实现:

// main.rs
mod cuda_runtime;
use cuda_runtime::*;
use std::ffi::c_void;
use std::fs;

fn main() {
    const N: usize = 1024;
    const BLOCK_SIZE: usize = 256;
    
    // 主机数据
    let mut h_a = vec![1.0f32; N];
    let mut h_b = vec![2.0f32; N];
    let mut h_c = vec![0.0f32; N];
    
    unsafe {
        // 设备指针
        let mut d_a: *mut f32 = std::ptr::null_mut();
        let mut d_b: *mut f32 = std::ptr::null_mut();
        let mut d_c: *mut f32 = std::ptr::null_mut();
        
        // 1. 初始化CUDA
        cudaDeviceReset().unwrap();
        
        // 2. 分配设备内存
        cudaMalloc(&mut d_a as *mut *mut f32, N * 4).unwrap();
        cudaMalloc(&mut d_b as *mut *mut f32, N * 4).unwrap();
        cudaMalloc(&mut d_c as *mut *mut f32, N * 4).unwrap();
        
        // 3. 拷贝数据到设备
        cudaMemcpy(
            d_a, 
            h_a.as_ptr() as *const _, 
            N * 4, 
            cudaMemcpyKind::cudaMemcpyHostToDevice
        ).unwrap();
        
        cudaMemcpy(
            d_b, 
            h_b.as_ptr() as *const _, 
            N * 4, 
            cudaMemcpyKind::cudaMemcpyHostToDevice
        ).unwrap();
        
        // 4. 加载PTX模块
        let mut module: *mut c_void = std::ptr::null_mut();
        let ptx_code = fs::read_to_string("vector_add.ptx").unwrap();
        cudaModuleLoadData(&mut module, ptx_code.as_ptr() as *const _).unwrap();
        
        // 5. 获取核函数
        let mut kernel: *mut c_void = std::ptr::null_mut();
        cudaModuleGetFunction(&mut kernel, module, b"vector_add\0".as_ptr() as *const _).unwrap();
        
        // 6. 设置核函数参数
        let mut args = [&d_a as *const _ as *mut c_void,
                        &d_b as *const _ as *mut c_void,
                        &d_c as *const _ as *mut c_void,
                        &N as *const _ as *mut c_void];
        
        // 7. 启动核函数
        cudaLaunchKernel(
            kernel,
            (N as u32 + BLOCK_SIZE as u32 - 1) / BLOCK_SIZE as u32, // 网格大小
            1,
            1,
            BLOCK_SIZE as u32, // 块大小
            1,
            1,
            0,
            std::ptr::null_mut(),
            args.as_mut_ptr(),
            std::ptr::null_mut()
        ).unwrap();
        
        // 8. 同步设备
        cudaDeviceSynchronize().unwrap();
        
        // 9. 拷贝结果回主机
        cudaMemcpy(
            h_c.as_mut_ptr() as *mut _, 
            d_c, 
            N * 4, 
            cudaMemcpyKind::cudaMemcpyDeviceToHost
        ).unwrap();
        
        // 10. 释放资源
        cudaFree(d_a).unwrap();
        cudaFree(d_b).unwrap();
        cudaFree(d_c).unwrap();
        cudaModuleUnload(module).unwrap();
    }
    
    // 验证结果
    println!("前10个结果: {:?}", &h_c[0..10]);
    assert!(h_c.iter().all(|&x| (x - 3.0).abs() < 1e-6));
}

项目结构说明

完整的项目结构应该如下:

.
├── Cargo.toml
├── src
│   ├── main.rs
│   └── cuda_runtime.rs (由bindgen_cuda生成)
└── vector_add.cu
└── vector_add.ptx (由nvcc编译生成)

构建和运行步骤

  1. 首先使用bindgen_cuda生成CUDA运行时绑定
  2. 编译CUDA核函数: nvcc -ptx vector_add.cu -o vector_add.ptx
  3. 使用Cargo运行项目: cargo run

这个完整示例展示了如何使用bindgen_cuda从Rust调用CUDA核函数进行GPU加速计算,包含了从内存分配、数据传输到核函数调用的完整流程。

回到顶部