cubek_convolution/launch/
strategy.rs

1use std::fmt::Display;
2
3use serde::{Deserialize, Serialize};
4
5use crate::{definition::ConvBlueprint, launch::UnitStrategy};
6
7/// Top-level user-facing strategy for `launch_ref`.
8///
9/// `Specific` selects an algorithm and tile-matmul kind, letting the routine
10/// infer the rest. `Forced` bypasses inference and uses the supplied blueprint
11/// directly (the algorithm tag must still be provided so the kernel-side
12/// generic dispatch can pick the right reading-strategy implementation).
13#[derive(Clone, Debug)]
14pub enum Strategy {
15    /// User picks the algorithm and tile-matmul kind. Tiling/swizzle/etc. are
16    /// inferred from the problem.
17    Inferred {
18        algorithm: ConvAlgorithm,
19        tile_kind: AcceleratedTileKind,
20    },
21    /// User supplies a pre-built blueprint. The algorithm tag tells the launcher
22    /// which kernel generic to instantiate; the tiling/swizzle/etc. come from
23    /// the blueprint. The tile-matmul kind comes from the blueprint as well.
24    Forced {
25        algorithm: ConvAlgorithm,
26        blueprint: ConvBlueprint,
27    },
28    /// F32 register-tiled implicit convolution, for devices without a usable
29    /// accelerated matmul. It takes its own strategy because it has no tile-matmul
30    /// kind to pick, and its forced tilings span regimes autotune chooses from.
31    Unit(UnitStrategy),
32}
33
34/// The convolution-side algorithm enum. Subsumes the previous
35/// `ReadingStrategy` axis and the Simple/Specialized split. A single value
36/// here picks one concrete `Routine` impl (see `crate::routines`).
37#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
38pub enum ConvAlgorithm {
39    SimpleSyncCyclic,
40    SimpleSyncStrided,
41    SimpleSyncTilewise,
42    SimpleAsyncCyclic,
43    SimpleAsyncStrided,
44    SimpleAsyncTma,
45    SpecializedAsyncCyclic,
46    SpecializedAsyncStrided,
47    SpecializedTma,
48}
49
50impl Display for ConvAlgorithm {
51    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52        let s = match self {
53            ConvAlgorithm::SimpleSyncCyclic => "simple_sync_cyclic",
54            ConvAlgorithm::SimpleSyncStrided => "simple_sync_strided",
55            ConvAlgorithm::SimpleSyncTilewise => "simple_sync_tilewise",
56            ConvAlgorithm::SimpleAsyncCyclic => "simple_async_cyclic",
57            ConvAlgorithm::SimpleAsyncStrided => "simple_async_strided",
58            ConvAlgorithm::SimpleAsyncTma => "simple_async_tma",
59            ConvAlgorithm::SpecializedAsyncCyclic => "specialized_async_cyclic",
60            ConvAlgorithm::SpecializedAsyncStrided => "specialized_async_strided",
61            ConvAlgorithm::SpecializedTma => "specialized_tma",
62        };
63        f.write_str(s)
64    }
65}
66
67#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
68/// Which tile matmul to use for accelerated algorithms.
69pub enum AcceleratedTileKind {
70    #[default]
71    Cmma,
72    Mma,
73}
74
75impl Display for AcceleratedTileKind {
76    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77        match self {
78            AcceleratedTileKind::Cmma => f.write_str("cmma"),
79            AcceleratedTileKind::Mma => f.write_str("mma"),
80        }
81    }
82}
83
84impl Display for Strategy {
85    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
86        match self {
87            Strategy::Inferred {
88                algorithm,
89                tile_kind,
90            } => write!(f, "{algorithm}_{tile_kind}"),
91            Strategy::Forced {
92                algorithm,
93                blueprint: _,
94            } => write!(f, "{algorithm}_forced"),
95            Strategy::Unit(strategy) => write!(f, "{strategy}"),
96        }
97    }
98}
cubek_convolution/launch/strategy.rs

cubek_convolution/launch/
strategy.rs
