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

t4a_cubecl_std/throughput/runners/
compute_direct.rs
