cubek-interpolate 0.3.0

CubeK: Interpolation
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

//! Benchmark catalogue for `cubek-interpolate`.
mod benchmark;
mod correctness;
mod problem;
mod strategy;

pub use benchmark::bench;
pub use correctness::InterpolateCorrectness;
pub use problem::{problems, problems_scaled};
pub use strategy::{BenchTarget, BenchTier, every_strategy, strategies, strategies_at};

use cubecl::benchmark::TimingMethod;
use cubecl::prelude::*;
use cubek_test_utils::{CatalogEntry, CategoryWork, ComputeWork, RunSamples};

use crate::{
    InterpolateStrategy,
    definition::{InterpolateCost, InterpolateProblem},
};

pub struct Category;

impl cubek_test_utils::Category for Category {
    type Problem = InterpolateProblem;
    type Strategy = InterpolateStrategy;

    fn id(&self) -> &'static str {
        "interpolate"
    }

    fn label(&self) -> &'static str {
        "Interpolate"
    }

    fn problems(&self) -> Vec<CatalogEntry<InterpolateProblem>> {
        problems()
    }

    fn strategies(&self) -> Vec<CatalogEntry<InterpolateStrategy>> {
        strategies(BenchTarget::Gpu)
    }

    fn bench(
        &self,
        strategy: &InterpolateStrategy,
        problem: &InterpolateProblem,
        num_samples: usize,
    ) -> Result<RunSamples, String> {
        bench(strategy, problem, num_samples)
    }

    fn timing_method(&self) -> TimingMethod {
        cubek_test_utils::timing_method(TimingMethod::Device)
    }
    fn correctness(
        &self,
    ) -> Option<
        &dyn cubek_test_utils::Correctness<
            Problem = InterpolateProblem,
            Strategy = InterpolateStrategy,
        >,
    > {
        Some(&InterpolateCorrectness)
    }
}

/// Every spatial extent divided by four, so one CPU sweep of the full strategy catalogue finishes
/// in minutes rather than hours. Sixteen times fewer output elements per problem, with the batch,
/// channel and scale-direction regimes left intact.
const CPU_DIVISOR: usize = 4;

/// The same sweep at CPU scale.
///
/// It is a separate category rather than a switch on the device because the two are not
/// comparable: a median from this set says which strategy wins, not what the operation costs.
pub struct CpuCategory;

impl cubek_test_utils::Category for CpuCategory {
    type Problem = InterpolateProblem;
    type Strategy = InterpolateStrategy;

    fn id(&self) -> &'static str {
        "interpolate_cpu"
    }

    fn label(&self) -> &'static str {
        "Interpolate (CPU scale)"
    }

    fn problems(&self) -> Vec<CatalogEntry<InterpolateProblem>> {
        problems_scaled(CPU_DIVISOR)
    }

    fn strategies(&self) -> Vec<CatalogEntry<InterpolateStrategy>> {
        strategies(BenchTarget::Cpu)
    }

    fn bench(
        &self,
        strategy: &InterpolateStrategy,
        problem: &InterpolateProblem,
        num_samples: usize,
    ) -> Result<RunSamples, String> {
        bench(strategy, problem, num_samples)
    }

    fn timing_method(&self) -> TimingMethod {
        cubek_test_utils::timing_method(TimingMethod::Device)
    }

    fn correctness(
        &self,
    ) -> Option<
        &dyn cubek_test_utils::Correctness<
            Problem = InterpolateProblem,
            Strategy = InterpolateStrategy,
        >,
    > {
        Some(&InterpolateCorrectness)
    }

    fn work(&self, problem: &InterpolateProblem) -> Option<CategoryWork> {
        let dtype = f32::elem_type_native();
        let cost = InterpolateCost::new(problem.clone(), dtype);
        let (bytes_read, bytes_written) = cost.traffic();

        Some(CategoryWork {
            compute: Some(ComputeWork {
                ops: cost.compute_ops(),
                key: cost.compute_key(),
            }),
            bytes_read,
            bytes_written,
        })
    }
}