burn-jit 0.16.1

Generic backend that can be compiled just-in-time to any shader language target
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

use crate::{tensor::JitTensor, FloatElement, JitAutotuneKey, JitRuntime};
use burn_tensor::{DType, Shape};
use core::fmt::Debug;
use cubecl::AutotuneKey;
use serde::{Deserialize, Serialize};
use std::{cmp::max, hash::Hash};

#[derive(Hash, Eq, PartialEq, Debug, Clone, Serialize, Deserialize, AutotuneKey)]
/// Autotune key representative of matmul versions
pub struct MatmulAutotuneKey {
    round: bool,     // True when all matmul dims are multiples of 64
    broadcast: bool, // True when there are differences in batch size
    #[autotune(anchor)]
    m: usize,
    #[autotune(anchor)]
    k: usize,
    #[autotune(anchor)]
    n: usize,
    #[autotune(anchor(max = 256))]
    batch: usize,
    dtype: DType,
}

impl MatmulAutotuneKey {
    fn from_shape(lhs_shape: &Shape, rhs_shape: &Shape, dtype: DType) -> Self {
        let ndims = lhs_shape.num_dims();
        let m = lhs_shape.dims[ndims - 2];
        let k = lhs_shape.dims[ndims - 1];
        let n = rhs_shape.dims[ndims - 1];

        let mut broadcast = false;
        let mut batch_product_lhs = 1;
        let mut batch_product_rhs = 1;

        for b in 0..ndims - 2 {
            batch_product_lhs *= lhs_shape.dims[b];
            batch_product_rhs *= rhs_shape.dims[b];
            if lhs_shape.dims[b] != rhs_shape.dims[b] {
                broadcast = true;
            }
        }
        let batch_product = max(batch_product_lhs, batch_product_rhs);

        let round = m % 64 == 0 && k % 64 == 0 && n % 64 == 0;

        Self::new(round, broadcast, m, k, n, batch_product, dtype)
    }
}

pub(crate) fn create_key<R: JitRuntime, E: FloatElement>(
    lhs: &JitTensor<R>,
    rhs: &JitTensor<R>,
    _out: &JitTensor<R>,
) -> JitAutotuneKey {
    JitAutotuneKey::Matmul(MatmulAutotuneKey::from_shape(
        &lhs.shape,
        &rhs.shape,
        E::dtype(),
    ))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn matmul_autotune_key_all_same_and_round() {
        let lhs_shape: Shape = [4, 512, 512].into();
        let rhs_shape: Shape = [4, 512, 512].into();
        let key = MatmulAutotuneKey::from_shape(&lhs_shape, &rhs_shape, DType::F32);

        assert!(key.round);
        assert!(!key.broadcast);
        assert_eq!(key.m, 512);
        assert_eq!(key.k, 512);
        assert_eq!(key.n, 512);
    }

    #[test]
    fn matmul_autotune_key_all_different() {
        let lhs_shape: Shape = [2, 3, 511, 512].into();
        let rhs_shape: Shape = [3, 2, 512, 513].into();
        let key = MatmulAutotuneKey::from_shape(&lhs_shape, &rhs_shape, DType::F32);

        assert!(!key.round);
        assert!(key.broadcast);
        assert_eq!(key.m, 512);
        assert_eq!(key.k, 512);
        assert_eq!(key.n, 1024);
        assert_eq!(key.batch, 8);
    }

    #[test]
    fn matmul_autotune_key_large_batch() {
        let lhs_shape: Shape = [128, 512, 511, 512].into();
        let rhs_shape: Shape = [200, 400, 512, 513].into();
        let key = MatmulAutotuneKey::from_shape(&lhs_shape, &rhs_shape, DType::F32);

        assert_eq!(key.batch, 256);
    }
}