oxgpu 0.1.0

A lightweight GPU compute library built on wgpu
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49

#[cfg(test)]
mod tests {
    use oxgpu::{Buffer, ComputeKernel, Context};
    // use wgpu; // No longer needed for setup

    #[tokio::test]
    async fn test_vector_add() {
        let ctx = Context::new().await.unwrap();

        let data_a = vec![1.0f32, 2.0, 3.0, 4.0];
        let data_b = vec![10.0f32, 20.0, 30.0, 40.0];

        let buf_a = Buffer::from_slice(&ctx, &data_a).await;
        let buf_b = Buffer::from_slice(&ctx, &data_b).await;
        let buf_c = Buffer::<f32>::zeros(&ctx, 4).await;

        let shader = r#"
            @group(0) @binding(0) var<storage, read> a: array<f32>;
            @group(0) @binding(1) var<storage, read> b: array<f32>;
            @group(0) @binding(2) var<storage, read_write> c: array<f32>;

            @compute @workgroup_size(1)
            fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
                let i = global_id.x;
                if (i < 4u) {
                    c[i] = a[i] + b[i];
                }
            }
        "#;

        let kernel = ComputeKernel::builder()
            .source(shader)
            .entry_point("main")
            .add_storage_read(0)
            .add_storage_read(1)
            .add_storage_read_write(2)
            .build(&ctx)
            .await
            .unwrap();

        // Run
        // Note: passing references to buffers, which now implement KernelArgument
        kernel.run(&ctx, 4, &[&buf_a, &buf_b, &buf_c]);

        // Read back
        let result = buf_c.read(&ctx).await.unwrap();
        assert_eq!(result, vec![11.0, 22.0, 33.0, 44.0]);
    }
}