cubek-std 0.3.0-pre.4

CubeK: Standard Library
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

mod benchmark;
mod problem;
mod strategy;

pub use benchmark::bench;
pub use problem::{UnaryProblem, problems};
pub use strategy::{UnaryStrategy, strategies};

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

pub struct Category;

impl cubek_test_utils::Category for Category {
    type Problem = UnaryProblem;
    type Strategy = UnaryStrategy;

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

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

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

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

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

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

    /// No honest op count: the kernel's cost is 256 rounds of `cos`, and a
    /// transcendental has no fixed multiply-add price to compare against
    /// `ComputeDirect`. Reads and writes are the elementwise shape: two
    /// operands in, one out.
    fn work(&self, problem: &UnaryProblem) -> Option<CategoryWork> {
        let dtype = f32::elem_type_native();
        let elems = problem.shape.iter().product::<usize>() * dtype.size();

        Some(CategoryWork {
            compute: None,
            bytes_read: 2 * elems,
            bytes_written: elems,
        })
    }
}