cubecl_std/throughput/runners/
memory_direct.rs

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

cubecl_std/throughput/runners/
memory_direct.rs
