#gpu #computation #gpgpu #wgpu #simple

easy-gpgpu

一个用于在GPU上计算的高层、易于使用的库

5次发布

0.1.6 2022年2月16日
0.1.5 2022年2月13日

#633 in 硬件支持

每月下载 28

MIT/Apache

63KB
912

Easy GPGPU

一个基于wgpu的高层、易于使用的gpgpu crate。它适用于在强大的GPU上进行非常大的计算

主要目标

  • 使通用计算非常简单
  • 尽可能容易地编写wgsl着色器
  • 自动处理绑定缓冲区

限制

  • 缓冲区可用的类型:bool、i32、u32、f32
  • 最大缓冲区字节大小:约134_000_000 (~33百万i32) 使用device.apply_on_vector可以扩展到10亿字节(26亿i32s)
  • 首次初始化设备需要时间(由于wgpu后端)

示例

重造wgpu的hello-compute(使用wgpu编写时为205 sloc)

use easy_gpgpu::*;
fn wgpu_hello_compute() {
    let mut device = Device::new();
    let v = vec![1u32, 4, 3, 295];
    let result = device.apply_on_vector(v.clone(), r"
    fn collatz_iterations(n_base: u32) -> u32{
        var n: u32 = n_base;
        var i: u32 = 0u;
        loop {
            if (n <= 1u) {
                break;
            }
            if (n % 2u == 0u) {
                n = n / 2u;
            }
            else {
                // Overflow? (i.e. 3*n + 1 > 0xffffffffu?)
                if (n >= 1431655765u) {   // 0x55555555u
                    return 4294967295u;   // 0xffffffffu
                }
                n = 3u * n + 1u;
            }
            i = i + 1u;
    }
    return i;
    }
    collatz_iterations(element)
    ");
    assert_eq!(result, vec![0, 2, 7, 55]);
}

=> 无绑定,无讨厌的global_id,无需使用低级API。

只需编写最少的wgsl着色器代码。

使用 .apply_on_vector 的最简单用法

use easy_gpgpu::*;
// create a device
let mut device = Device::new();
// create the vector we want to apply a computation on
let v1 = vec![1.0f32, 2.0, 3.0];
// the next line reads : for every element in v1, perform : element = element * 2.0
let v1 = device.apply_on_vector(v1, "element * 2.0");
println!("{v1:?}");

使用 .execute_shader_code 的用法

//First create a device :
use easy_gpgpu::*;
let mut device = Device::new();
// Then create some buffers, specify if you want to get their content after the execution :
let v1 = vec![1i32, 2, 3, 4, 5, 6];
// from a vector
device.create_buffer_from("v1", &v1, BufferUsage::ReadOnly, false);
// creates an empty buffer
device.create_buffer("output", BufferType::I32, v1.len(), BufferUsage::WriteOnly, true);
// Finaly, execute a shader :
let result = device.execute_shader_code(Dispatch::Linear(v1.len()), r"
fn main() {
    output[index] = v1[index] * 2;
}").into_iter().next().unwrap().unwrap_i32();
assert_eq!(result, vec![2i32, 4, 6, 8, 10, 12])

缓冲区在着色器中以创建设备时提供的名称可用。

index变量是通过使用Dispatch::Linear提供的(索引是u32)。

我们只指定了一个具有is_output: true的缓冲区,所以我们只得到一个输出向量。

我们只需使用.unwrap_i32()将数据解包为i32向量。

更多示例

最简单的方法来使用

pub fn simplest_apply() {
    let mut device = Device::new();
    let v1 = vec![1.0f32, 2.0, 3.0];
    // the next line reads : for every element in v1, perform : element = element * 2.0
    let v1 = device.apply_on_vector(v1, "element * 2.0");
    println!("{v1:?}");
}

使用先前创建的缓冲区的device.apply_on_vector示例

pub fn apply_with_buf() {
    let mut device = Device::new();
    let v1 = vec![2.0f32, 3.0, 5.0, 7.0, 11.0];
    let exponent = vec![3.0];
    device.create_buffer_from("exponent", &exponent, BufferUsage::ReadOnly, false);
    let cubes = device.apply_on_vector(v1, "pow(element, exponent[0u])");
    println!("{cubes:?}")
}

使用device.execute_shader_code的最简单示例:将向量的每个元素乘以2。

pub fn simplest_execute_shader() {
    let mut device = Device::new();
    let v1 = vec![1i32, 2, 3, 4, 5, 6];
    device.create_buffer_from("v1", &v1, BufferUsage::ReadOnly, false);
    device.create_buffer("output", BufferType::I32, v1.len(), BufferUsage::WriteOnly, true);
    let result = device.execute_shader_code(Dispatch::Linear(v1.len()), r"
    fn main() {
        output[index] = v1[index] * 2;
    }
    ").into_iter().next().unwrap().unwrap_i32();
    assert_eq!(result, vec![2, 4, 6, 8, 10, 12]);
}

使用device.execute_shader_code返回多个缓冲区的示例

pub fn multiple_returned_buffers() {
    let mut device = Device::new();
    let v = vec![1u32, 2, 3];
    let v2 = vec![3u32, 4, 5];
    let v3 = vec![7u32, 8, 9];
    // we specify is_output: true so this will be our first returned buffer
    device.create_buffer_from(
        "buf",
        &v,
        BufferUsage::ReadWrite,
        true
    );
    device.create_buffer_from(
        "buf2",
        &v2,
        BufferUsage::ReadOnly,
        false
    );
    // we specify is_output: true so this will be our second returned buffer
    device.create_buffer_from(
        "buf3",
        &v3,
        BufferUsage::ReadWrite,
        true
    );
    let mut result = device.execute_shader_code(Dispatch::Linear(v.len()), r"
        fn main() {
            buf[index] = buf[index] + buf2[index] + buf3[index];
            buf3[index] = buf[index] * buf2[index] * buf3[index];
        }
    ").into_iter();

    let sum = result.next().unwrap().unwrap_u32(); // first returned buffer
    let product = result.next().unwrap().unwrap_u32(); //second returned buffer
    println!("{:?}", sum);
    println!("{:?}", product);
}

使用自定义调度器访问全局_id变量的示例

pub fn global_id() {
    let mut device = Device::new();
    let vec = vec![2u32, 3, 5, 7, 11, 13, 17];
    // "vec" is actually a reserved keyword in wgsl.
    device.create_buffer_from("vec1", &vec, BufferUsage::ReadWrite, true);
    let result = device.execute_shader_code(Dispatch::Custom(1, vec.len() as u32, 1), r"
    fn main() {
        vec1[global_id.y] = vec1[global_id.y] + global_id.x + global_id.z;
    }
    ").into_iter().next().unwrap().unwrap_u32();
    // since the Dispatch was (1, 7, 1), the global_id.x and global.y are always 0
    // so our primes stay primes :
    assert_eq!(result, vec![2u32, 3, 5, 7, 11, 13, 17]);
}

一个完整的管道示例,如您所见,只需将向量乘以2就相当麻烦。

pub fn complete_pipeline() {
    let mut device = Device::new();
    let v1 = vec![1u32, 2, 3, 4, 5];
    device.create_buffer_from("v1", &v1, BufferUsage::ReadWrite, true);
    let shader_module = device.create_shader_module(Dispatch::Linear(v1.len()), "
    fn main() {
        v1[index] = v1[index] * 2u;
    }
    ");
    let mut commands = vec![];
    commands.push(Command::Shader(shader_module));
    commands.push(Command::Retrieve("v1"));
    device.execute_commands(commands);
    let result = device.get_buffer_data(vec!["v1"]).into_iter().next().unwrap().unwrap_u32();
    assert_eq!(result, vec![2u32, 4, 6, 8, 10]);
}

一个示例,我们在同一设备上执行两个着色器,使用相同的缓冲区。

pub fn reusing_device() {
    let mut device = Device::new();
    let v1 = vec![1i32, 2, 3, 4, 5, 6];
    device.create_buffer_from("v1", &v1, BufferUsage::ReadOnly, false);
    device.create_buffer("output", BufferType::I32, v1.len(), BufferUsage::WriteOnly, true);
    let result = device.execute_shader_code(Dispatch::Linear(v1.len()), r"
    fn main() {
        output[index] = v1[index] * 2;
    }
    ").into_iter().next().unwrap().unwrap_i32();
    assert_eq!(result, vec![2i32, 4, 6, 8, 10, 12]);

    let result2 = device.execute_shader_code(Dispatch::Linear(v1.len()), r"
    fn main() {
        output[index] = v1[index] * 10;
    }
    ").into_iter().next().unwrap().unwrap_i32();
    assert_eq!(result2, vec![10, 20, 30, 40, 50, 60]);
}

更多示例可以在文档中的 examples 模块中找到

有用的文档: wgsl

警告:本文档中描述的 wgsl 语言与该 crate 使用的语言不完全相同

(本 crate 使用的 wgsl 语言与 wgpu crate 中使用的语言相同)

-> 文档中的属性使用 @attribute 指定,而在此 crate 中使用 [[attribute]],所以 @group(0) @binding(0) 变为 [[group(0), binding(0)]]

还有一些其他细微的差异,但这份文档对所有 内置函数 都非常有用

依赖关系

~3–15MB
~183K SLoC