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

cubecl_std/throughput/runners/
compute_direct.rs
