cubecl_std/throughput/runners/
compute_direct.rs

1use cubecl::prelude::*;
2use cubecl_core::{self as cubecl, frontend::fma, ir::ElemType};
3use cubecl_runtime::throughput::{KernelConfig, ThroughputKey};
4
5use crate::throughput::LaunchConfig;
6
7pub fn build_kernel<R: Runtime>(
8    client: &ComputeClient<R>,
9    key: ThroughputKey,
10    config: LaunchConfig,
11) -> KernelConfig {
12    let client = client.clone();
13    let dtype = key.dtype;
14
15    let use_fma = matches!(dtype, ElemType::Float(_));
16
17    let kernel = Box::new(move |iterations: usize| unsafe {
18        let out = client.empty(config.vector_size * dtype.size());
19
20        compute_direct_throughput::launch_unchecked(
21            &client,
22            CubeCount::Static(config.cube_count as u32, 1, 1),
23            CubeDim::new(&client, config.cube_dim),
24            config.vector_size,
25            BufferArg::from_raw_parts(out, 1),
26            iterations,
27            use_fma,
28            dtype.into(),
29        )
30    });
31
32    // `CHAINS` independent accumulators per lane, each retiring one fma (two flops) or one mul.
33    let ops_per_chain = if use_fma { 2 } else { 1 };
34    let ops_count =
35        ops_per_chain * CHAINS * config.cube_count * config.cube_dim * config.vector_size;
36
37    KernelConfig { kernel, ops_count }
38}
39
40/// Independent accumulator chains per lane to hide arithmetic latency.
41const CHAINS: usize = 4;
42
43#[cube(launch_unchecked)]
44pub fn compute_direct_throughput<I: Numeric, N: Size>(
45    output: &mut [Vector<I, N>],
46    n_iter: usize,
47    #[comptime] use_fma: bool,
48    #[define(I)] _dtype: StorageType,
49) {
50    let tid = I::cast_from(ABSOLUTE_POS);
51
52    let mut b = Vector::<I, N>::empty();
53    let mut c = Vector::<I, N>::empty();
54
55    let mut s0 = Vector::<I, N>::empty();
56    let mut s1 = Vector::<I, N>::empty();
57    let mut s2 = Vector::<I, N>::empty();
58    let mut s3 = Vector::<I, N>::empty();
59
60    // Give every lane and chain a distinct seed to prevent folding.
61    let lanes = b.vector_size();
62    #[unroll]
63    for lane in 0..lanes {
64        let offset = I::cast_from(lane);
65        b.insert(lane, tid + offset + I::cast_from(1));
66        c.insert(lane, tid + offset);
67
68        s0.insert(lane, offset + I::cast_from(1));
69        s1.insert(lane, offset + I::cast_from(2));
70        s2.insert(lane, offset + I::cast_from(3));
71        s3.insert(lane, offset + I::cast_from(4));
72    }
73
74    for _ in 0..n_iter {
75        s0 = step(s0, b, c, use_fma);
76        s1 = step(s1, b, c, use_fma);
77        s2 = step(s2, b, c, use_fma);
78        s3 = step(s3, b, c, use_fma);
79    }
80
81    let sum = s0 + s1 + s2 + s3;
82
83    if ABSOLUTE_POS == 0 {
84        output[0] = sum;
85    }
86}
87
88/// Retires one arithmetic op per chain: an fma (two flops) for floats, otherwise a mul
89/// (the slowest integer op, giving a lower bound).
90#[cube]
91fn step<I: Numeric, N: Size>(
92    s: Vector<I, N>,
93    b: Vector<I, N>,
94    c: Vector<I, N>,
95    #[comptime] use_fma: bool,
96) -> Vector<I, N> {
97    if use_fma { fma(s, b, c) } else { s * b }
98}
cubecl_std/throughput/runners/compute_direct.rs

cubecl_std/throughput/runners/
compute_direct.rs
