burn-cubecl 0.22.0-pre.2

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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

use crate::{
    CubeRuntime, CubeTuneId, kernel::interpolate::execute_interpolate,
    kernel::interpolate::map_mode, tensor::CubeTensor,
};
use burn_backend::cubecl::dtype_to_elem_type;
use burn_backend::ops::InterpolateOptions;
use cubecl::tune::{LocalTuner, Tunable, TunableSet, TuneGroup, local_tuner};
use cubek::interpolate::{
    definition::TileSize,
    launch::{InterpolateAutotuneKey, InterpolateStrategy},
    routines::{
        BlueprintStrategy, GlobalMemoryRoutine, GlobalMemoryStrategy, SharedMemoryRoutine,
        SharedMemoryStrategy,
    },
};

/// Interpolate operation with autotuning. This benchmarks multiple strategies and selects the best one at runtime.
pub fn interpolate_autotune<R: CubeRuntime>(
    input: CubeTensor<R>,
    output_size: [usize; 2],
    options: InterpolateOptions,
) -> CubeTensor<R> {
    let client = input.client.clone();

    static TUNER: LocalTuner<InterpolateAutotuneKey, CubeTuneId> = local_tuner!();

    let tunables = TUNER.init(|| {
        const PRIORITY: i8 = 0;

        let global_memory =
            TuneGroup::<InterpolateAutotuneKey>::new("global_memory", |_key| PRIORITY);
        let shared_memory =
            TuneGroup::<InterpolateAutotuneKey>::new("shared_memory", |_key| PRIORITY);

        let mut set = TunableSet::new(create_key::<R>, input_gen::<R>);

        let tile_sizes: [TileSize; 16] = [
            // Square shapes
            TileSize::new(8, 8),
            TileSize::new(16, 16),
            TileSize::new(32, 32),
            // Rectangular shapes
            TileSize::new(8, 16),
            TileSize::new(16, 8),
            TileSize::new(16, 32),
            TileSize::new(32, 16),
            // Flat horizontal shapes
            TileSize::new(1, 64),
            TileSize::new(1, 128),
            TileSize::new(1, 256),
            TileSize::new(2, 128),
            TileSize::new(1, 512),
            TileSize::new(1, 1024),
            // Flat vertical shapes
            TileSize::new(64, 1),
            TileSize::new(128, 1),
            TileSize::new(256, 1),
        ];

        for tile_size in tile_sizes {
            let name = format!(
                "global_memory_tile_size_{}_{}",
                tile_size.height(),
                tile_size.width()
            );
            set = set.with(
                Tunable::new(&name, move |(input, output_size, options)| {
                    execute_interpolate::<R>(
                        input,
                        output_size,
                        options,
                        InterpolateStrategy::GlobalMemoryStrategy(BlueprintStrategy::<
                            GlobalMemoryRoutine,
                        >::Inferred(
                            GlobalMemoryStrategy { tile_size },
                        )),
                    )
                })
                .group(&global_memory, |_key| PRIORITY),
            );

            let name = format!(
                "shared_memory_tile_size_{}_{}",
                tile_size.height(),
                tile_size.width()
            );
            set = set.with(
                Tunable::new(&name, move |(input, output_size, options)| {
                    execute_interpolate::<R>(
                        input,
                        output_size,
                        options,
                        InterpolateStrategy::SharedMemoryStrategy(BlueprintStrategy::<
                            SharedMemoryRoutine,
                        >::Inferred(
                            SharedMemoryStrategy { tile_size },
                        )),
                    )
                })
                .group(&shared_memory, |_key| PRIORITY),
            );
        }

        set
    });

    TUNER.execute(
        &CubeTuneId::new(&client, &input.device),
        &client,
        tunables,
        (input, output_size, options),
    )
}

fn create_key<R: CubeRuntime>(
    (input, output_size, options): &(CubeTensor<R>, [usize; 2], InterpolateOptions),
) -> InterpolateAutotuneKey {
    let elem_input = dtype_to_elem_type(input.dtype);
    let elem_output = dtype_to_elem_type(input.dtype);
    let mode = map_mode(options.mode.clone());

    InterpolateAutotuneKey::generate(
        elem_input,
        elem_output,
        mode,
        input.meta.shape(),
        output_size,
    )
}

fn input_gen<R: CubeRuntime>(
    _key: &InterpolateAutotuneKey,
    (input, output_size, options): &(CubeTensor<R>, [usize; 2], InterpolateOptions),
) -> (CubeTensor<R>, [usize; 2], InterpolateOptions) {
    (input.clone(), *output_size, options.clone())
}