cubek_convolution/launch/
unit_strategy.rs

1use std::fmt::Display;
2
3use crate::routines::unit::UnitMatmul;
4use cubecl::client::Client;
5use cubek_matmul::{
6    definition::{MatmulElems, MatmulProblem},
7    multi_level::{
8        components::{stage::PartitionBuffering, tile::TileMatmulKind},
9        definition::BatchMatmulBlueprint,
10        routines::{
11            TileSizeSelection, UnitSmemBudget, batch::simple_unit::SimpleUnitSelectionArgs,
12            closest_factor_pair,
13        },
14        stage::{SwizzleMode, SwizzleModes},
15    },
16    routine::BlueprintStrategy,
17};
18
19/// Register-tiled F32 convolution candidates. Selection is left to the caller's tuner.
20///
21/// A forced candidate is a fixed [`UnitTiling`]. The presets on that type span
22/// the output-rows (`M`), output-width (`N`) and reduction (`K`) axes so that
23/// autotune has a candidate near every regime a convolution's matmul lands in,
24/// rather than only the two dense shapes the presets were first measured on.
25#[derive(Clone, Copy, Debug)]
26pub enum UnitStrategy {
27    /// Use the general matmul selector.
28    Inferred(TileSizeSelection),
29    /// A fixed register tiling, one of the [`UnitTiling`] presets.
30    Fixed(UnitTiling),
31}
32
33/// A forced register tiling, expressed as the tiles each unit's partition holds
34/// along `(M, N, K)`.
35///
36/// The tile `(4, 4, 4)` and stage `(16, 8, 1)` are held constant, so the stage
37/// a preset asks for is `(64*m) x (32*n) x (4*k)` output/reduction elements.
38/// A device whose cube holds fewer than 128 units gets the largest
39/// power-of-two stage it holds instead.
40/// A preset that overflows shared memory is trimmed by [`UnitSmemBudget`], the
41/// same budget the matmul unit selector fits its tilings to.
42#[derive(Clone, Copy, Debug, PartialEq, Eq)]
43pub struct UnitTiling {
44    partition: (u32, u32, u32),
45}
46
47impl UnitTiling {
48    /// `128 x 32 x 32`: few output channels, as in a 32-wide dense layer.
49    pub const NARROW: Self = Self::new(2, 1, 8);
50    /// `128 x 64 x 32`: a moderate output width.
51    pub const BALANCED: Self = Self::new(2, 2, 8);
52    /// `128 x 128 x 32`: wide outputs and the wide weight-gradient matrices a
53    /// large channel count produces (`taps * C_in`), e.g. 128–512 channels.
54    pub const WIDE: Self = Self::new(2, 4, 8);
55    /// `128 x 64 x 64`: a long reduction, as a large `C_in * taps` contraction
56    /// gives on a short feature map (the short-spatial/long-channel regime).
57    pub const DEEP: Self = Self::new(2, 2, 16);
58    /// `256 x 64 x 32`: many output rows, as a large spatial map times batch
59    /// gives in forward and data-gradient passes.
60    pub const TALL: Self = Self::new(4, 2, 8);
61
62    const fn new(m: u32, n: u32, k: u32) -> Self {
63        Self {
64            partition: (m, n, k),
65        }
66    }
67}
68
69impl Display for UnitStrategy {
70    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
71        match self {
72            Self::Inferred(tile_size) => write!(f, "unit_{tile_size}"),
73            Self::Fixed(UnitTiling {
74                partition: (m, n, k),
75            }) => write!(f, "unit_{m}x{n}x{k}"),
76        }
77    }
78}
79
80impl UnitStrategy {
81    pub(crate) fn blueprint(
82        self,
83        client: &Client,
84        problem: &MatmulProblem,
85        dtypes: &MatmulElems,
86    ) -> BlueprintStrategy<UnitMatmul> {
87        let partition = match self {
88            Self::Inferred(tile_size) => {
89                return BlueprintStrategy::Inferred(SimpleUnitSelectionArgs { tile_size });
90            }
91            Self::Fixed(tiling) => tiling.partition,
92        };
93        let hardware = &client.properties().hardware;
94        let plane_dim = match hardware.plane_size_min {
95            0 => 32,
96            plane_dim => plane_dim,
97        };
98        // One unit per stage partition. The cpu runtime runs a barrier kernel on
99        // one thread per unit, so a 128-unit stage there crawls.
100        let stage = closest_factor_pair(1 << hardware.max_units_per_cube.min(128).ilog2());
101        let tiling = UnitSmemBudget {
102            dtypes,
103            stage_buffering: 1,
104            accumulator: problem.accumulator,
105            max_smem: hardware.max_shared_memory_size,
106            plane_dim,
107        }
108        .fit((4, 4, 4), partition, stage);
109        let mut builder =
110            BatchMatmulBlueprint::builder(TileMatmulKind::Register, tiling, plane_dim, problem)
111                .partition_buffering(PartitionBuffering::Single);
112        if client.properties().features.alignment && tiling.elements_per_stage_along_k() >= 32 {
113            builder = builder.shared_swizzle(SwizzleModes {
114                lhs: SwizzleMode::B128,
115                rhs: SwizzleMode::B128,
116                ..Default::default()
117            });
118        }
119        BlueprintStrategy::Forced(builder.build())
120    }
121}
cubek_convolution/launch/unit_strategy.rs

cubek_convolution/launch/
unit_strategy.rs
