t4a_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
91    for _ in 0..n_iter {
92        for step in 0..steps {
93            // Coalesced spreads one step's addresses across adjacent threads,
94            // which is only fast where those threads share a real plane. A
95            // CPU worker has no such neighbour, so it instead gets a run of
96            // `steps` lines entirely its own.
97            let base = if blocked {
98                ABSOLUTE_POS * steps + step
99            } else {
100                ABSOLUTE_POS + (step * stride)
101            };
102
103            if base < window {
104                let mut idx = start + base;
105                if idx >= len {
106                    idx -= len;
107                }
108
109                output[idx] = input[idx];
110            }
111        }
112
113        start += window;
114        if start >= len {
115            start -= len;
116        }
117    }
118}
t4a_cubecl_std/throughput/runners/memory_direct.rs

t4a_cubecl_std/throughput/runners/
memory_direct.rs
