cudarc：CUDA 工具箱的极简和安全的 API

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

对以下进行封装

1. CUDA 驱动 API
2. NVRTC API
3. cuRAND API
4. cuBLAS API
5. cuBLASLt API
6. NCCL 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};
use cudarc::nccl::{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许可证，您可任选其一。未经许可，不得复制、修改或分发本文件。