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

cubecl_std/throughput/runners/
memory_direct.rs
