cubek-reduce 0.3.0

CubeK: Reduce Kernels
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

//! Benchmark catalogue for `cubek-reduce`.

mod benchmark;
mod correctness;
mod problem;
mod strategy;

pub use benchmark::bench;
pub use correctness::ReduceCorrectness;
pub use problem::{ReduceBenchKind, ReduceBenchPrecision, ReduceProblem, precisions, problems};
pub use strategy::strategies;

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

use crate::ReduceStrategy;
use crate::eval::cpu_reference::{accumulation_dtype, output_dtype_for};
use crate::launch::ReduceDtypes;
use crate::routines::ReduceCost;

pub struct Category;

impl cubek_test_utils::Category for Category {
    type Problem = ReduceProblem;
    type Strategy = ReduceStrategy;

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

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

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

    fn strategies(&self) -> Vec<CatalogEntry<ReduceStrategy>> {
        strategies()
    }

    fn bench(
        &self,
        strategy: &ReduceStrategy,
        problem: &ReduceProblem,
        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 = ReduceProblem, Strategy = ReduceStrategy>,
    > {
        Some(&ReduceCorrectness)
    }

    /// Scaled by the passes the harness makes over the cost model: a two-launch
    /// reduction reads and writes twice, which is the harness's doing.
    fn work(&self, problem: &ReduceProblem) -> Option<CategoryWork> {
        let value_dtype = problem.precision.dtype();
        let input_elems: usize = problem.shape.iter().product();
        let reduce_len = problem.shape[problem.axis];

        let cost = ReduceCost {
            reduce_len,
            reduce_count: input_elems / reduce_len,
            instruction: problem.config,
            dtypes: ReduceDtypes {
                input: value_dtype,
                output: output_dtype_for(&problem.config, value_dtype),
                accumulation: accumulation_dtype(),
            },
        };

        let (read_passes, write_passes) = match problem.kind {
            ReduceBenchKind::Single => (1, 1),
            ReduceBenchKind::TwoLaunch => (2, 2),
            ReduceBenchKind::Fused => (1, 2),
        };

        let (bytes_read, bytes_written) = cost.traffic();

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