cubek-fft 0.3.0

CubeK: Fast Fourier Transform
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

//! Per-device limits used when launching the FFT kernels.

use cubecl::prelude::*;

/// Largest power-of-two `n_fft` such that a shared-memory radix-2 butterfly
/// over `dtype` fits in this device's per-cube shared memory.
///
/// Every shared-memory FFT kernel in this crate allocates two shared buffers of
/// `n_fft` elements (one for the real part, one for the imaginary part), so the
/// byte budget is `2 * dtype.size() * n_fft <= hardware.max_shared_memory_size`.
/// We floor to a power of two because the butterfly requires it.
pub(crate) fn max_shared_fft_n(client: &Client, dtype: ElemType) -> usize {
    let max_smem = client.properties().hardware.max_shared_memory_size;
    let max_elems = max_smem / (2 * dtype.size());
    floor_power_of_two(max_elems)
}

/// Hardware-reported maximum number of units (threads) per cube.
pub(crate) fn max_units_per_cube(client: &Client) -> usize {
    client.properties().hardware.max_units_per_cube as usize
}

fn floor_power_of_two(n: usize) -> usize {
    assert!(n > 0, "device reports zero shared memory / units");
    if n.is_power_of_two() {
        n
    } else {
        n.next_power_of_two() >> 1
    }
}