cudarc: 在 CUDA 工具包之上的最小化和安全 API

在 crates.io 和 docs.rs 上检查 cudarc。

在以下方面提供安全抽象

1. CUDA 驱动 API
2. NVRTC API
3. cuRAND API
4. cuBLAS API
5. cuBLASLt API

预-alpha 状态，预期会有破坏性更改，并且不是所有 CUDA 函数都包含安全封装。 欢迎为未包含的功能做出贡献！

设计

目标是

1. 尽可能安全（由于 ffi 和 async，仍然会有很多不安全的地方）
2. 尽可能舒适
3. 允许混合高级 safe API 和低级 sys API

为此，每个封装都有三个级别（默认情况下导出安全 API）

use cudarc::driver::{safe, result, sys};
use cudarc::nvrtc::{safe, result, sys};
use cudarc::cublas::{safe, result, sys};
use cudarc::cublaslt::{safe, result, sys};
use cudarc::curand::{safe, result, sys};

其中

1. sys 是 bindgen 生成的原始 ffi API
2. result 是围绕 sys 的小型封装，以从每个函数返回 Result
3. safe 是围绕 result/sys 提供安全抽象的封装

强烈建议坚持使用安全 API

API 预览

创建新设备并将数据传输到 GPU 很简单

let dev = cudarc::driver::CudaDevice::new(0)?;

// allocate buffers
let inp = dev.htod_copy(vec![1.0f32; 100])?;
let mut out = dev.alloc_zeros::<f32>(100)?;

您还可以使用 nvrtc API 在运行时编译内核

let ptx = cudarc::nvrtc::compile_ptx("
extern \"C\" __global__ void sin_kernel(float *out, const float *inp, const size_t numel) {
    unsigned int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < numel) {
        out[i] = sin(inp[i]);
    }
}")?;

// and dynamically load it into the device
dev.load_ptx(ptx, "my_module", &["sin_kernel"])?;

cudarc 提供了一种非常简单的接口来启动内核，元组是参数！

let sin_kernel = dev.get_func("my_module", "sin_kernel").unwrap();
let cfg = LaunchConfig::for_num_elems(100);
unsafe { sin_kernel.launch(cfg, (&mut out, &inp, 100usize)) }?;

当然，完成操作后将事物复制回主机也很简单

let out_host: Vec<f32> = dev.dtoh_sync_copy(&out)?;
assert_eq!(out_host, [1.0; 100].map(f32::sin));

许可协议

双重许可以与 Rust 项目兼容。

许可协议为 Apache License，版本 2.0 http://www.apache.org/licenses/LICENSE-2.0 或 MIT 许可证 http://opensource.org/licenses/MIT，任选其一。此文件不得根据这些条款进行复制、修改或分发。