如何使用rust调用cuda代码

基本原理

使用已有rust cuda库 cust 来调用cuda代码。

1. 将cuda代码编译为.ptx或.fatbin文件
2. 使用cust读取cuda 编译的 .ptx 或者 .fatbin 文件，然后调用其中对应的cuda函数

示例

环境：ubuntu 22.04（rust和cuda要提前安装）

1. 书写gpu hello-world代码 hello.cu

#include <stdio.h>
#include <cuda_runtime.h>

extern "C"
__global__ void gpu_hello_world()
{
    printf("Hello world GPU!\n");
}

2. 编译为.ptx或者.fatbin

nvcc --fatbin hello.cu  //输出 hello.fatbin
nvcc --ptx hello.cu		//输出 hello.ptx

3. 书写rust代码 main.rs

use cust::{launch, module::Module, stream::{Stream, StreamFlags}};


static TEST_FATBIN: &[u8] = include_bytes!("../../hello.fatbin");
static TEST_PTX: &str = include_str!("../../hello.ptx");

fn main() {
    println!("Hello, world!");
    let _ctx = cust::quick_init().unwrap();
    //let module = Module::from_fatbin(TEST_FATBIN, &[]).unwrap();
    let module = Module::from_ptx(TEST_PTX, &[]).unwrap();
    let stream = Stream::new(StreamFlags::DEFAULT, None).unwrap();
    let kernel = module.get_function("gpu_hello_world").unwrap();
    unsafe {
        launch!(kernel<<<1,1,0,stream>>>()).unwrap();
    }
    stream.synchronize().unwrap();
}

Cargo.toml

[package]
name = "cuda-demo"
version = "0.1.0"
edition = "2021"

[dependencies]
cust = { version = "0.3" }  #引入cust

4. cargo build 编译运行

>>> cargo build 
>>> ./target/debug/cuda-demo 
Hello, world!
Hello from GPU!

总结

cust库做了一层封装，使rust调用cuda代码变的十分简单，只需要专注于业务代码就好。

上面的cuda nvcc的编译流程也可以通过rust的build.rs来实现，让代码更加完善。