tenferro-gpu 0.3.0

CubeCL-backed CUDA and WebGPU provider backends for tenferro tensors.
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

// Portions of this file are adapted from cubek-reduce:
// https://github.com/tracel-ai/cubek/tree/9cf90b797107d46829e1c9d9355ce801c3dd4a7d/crates/cubek-reduce
//
// Original copyright:
// Copyright (c) 2022 Nathaniel Simard & CubeCL Framework Contributors
//
// Original source paths:
// - crates/cubek-reduce/src/launch/base.rs
// - crates/cubek-reduce/src/launch/strategy.rs
// - crates/cubek-reduce/src/routines/unit.rs
// - crates/cubek-reduce/src/routines/blueprint.rs
//
// Original license: MIT OR Apache-2.0.
// See tenferro-gpu/THIRD_PARTY_NOTICES.md for license notice text.
// Tenferro changes: narrowed to tenferro reduction ops, current CubeCL fork,
// single-axis keepdims output, and explicit tenferro column-major bindings.

//! Shared reduction launch routines.

// INVARIANT: CubeCL's kernel DSL lowers modulo checks differently from host
// Rust; the reduction dimensions were validated before launch.
#![allow(clippy::manual_is_multiple_of)]

use cubecl::prelude::*;

use crate::kernels::{CubeclKernelError, Result};

/// Validated single-axis reduction dimensions.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct ReduceProblem {
    /// Length of the axis being reduced.
    pub reduce_len: usize,
    /// Number of keepdims output elements.
    pub reduce_count: usize,
    /// Axis being reduced.
    pub axis: usize,
}

/// Unit-kernel launch blueprint.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct UnitReduceBlueprint {
    /// Whether some units in the last cube do not map to an output element.
    pub idle_units: bool,
}

/// Concrete CubeCL launch settings for a reduction kernel.
#[derive(Clone, Debug)]
pub struct ReduceLaunchSettings {
    /// Number of cubes to dispatch.
    pub cube_count: CubeCount,
    /// Number of units per cube.
    pub cube_dim: CubeDim,
    /// Unit-kernel metadata.
    pub blueprint: UnitReduceBlueprint,
}

/// Choose launch settings for the first unit-kernel implementation.
pub(crate) fn unit_launch_settings<R: Runtime>(
    client: &ComputeClient<R>,
    problem: ReduceProblem,
) -> Result<ReduceLaunchSettings> {
    unit_launch_settings_for_plane_size(client.properties().hardware.plane_size_max, problem)
}

fn unit_launch_settings_for_plane_size(
    plane_size: u32,
    problem: ReduceProblem,
) -> Result<ReduceLaunchSettings> {
    let cube_dim = CubeDim::new_1d(plane_size);
    let units_per_cube = cube_dim.num_elems() as usize;
    let cubes =
        u32::try_from(problem.reduce_count.div_ceil(units_per_cube).max(1)).map_err(|_| {
            CubeclKernelError::InvalidStrategy {
                reason: format!(
                    "reduction output element count {} exceeds static CubeCL launch limit",
                    problem.reduce_count
                ),
            }
        })?;

    Ok(ReduceLaunchSettings {
        cube_count: CubeCount::Static(cubes, 1, 1),
        cube_dim,
        blueprint: UnitReduceBlueprint {
            idle_units: problem.reduce_count % units_per_cube != 0,
        },
    })
}

#[cfg(test)]
mod tests;