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

mod benchmark;
mod problem;
mod strategy;

pub use benchmark::bench;
pub use problem::{MemcpyAsyncProblem, problems};
pub use strategy::{CopyStrategyEnum, strategies};

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

pub struct Category;

impl cubek_test_utils::Category for Category {
    type Problem = MemcpyAsyncProblem;
    type Strategy = CopyStrategyEnum;

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

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

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

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

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

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

    /// Copy-shaped: the kernel streams `data_count` elements through shared memory
    /// and folds them into a `window_size`-sized accumulator that becomes the
    /// output, so the global write is the small window, not the full stream.
    fn work(&self, problem: &MemcpyAsyncProblem) -> Option<CategoryWork> {
        let dtype = f32::elem_type_native();
        let elem_size = dtype.size();

        Some(CategoryWork {
            compute: None,
            bytes_read: problem.data_count * elem_size,
            bytes_written: problem.window_size * elem_size,
        })
    }
}