cubecl_std/throughput/runners/
memory_direct.rs

1use cubecl::prelude::*;
2use cubecl_core as cubecl;
3use cubecl_runtime::throughput::{KernelConfig, MemoryAccess, ThroughputKey};
4
5use crate::throughput::{
6    LaunchConfig,
7    memory_probe::{self, MemoryProbe},
8};
9
10/// Builds the copy kernel, moving `working_set` bytes per pass: half read out
11/// of the input buffer, half written into the output one.
12pub fn build_kernel<R: Runtime>(
13    client: &ComputeClient<R>,
14    key: ThroughputKey,
15    config: LaunchConfig,
16    working_set: usize,
17) -> KernelConfig {
18    let client = client.clone();
19    let dtype = key.dtype();
20
21    let line_bytes = config.vector_size * dtype.size();
22    let probe = MemoryProbe::new(&client, config, line_bytes, MemoryAccess::Copy, working_set);
23
24    let in_handle = client.empty(probe.buffer_bytes);
25    memory_probe::prime(&client, &in_handle, probe.pool_lines, config, dtype);
26    let out_handle = client.empty(probe.buffer_bytes);
27
28    let sample = Box::new(move |iterations: usize| {
29        let start = cubecl_common::profile::Instant::now();
30        unsafe {
31            memory_direct_throughput::launch_unchecked(
32                &client,
33                CubeCount::Static(probe.cube_count as u32, 1, 1),
34                CubeDim::new(&client, config.cube_dim),
35                config.vector_size,
36                BufferArg::from_raw_parts(in_handle.clone(), probe.pool_lines),
37                BufferArg::from_raw_parts(out_handle.clone(), probe.pool_lines),
38                probe.window_lines,
39                iterations,
40                probe.blocked,
41                dtype,
42            )
43        };
44        let _ = cubecl_core::future::block_on(client.sync());
45        start.elapsed()
46    });
47
48    // One pass moves the window twice: once in, once out.
49    let ops_count = 2 * probe.window_lines * config.vector_size;
50
51    KernelConfig { sample, ops_count }
52}
53
54#[cube(launch_unchecked)]
55pub fn memory_direct_throughput<I: Numeric, N: Size>(
56    input: &[Vector<I, N>],
57    output: &mut [Vector<I, N>],
58    window: usize,
59    n_iter: usize,
60    #[comptime] blocked: bool,
61    #[define(I)] _dtype: ElemType,
62) {
63    let len = output.len();
64    let stride = CUBE_DIM as usize * CUBE_COUNT;
65
66    // From `window` alone rather than from `window - ABSOLUTE_POS`, which
67    // underflows for a thread past the end of a window smaller than the launch.
68    // High threads get one step too many and the bounds check drops it.
69    let steps = window.div_ceil(stride).max(1);
70
71    // Each pass copies the *next* window of the buffers, not the same one
72    // again. A window read repeatedly would be served from cache after the
73    // first pass, and every working set below the cache would report cache
74    // bandwidth instead of what a kernel of that size moves; coming back to a
75    // window only after a whole buffer of traffic keeps it cold.
76    //
77    // It also keeps the addresses moving. A window small enough that every
78    // thread copies a single line would otherwise be loop-invariant, and the
79    // compiler is free to sink such a copy out of the loop and perform it once
80    // — leaving the probe reporting a bandwidth the hardware never moved.
81    let mut start = 0;
82    let mut wrap = 0;
83
84    for _ in 0..n_iter {
85        for step in 0..steps {
86            // Coalesced spreads one step's addresses across adjacent threads,
87            // which is only fast where those threads share a real plane. A
88            // CPU worker has no such neighbour, so it instead gets a run of
89            // `steps` lines entirely its own.
90            let base = if blocked {
91                ABSOLUTE_POS * steps + step
92            } else {
93                ABSOLUTE_POS + (step * stride)
94            };
95
96            if base < window {
97                let mut idx = start + base;
98                if idx >= len {
99                    idx -= len;
100                }
101
102                output[idx] = input[idx];
103            }
104        }
105
106        start += window;
107        // Back to the beginning, one line further along each time round, so a
108        // window that fills the whole buffer still moves between passes.
109        if start + window > len {
110            wrap += 1;
111            if wrap >= window {
112                wrap = 0;
113            }
114            start = wrap;
115        }
116    }
117}
cubecl_std/throughput/runners/memory_direct.rs

cubecl_std/throughput/runners/
memory_direct.rs
