cubek_reduce/launch/
strategy.rs

1use crate::routines::{
2    BlueprintStrategy, cube::CubeRoutine, plane::PlaneRoutine, unit::UnitRoutine,
3};
4use cubecl::{features::Plane, prelude::*};
5
6#[derive(Debug, Clone)]
7pub struct ReduceStrategy {
8    pub routine: RoutineStrategy,
9    pub line_size: LineSizeStrategy,
10}
11
12#[derive(Debug, Clone)]
13pub enum RoutineStrategy {
14    /// A unit is responsible to reduce a full vector.
15    Unit(BlueprintStrategy<UnitRoutine>),
16    /// A plane is responsible to reduce a full vector.
17    Plane(BlueprintStrategy<PlaneRoutine>),
18    /// A cube is responsible to reduce a full vector.
19    Cube(BlueprintStrategy<CubeRoutine>),
20}
21
22#[derive(Debug, Clone, Copy)]
23pub struct LineSizeStrategy {
24    /// When the vectorization is parallel, enable vectorization of the output so that each
25    /// unit can perform N reductions, where N is the output `line_size`.
26    pub parallel_output_vectorization: bool,
27}
28
29pub(crate) fn support_plane<R: Runtime>(client: &ComputeClient<R>) -> bool {
30    client.properties().features.plane.contains(Plane::Ops)
31}