stensor 0.1.1

Cross-platform GPU tensor library with Slang and Rust.
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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

use nalgebra::Matrix3;
use slang_hal::{test_shader_compilation, Shader};
#[cfg(test)]
use {
    naga_oil::compose::NagaModuleDescriptor,
    wgpu::{ComputePipeline, Device},
};

#[derive(Copy, Clone, Debug, encase::ShaderType)]
#[repr(C)]
/// GPU representation of a 3x3 matrix QR decomposition.
///
/// See the [nalgebra](https://nalgebra.rs/docs/user_guide/decompositions_and_lapack#qr)
/// documentation for details on the QR decomposition.
pub struct GpuQR3 {
    /// The QR decomposition’s 3x3 unitary matrix.
    pub q: Matrix3<f32>,
    /// The QR decomposition’s 3x3 upper-triangular matrix.
    pub r: Matrix3<f32>,
}

#[derive(Shader)]
#[shader(src = "qr3.wgsl")]
/// Shader for computing the Singular Value Decomposition of 3x3 matrices.
pub struct WgQR3;

test_shader_compilation!(WgQR3);

impl WgQR3 {
    #[doc(hidden)]
    #[cfg(test)]
    pub fn tests(device: &Device) -> ComputePipeline {
        let test_kernel = r#"
 @group(0) @binding(0)
 var<storage, read_write> in: array<mat3x3<f32>>;
 @group(0) @binding(1)
 var<storage, read_write> out: array<QR>;

 @compute @workgroup_size(1, 1, 1)
 fn test(@builtin(global_invocation_id) invocation_id: vec3<u32>) {
     let i = invocation_id.x;
     out[i] = qr(in[i]);
 }
        "#;

        let src = format!("{}\n{}", Self::src(), test_kernel);
        let module = WgQR3::composer()
            .unwrap()
            .make_naga_module(NagaModuleDescriptor {
                source: &src,
                file_path: Self::FILE_PATH,
                ..Default::default()
            })
            .unwrap();
        slang_hal::utils::load_module(device, "test", module)
    }
}

#[cfg(test)]
mod test {
    use super::GpuQR3;
    use approx::{assert_relative_eq, relative_eq};
    use nalgebra::{DVector, Matrix3};
    use slang_hal::gpu::GpuInstance;
    use slang_hal::kernel::{CommandEncoderExt, KernelDispatch};
    use crate::tensor::GpuTensor;
    use wgpu::BufferUsages;

    #[futures_test::test]
    #[serial_test::serial]
    async fn gpu_qr3() {
        let gpu = GpuInstance::new().await.unwrap();
        let svd = super::WgQR3::tests(gpu.device());
        let mut encoder = gpu.device().create_command_encoder(&Default::default());

        const LEN: usize = 345;
        let matrices: DVector<Matrix3<f32>> = DVector::new_random(LEN);
        let inputs = GpuTensor::encase(gpu.device(), &matrices, BufferUsages::STORAGE);
        let result: GpuTensor<GpuQR3> = GpuTensor::vector_uninit_encased(
            gpu.device(),
            matrices.len() as u32,
            BufferUsages::STORAGE | BufferUsages::COPY_SRC,
        );
        let staging: GpuTensor<GpuQR3> = GpuTensor::vector_uninit_encased(
            gpu.device(),
            matrices.len() as u32,
            BufferUsages::MAP_READ | BufferUsages::COPY_DST,
        );

        // Dispatch the test.
        let mut pass = encoder.compute_pass("test", None);
        KernelDispatch::new(gpu.device(), &mut pass, &svd)
            .bind0([inputs.buffer(), result.buffer()])
            .dispatch(matrices.len() as u32);
        drop(pass); // Ensure the pass is ended before the encoder is borrowed again.

        staging.copy_from_encased(&mut encoder, &result);
        gpu.queue().submit(Some(encoder.finish()));

        // Check the result is correct.
        let gpu_result = staging.read_encased(gpu.device()).await.unwrap();
        let mut allowed_fails = 0;

        for (m, qr) in matrices.iter().zip(gpu_result.iter()) {
            let qr_na = m.qr();

            // NOTE: we allow about 1% of the decompositions to fail, to account for occasionally
            //       bad random matrices that will fail the test due to an unsuitable epsilon.
            //       Ideally this percentage should be kept as low as possible, but likely not
            //       removable entirely.
            if allowed_fails == matrices.len() * 2 / 100 {
                assert_relative_eq!(qr_na.q(), qr.q, epsilon = 1.0e-4);
                assert_relative_eq!(qr_na.r(), qr.r, epsilon = 1.0e-4);
            } else if !relative_eq!(qr_na.q(), qr.q, epsilon = 1.0e-4)
                || !relative_eq!(qr_na.r(), qr.r, epsilon = 1.0e-4)
            {
                allowed_fails += 1;
            }
        }

        println!("Num fails: {}/{}", allowed_fails, matrices.len());
    }
}