cubecl_matmul/components/
error.rs

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