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

mod benchmark;
mod problem;
mod strategy;

pub use benchmark::bench;
pub use problem::{ContiguousProblem, problems};
pub use strategy::{ContiguousStrategy, strategies};

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

pub struct Category;

impl cubek_test_utils::Category for Category {
    type Problem = ContiguousProblem;
    type Strategy = ContiguousStrategy;

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

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

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

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

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

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

    /// A layout copy: reads the strided input and writes a fresh contiguous
    /// buffer of the same element count. No compute beyond addressing.
    fn work(&self, problem: &ContiguousProblem) -> Option<CategoryWork> {
        let dtype = f32::elem_type_native();
        let elems = problem.shape.iter().product::<usize>() * dtype.size();

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