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

use super::{
    algorithm::{Algorithm, ImplicitCmmaConv},
    precision::ConvPrecision,
};
use crate::JitRuntime;
use cubecl::linalg::matmul::components::{
    stage::CommonStageInput,
    tile::{accelerated::Accelerated, TileMatmulFamily},
    MatmulSelection, MatmulSize,
};

pub struct ConvSelection {
    pub matmul: MatmulSelection,
}

pub trait ConvSelector<A: Algorithm> {
    fn select_kernel<R: JitRuntime, CS: ConvPrecision>(plane_dim: u32) -> (A::Selection, A::Input);
}

/// Large m stage size for the usual case where `batch_size * out_h * out_w` is significantly larger
/// than `out_channels`
pub struct Large;
/// Balanced stage size for cases where `batch_size * out_h * out_w` is relatively small and `k` or
/// `out_channels` is relatively large
pub struct Balanced;

impl ConvSelector<ImplicitCmmaConv> for Large {
    fn select_kernel<R: JitRuntime, CS: ConvPrecision>(
        plane_dim: u32,
    ) -> (
        <ImplicitCmmaConv as Algorithm>::Selection,
        <ImplicitCmmaConv as Algorithm>::Input,
    ) {
        let selection = MatmulSelection {
            tile: MatmulSize {
                m: 16,
                n: 16,
                k: 16,
            },
            num_stagess: MatmulSize { m: 8, n: 4, k: 2 },
            plane_dim,
        };
        let config_input = CommonStageInput {
            tile: Accelerated::input(selection.tile),
            num_stages: selection.num_stagess,
        };

        let selection = ConvSelection { matmul: selection };

        (selection, config_input)
    }
}

impl ConvSelector<ImplicitCmmaConv> for Balanced {
    fn select_kernel<R: JitRuntime, CS: ConvPrecision>(
        plane_dim: u32,
    ) -> (
        <ImplicitCmmaConv as Algorithm>::Selection,
        <ImplicitCmmaConv as Algorithm>::Input,
    ) {
        let selection = MatmulSelection {
            tile: MatmulSize {
                m: 16,
                n: 16,
                k: 16,
            },
            num_stagess: MatmulSize { m: 4, n: 2, k: 4 },
            plane_dim,
        };
        let config_input = CommonStageInput {
            tile: Accelerated::input(selection.tile),
            num_stages: selection.num_stagess,
        };

        let selection = ConvSelection { matmul: selection };

        (selection, config_input)
    }
}