cubecl-matmul 0.8.1

CubeCL Matrix Multiplication Kernels Engine
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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89

use cubecl_core::{CubeElement, Runtime, client::ComputeClient, prelude::Float};
use cubecl_std::tensor::TensorHandle;

use crate::tests::test_utils::Sample;

pub(crate) struct MatmulTestCase {
    pub m: usize,
    pub k: usize,
    pub n: usize,
    pub batch: usize,
}

impl MatmulTestCase {
    pub(crate) fn matmul_cpu<R: Runtime, F: Float + CubeElement>(
        &self,
        lhs: &TensorHandle<R, F>,
        rhs: &TensorHandle<R, F>,
        client: &ComputeClient<R::Server>,
    ) -> Vec<F> {
        let lhs_binding = &client.read_one_tensor(lhs.handle.clone().copy_descriptor(
            &lhs.shape,
            &lhs.strides,
            size_of::<F>(),
        ));
        let rhs_binding = &client.read_one_tensor(rhs.handle.clone().copy_descriptor(
            &rhs.shape,
            &rhs.strides,
            size_of::<F>(),
        ));

        let lhs = F::from_bytes(lhs_binding);
        let rhs = F::from_bytes(rhs_binding);

        self.matmul_cpu_algorithm(lhs, rhs)
    }

    fn matmul_cpu_algorithm<F: Float + CubeElement>(&self, lhs: &[F], rhs: &[F]) -> Vec<F> {
        let mut out = vec![F::from_int(0); self.batch * self.m * self.n];
        let lhs_batch_offset = self.m * self.k;
        let rhs_batch_offset = self.k * self.n;
        let out_batch_offset = self.m * self.n;

        for b in 0..self.batch {
            for i in 0..self.m {
                for j in 0..self.n {
                    for k_ in 0..self.k {
                        let lhs_value = lhs[b * lhs_batch_offset + i * self.k + k_];
                        let rhs_value = rhs[b * rhs_batch_offset + j + k_ * self.n];

                        let result = lhs_value * rhs_value;

                        out[b * out_batch_offset + i * self.n + j] += result;
                    }
                }
            }
        }

        out
    }

    pub(crate) fn random_lhs<R: Runtime, F: Float + CubeElement + Sample>(
        &self,
        client: &ComputeClient<R::Server>,
    ) -> TensorHandle<R, F> {
        self.random_tensor(client, vec![self.batch, self.m, self.k])
    }

    pub(crate) fn random_rhs<R: Runtime, F: Float + CubeElement + Sample>(
        &self,
        client: &ComputeClient<R::Server>,
    ) -> TensorHandle<R, F> {
        self.random_tensor(client, vec![self.batch, self.k, self.n])
    }

    pub(crate) fn empty_out<R: Runtime, F: Float + CubeElement + Sample>(
        &self,
        client: &ComputeClient<R::Server>,
    ) -> TensorHandle<R, F> {
        TensorHandle::empty(client, vec![self.batch, self.m, self.n])
    }

    pub(crate) fn random_tensor<R: Runtime, F: Float + CubeElement + Sample>(
        &self,
        client: &ComputeClient<R::Server>,
        shape: Vec<usize>,
    ) -> TensorHandle<R, F> {
        F::sample::<R>(client, &shape, 999)
    }
}