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

cubecl_std/throughput/runners/
compute_direct.rs
