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
90
91
92
93
94
95
96
97
98
99
100
101
102
103

use std::fmt::Display;

use cubecl_core::{CubeElement, Runtime, prelude::Float};

use crate::{
    MatmulInputHandle,
    kernels::naive,
    tests::{
        naive::utils::MatmulTestCase,
        test_utils::{Sample, assert_equals_approx},
    },
};
use cubecl_std::tensor::TensorHandle;

pub fn test_small<R: Runtime, F: Float + CubeElement + Display + Sample>(device: &R::Device) {
    let case = MatmulTestCase {
        m: 64,
        k: 64,
        n: 64,
        batch: 1,
    };

    test_simple::<R, F>(case, device);
}

pub fn test_odd<R: Runtime, F: Float + CubeElement + Display + Sample>(device: &R::Device) {
    let case = MatmulTestCase {
        m: 1,
        k: 101,
        n: 255,
        batch: 1,
    };

    test_simple::<R, F>(case, device);
}

pub fn test_large<R: Runtime, F: Float + CubeElement + Display + Sample>(device: &R::Device) {
    let case = MatmulTestCase {
        m: 256,
        k: 256,
        n: 256,
        batch: 1,
    };

    test_simple::<R, F>(case, device);
}

pub fn test_with_check_bounds<R: Runtime, F: Float + CubeElement + Display + Sample>(
    device: &R::Device,
) {
    let case = MatmulTestCase {
        m: 60,
        k: 60,
        n: 60,
        batch: 1,
    };

    test_simple::<R, F>(case, device);
}

pub fn test_with_batches<R: Runtime, F: Float + CubeElement + Display + Sample>(
    device: &R::Device,
) {
    let case = MatmulTestCase {
        m: 64,
        k: 64,
        n: 64,
        batch: 3,
    };

    test_simple::<R, F>(case, device);
}

fn test_simple<R: Runtime, F: Float + CubeElement + Display + Sample>(
    case: MatmulTestCase,
    device: &R::Device,
) {
    let client = R::client(device);
    let lhs = case.random_lhs::<R, F>(&client);
    let rhs = case.random_rhs::<R, F>(&client);

    let expected = case.matmul_cpu::<R, F>(&lhs, &rhs, &client);

    let out: TensorHandle<R, F> = case.empty_out(&client);
    naive::launch::<R, F, F>(
        &client,
        MatmulInputHandle::Normal(lhs),
        MatmulInputHandle::Normal(rhs),
        &out.as_ref(),
    )
    .unwrap();

    if let Err(e) = assert_equals_approx::<R, F>(
        &client,
        out.handle,
        &out.shape,
        &out.strides,
        &expected,
        10e-4,
    ) {
        panic!("{}", e);
    }
}