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
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70

#[cfg(test)]
mod tests {
    use oxgpu::{Buffer, ComputeKernel, Context};

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

        // 2. Buffer Creation (from slice)
        let rows = 4usize;
        let cols = 4usize;
        let input_data: Vec<f32> = (0..rows * cols).map(|x| x as f32).collect();
        let buf_input = Buffer::from_slice(&ctx, &input_data).await;

        // 3. Buffer Creation (zeros)
        let buf_output = Buffer::<f32>::zeros(&ctx, rows * cols).await;

        // 4. Buffer Creation (write after init)
        let scale_factor_data = [2.0f32];
        let buf_scale = Buffer::from_slice(&ctx, &[0.0f32]).await; // Init with dummy
        buf_scale.write(&ctx, &scale_factor_data); // Overwrite with actual value

        // 5. Kernel Builder
        let shader = r#"
            @group(0) @binding(0) var<storage, read> input: array<f32>;
            @group(0) @binding(1) var<storage, read_write> output: array<f32>;
            @group(0) @binding(2) var<storage, read> scale: array<f32>;

            @compute @workgroup_size(1, 1)
            fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
                let x = global_id.x;
                let y = global_id.y;
                let width = 4u;
                
                if (x < width && y < 4u) {
                    let index = y * width + x;
                    output[index] = input[index] * scale[0];
                }
            }
        "#;

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

        // 6. Run with Flexible Dispatch (2D)
        // Global size is (4, 4), which matches our data
        kernel.run(&ctx, (4, 4), &[&buf_input, &buf_output, &buf_scale]);

        // 7. Read back
        let result = buf_output.read(&ctx).await.unwrap();

        // Verification
        let expected: Vec<f32> = input_data.iter().map(|&x| x * 2.0).collect();
        assert_eq!(result, expected);

        println!(
            "Matrix scaling test passed: Input {:?} * 2.0 = {:?}",
            input_data, result
        );
    }
}