cubecl_matmul/components/
error.rs

1use cubecl_core::{CubeCount, CubeDim, LineSizeError, ir::StorageType, server::LaunchError};
2use std::fmt::{Debug, Display};
3
4use crate::components::{MatrixLayout, TileSize};
5
6/// Errors that can occur during the setup phase of a matmul operation.
7pub enum MatmulSetupError {
8    /// A required hardware or runtime feature is not available.
9    Unavailable(MatmulAvailabilityError),
10
11    /// The provided configuration is invalid or rejected by a component.
12    InvalidConfig(InvalidConfigError),
13
14    /// No compatible line size could be found for the given constraints.
15    LineSize(LineSizeError),
16
17    /// An error happened during launch.
18    Launch(LaunchError),
19}
20
21impl From<LaunchError> for MatmulSetupError {
22    fn from(value: LaunchError) -> Self {
23        Self::Launch(value)
24    }
25}
26
27/// A specific feature required for matmul is not available in the current runtime or hardware.
28pub enum MatmulAvailabilityError {
29    /// The requested cube count exceeds what the runtime or hardware supports.
30    CubeCountTooBig(CubeCount),
31
32    /// The requested cube dimensions are too large for the current runtime or hardware.
33    CubeDimTooBig(CubeDim),
34
35    /// The requested plane dimension is not supported.
36    PlaneDimUnsupported { plane_dim: u32 },
37
38    /// The required data types for input or output are not supported.
39    TypesUnavailable {
40        lhs: StorageType,
41        rhs: StorageType,
42        output: StorageType,
43    },
44
45    /// The required CMMA instruction is not supported for the given element types and tile size.
46    CmmaInstructionUnavailable {
47        lhs: StorageType,
48        rhs: StorageType,
49        output: StorageType,
50        size: Option<TileSize>,
51    },
52
53    /// Impossible to find a supported tile size for the problem.
54    TileSizeNotFound,
55
56    /// The layout of the matmul is unsupported
57    LayoutUnsupported {
58        lhs: MatrixLayout,
59        rhs: MatrixLayout,
60    },
61
62    /// Barrier synchronization is not available in the runtime.
63    BarrierUnavailable,
64
65    /// TMA (Tensor Memory Access) is not available in the runtime.
66    TmaUnavailable,
67
68    /// Dynamic selection of line size is unsupported in the current runtime.
69    DynamicLineSizeUnavailable,
70
71    /// Plane operations like plane_sum are unavailable
72    PlaneOpsUnavailable,
73}
74impl From<MatmulAvailabilityError> for MatmulSetupError {
75    fn from(value: MatmulAvailabilityError) -> Self {
76        Self::Unavailable(value)
77    }
78}
79
80impl From<InvalidConfigError> for MatmulSetupError {
81    fn from(value: InvalidConfigError) -> Self {
82        Self::InvalidConfig(value)
83    }
84}
85
86impl From<LineSizeError> for MatmulSetupError {
87    fn from(value: LineSizeError) -> Self {
88        Self::LineSize(value)
89    }
90}
91
92impl Display for MatmulSetupError {
93    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
94        write!(f, "{self:?}")
95    }
96}
97
98impl Debug for MatmulSetupError {
99    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
100        match self {
101            MatmulSetupError::Unavailable(err) => {
102                writeln!(
103                    f,
104                    "Unable to launch matmul because a required feature is unavailable: {err:?}"
105                )
106            }
107            MatmulSetupError::InvalidConfig(err) => {
108                writeln!(
109                    f,
110                    "Unable to launch matmul because the config is invalid: {:?}",
111                    err.to_string()
112                )
113            }
114            MatmulSetupError::LineSize(err) => {
115                writeln!(
116                    f,
117                    "Unable to launch matmul because could not find supported line size: {err:?}"
118                )
119            }
120            MatmulSetupError::Launch(err) => {
121                writeln!(f, "Unable to launch matmul with err: {err:?}")
122            }
123        }
124    }
125}
126
127impl Debug for MatmulAvailabilityError {
128    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
129        match self {
130            MatmulAvailabilityError::CubeCountTooBig(count) => {
131                writeln!(f, "Cube count too big {count:?}")
132            }
133            MatmulAvailabilityError::CubeDimTooBig(dim) => {
134                writeln!(f, "Cube dim too big {dim:?}")
135            }
136            MatmulAvailabilityError::PlaneDimUnsupported { plane_dim } => {
137                writeln!(
138                    f,
139                    "Plane dimension unsupported: {plane_dim}. Only 32 & 64 are supported."
140                )
141            }
142            MatmulAvailabilityError::TypesUnavailable { lhs, rhs, output } => {
143                writeln!(
144                    f,
145                    "Types lhs={lhs:?}, rhs={rhs:?} and/or output={output:?} not supported.",
146                )
147            }
148            MatmulAvailabilityError::CmmaInstructionUnavailable {
149                lhs,
150                rhs,
151                output,
152                size: Some(size),
153            } => writeln!(
154                f,
155                "Cmma on lhs {:?} rhs {:?} and output {:?} with shape m={:?}, n={:?}, k={:?} not supported.",
156                lhs,
157                rhs,
158                output,
159                size.m(),
160                size.n(),
161                size.k()
162            ),
163            MatmulAvailabilityError::LayoutUnsupported { lhs, rhs } => {
164                writeln!(
165                    f,
166                    "Cmma with layouts lhs {lhs:?} and rhs {rhs:?} not supported."
167                )
168            }
169            MatmulAvailabilityError::CmmaInstructionUnavailable {
170                lhs,
171                rhs,
172                output,
173                size: None,
174            } => writeln!(
175                f,
176                "Cmma on inputs lhs {lhs:?} rhs {rhs:?} and output {output:?} not supported.",
177            ),
178            MatmulAvailabilityError::BarrierUnavailable => {
179                writeln!(f, "Barrier is not available.")
180            }
181            MatmulAvailabilityError::TmaUnavailable => {
182                writeln!(f, "TMA is not available.")
183            }
184            MatmulAvailabilityError::DynamicLineSizeUnavailable => {
185                writeln!(f, "Dynamic line size is not available.")
186            }
187            MatmulAvailabilityError::PlaneOpsUnavailable => {
188                writeln!(f, "Plane-wide operations like plane_sum are not available.")
189            }
190            MatmulAvailabilityError::TileSizeNotFound => {
191                writeln!(f, "No tile size is available for the problem.")
192            }
193        }
194    }
195}
196
197/// Error that arises from invalid configurations
198pub type InvalidConfigError = Box<dyn Display>;
199
200/// Error that arises from invalid configurations
201pub struct FormattedConfigError {
202    func: Box<dyn Fn() -> String>,
203}
204
205impl FormattedConfigError {
206    #[allow(clippy::new_ret_no_self)]
207    pub fn new<F: Fn() -> String + 'static>(func: F) -> Box<dyn Display> {
208        Box::new(Self {
209            func: Box::new(func),
210        })
211    }
212}
213
214impl Display for FormattedConfigError {
215    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
216        let string = (self.func)();
217        write!(f, "{string}")
218    }
219}