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