cubecl_std/throughput/runners/
memory_direct.rs

1use cubecl::prelude::*;
2use cubecl_core as cubecl;
3use cubecl_runtime::throughput::{KernelConfig, ThroughputKey};
4
5use crate::throughput::LaunchConfig;
6
7/// Per-buffer size, clamped to the device's maximum allocation.
8const TARGET_BYTES: usize = 512 * 1024 * 1024;
9
10/// Builds the copy kernel.
11pub fn build_kernel<R: Runtime>(
12    client: &ComputeClient<R>,
13    key: ThroughputKey,
14    config: LaunchConfig,
15) -> KernelConfig {
16    let client = client.clone();
17    let dtype = key.dtype;
18
19    let line_bytes = config.vector_size * dtype.size();
20
21    let max_alloc = client.properties().memory.max_page_size as usize;
22    let target = TARGET_BYTES.min(max_alloc);
23
24    let total_threads = config.cube_count * config.cube_dim;
25    let num_lines = (target / line_bytes).max(total_threads);
26    let bytes = num_lines * line_bytes;
27
28    let in_handle = client.empty(bytes);
29    let out_handle = client.empty(bytes);
30
31    let kernel = Box::new(move |iterations: usize| unsafe {
32        memory_direct_throughput::launch_unchecked(
33            &client,
34            CubeCount::Static(config.cube_count as u32, 1, 1),
35            CubeDim::new(&client, config.cube_dim),
36            config.vector_size,
37            BufferArg::from_raw_parts(in_handle.clone(), num_lines),
38            BufferArg::from_raw_parts(out_handle.clone(), num_lines),
39            iterations,
40            dtype.into(),
41        )
42    });
43
44    let ops_count = 2 * num_lines * config.vector_size;
45
46    KernelConfig { kernel, ops_count }
47}
48
49#[cube(launch_unchecked)]
50pub fn memory_direct_throughput<I: Numeric, N: Size>(
51    input: &[Vector<I, N>],
52    output: &mut [Vector<I, N>],
53    n_iter: usize,
54    #[define(I)] _dtype: StorageType,
55) {
56    let len = output.len();
57    let stride = CUBE_DIM as usize * CUBE_COUNT;
58
59    let steps = (len - ABSOLUTE_POS).div_ceil(stride).max(1);
60
61    for _ in 0..n_iter {
62        for step in 0..steps {
63            let idx = ABSOLUTE_POS + (step * stride);
64
65            if idx < len {
66                output[idx] = input[idx];
67            }
68        }
69    }
70}
cubecl_std/throughput/runners/memory_direct.rs

cubecl_std/throughput/runners/
memory_direct.rs
