cubecl_std/throughput/runners/
memory_write.rs

1use cubecl::prelude::*;
2use cubecl_core as cubecl;
3use cubecl_runtime::throughput::{KernelConfig, MemoryAccess, ThroughputKey};
4
5use crate::throughput::{LaunchConfig, memory_probe::MemoryProbe};
6
7/// Builds the write-only streaming kernel, moving `working_set` bytes per
8/// pass, all of them written.
9///
10/// This is [`memory_direct`](super::memory_direct) with the load removed. The
11/// copy kernel moves a line in and a line back out, and counts both
12/// directions in `ops_count`, so what it reports is total traffic across the
13/// memory interface. That is the right ceiling for a kernel that also reads
14/// what it writes, and the wrong one for a kernel that only writes: an RNG
15/// fill, a memset, a broadcast. Those legitimately exceed the copy figure,
16/// because half of the copy's traffic is a direction they never use.
17///
18/// Reported `ops_count` is the write count alone.
19pub fn build_kernel<R: Runtime>(
20    client: &ComputeClient<R>,
21    key: ThroughputKey,
22    config: LaunchConfig,
23    working_set: usize,
24) -> KernelConfig {
25    let client = client.clone();
26    let dtype = key.dtype();
27
28    let line_bytes = config.vector_size * dtype.size();
29    let probe = MemoryProbe::new(
30        &client,
31        config,
32        line_bytes,
33        MemoryAccess::Write,
34        working_set,
35    );
36
37    let out_handle = client.empty(probe.buffer_bytes);
38
39    let sample = Box::new(move |iterations: usize| {
40        let start = cubecl_common::profile::Instant::now();
41        unsafe {
42            memory_write_throughput::launch_unchecked(
43                &client,
44                CubeCount::Static(probe.cube_count as u32, 1, 1),
45                CubeDim::new(&client, config.cube_dim),
46                config.vector_size,
47                BufferArg::from_raw_parts(out_handle.clone(), probe.pool_lines),
48                probe.window_lines,
49                iterations,
50                probe.blocked,
51                dtype,
52            )
53        };
54        let _ = cubecl_core::future::block_on(client.sync());
55        start.elapsed()
56    });
57
58    // Writes only, no `2 *`. That factor is the whole difference from the copy.
59    let ops_count = probe.window_lines * config.vector_size;
60
61    KernelConfig { sample, ops_count }
62}
63
64#[cube(launch_unchecked)]
65pub fn memory_write_throughput<I: Numeric, N: Size>(
66    output: &mut [Vector<I, N>],
67    window: usize,
68    n_iter: usize,
69    #[comptime] blocked: bool,
70    #[define(I)] _dtype: ElemType,
71) {
72    let len = output.len();
73    let stride = CUBE_DIM as usize * CUBE_COUNT;
74
75    // From `window` alone rather than from `window - ABSOLUTE_POS`, which
76    // underflows for a thread past the end of a window smaller than the launch.
77    // High threads get one step too many and the bounds check drops it.
78    let steps = window.div_ceil(stride).max(1);
79
80    // Read once, write everywhere: the reverse of `memory_read`, which reads
81    // everywhere and writes once. `n_iter` is a scalar the kernel compiler
82    // sees only at launch time, so it cannot fold the stores into a single
83    // known-pattern fill, and the per-lane offset keeps the written line from
84    // being a uniform value even within one call.
85    let seed = I::cast_from(n_iter);
86    let mut line = Vector::<I, N>::empty();
87    let lanes = line.vector_size();
88    #[unroll]
89    for lane in 0..lanes {
90        line.insert(lane, seed + I::cast_from(lane));
91    }
92
93    // Each pass writes the *next* window of the buffer, not the same one
94    // again. A window written repeatedly would stay resident in cache after
95    // the first pass, and every working set below the cache would report
96    // cache bandwidth instead of what a kernel of that size moves; coming
97    // back to a window only after a whole buffer of traffic keeps it cold.
98    //
99    // It also keeps the addresses moving. A window small enough that every
100    // thread writes a single line would otherwise be loop-invariant, and the
101    // compiler is free to sink such a store out of the loop and perform it
102    // once, leaving the probe reporting a bandwidth the hardware never moved.
103    let mut start = 0;
104    let mut wrap = 0;
105
106    for _ in 0..n_iter {
107        for step in 0..steps {
108            // Coalesced spreads one step's addresses across adjacent threads,
109            // which is only fast where those threads share a real plane. A
110            // CPU worker has no such neighbour, so it instead gets a run of
111            // `steps` lines entirely its own.
112            let base = if blocked {
113                ABSOLUTE_POS * steps + step
114            } else {
115                ABSOLUTE_POS + (step * stride)
116            };
117
118            if base < window {
119                let mut idx = start + base;
120                if idx >= len {
121                    idx -= len;
122                }
123
124                output[idx] = line;
125            }
126        }
127
128        start += window;
129        // Back to the beginning, one line further along each time round, so a
130        // window that fills the whole buffer still moves between passes.
131        if start + window > len {
132            wrap += 1;
133            if wrap >= window {
134                wrap = 0;
135            }
136            start = wrap;
137        }
138    }
139}
cubecl_std/throughput/runners/memory_write.rs

cubecl_std/throughput/runners/
memory_write.rs
