t4a-cubecl-std 0.11.0

CubeCL Standard Library.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114

use cubecl::prelude::*;
use cubecl_core::{self as cubecl, frontend::fma, ir::ElemType};
use cubecl_runtime::throughput::{KernelConfig, ThroughputError};

use crate::throughput::{LaunchConfig, verify::verify};

pub fn build_kernel(
    client: &Client,
    dtype: ElemType,
    config: LaunchConfig,
) -> Result<KernelConfig, ThroughputError> {
    let client = client.clone();

    let use_fma = matches!(dtype, ElemType::Float(_));
    let out = client.empty(config.vector_size * dtype.size());

    let (verifier, written) = (client.clone(), out.clone());
    let sample = Box::new(move |iterations: usize| {
        let start = cubecl_common::profile::Instant::now();
        unsafe {
            compute_direct_throughput::launch_unchecked(
                &client,
                CubeCount::Static(config.cube_count as u32, 1, 1),
                config.cube_dim,
                config.vector_size,
                BufferArg::from_raw_parts(out.clone(), 1),
                iterations,
                use_fma,
                dtype,
            )
        };
        // A failure is not this sync's to report: `verify` asked the output.
        let _ = cubecl_core::future::block_on(client.sync());
        start.elapsed()
    });
    verify(&verifier, &sample, &written)?;

    // `CHAINS` independent accumulators per lane, each retiring one fma (two flops) or one mul.
    let ops_per_chain = if use_fma { 2 } else { 1 };
    let ops_count = ops_per_chain
        * CHAINS
        * config.cube_count
        * config.cube_dim.num_elems() as usize
        * config.vector_size;

    Ok(KernelConfig {
        sample,
        ops_count,
        min_iterations: 1,
    })
}

/// Independent accumulator chains per lane to hide arithmetic latency.
const CHAINS: usize = 4;

/// Vectors the kernel keeps live across its loop: the chains and their two operands.
pub(crate) const LIVE_VECTORS: usize = CHAINS + 2;

#[cube(launch_unchecked)]
pub fn compute_direct_throughput<I: Numeric, N: Size>(
    output: &mut [Vector<I, N>],
    n_iter: usize,
    #[comptime] use_fma: bool,
    #[define(I)] _dtype: ElemType,
) {
    let tid = I::cast_from(ABSOLUTE_POS);

    let mut b = Vector::<I, N>::empty();
    let mut c = Vector::<I, N>::empty();

    let mut s0 = Vector::<I, N>::empty();
    let mut s1 = Vector::<I, N>::empty();
    let mut s2 = Vector::<I, N>::empty();
    let mut s3 = Vector::<I, N>::empty();

    // Give every lane and chain a distinct seed to prevent folding.
    let lanes = b.vector_size();
    #[unroll]
    for lane in 0..lanes {
        let offset = I::cast_from(lane);
        b.insert(lane, tid + offset + I::cast_from(1));
        c.insert(lane, tid + offset);

        s0.insert(lane, offset + I::cast_from(1));
        s1.insert(lane, offset + I::cast_from(2));
        s2.insert(lane, offset + I::cast_from(3));
        s3.insert(lane, offset + I::cast_from(4));
    }

    for _ in 0..n_iter {
        s0 = step(s0, b, c, use_fma);
        s1 = step(s1, b, c, use_fma);
        s2 = step(s2, b, c, use_fma);
        s3 = step(s3, b, c, use_fma);
    }

    let sum = s0 + s1 + s2 + s3;

    if ABSOLUTE_POS == 0 {
        output[0] = sum;
    }
}

/// Retires one arithmetic op per chain: an fma (two flops) for floats, otherwise a mul
/// (the slowest integer op, giving a lower bound).
#[cube]
fn step<I: Numeric, N: Size>(
    s: Vector<I, N>,
    b: Vector<I, N>,
    c: Vector<I, N>,
    #[comptime] use_fma: bool,
) -> Vector<I, N> {
    if use_fma { fma(s, b, c) } else { s * b }
}