cubecl_attention/components/
error.rs

1use cubecl_core::{CubeCount, CubeDim, LineSizeError};
2use cubecl_matmul::components::MatmulSetupError;
3use std::fmt::{Debug, Display};
4
5/// Errors that can occur during the setup phase of an attention operation.
6pub enum AttentionSetupError {
7    /// A required hardware or runtime feature is not available.
8    Unavailable(AttentionAvailabilityError),
9
10    /// The provided configuration is invalid or rejected by a component.
11    InvalidConfig(InvalidConfigError),
12
13    /// No compatible line size could be found for the given constraints.
14    LineSize(LineSizeError),
15
16    /// Error in underlying matmul
17    MatmulSetup(MatmulSetupError),
18}
19
20/// A specific feature required for attention is not available in the current runtime or hardware.
21pub enum AttentionAvailabilityError {
22    /// The requested cube count exceeds what the runtime or hardware supports.
23    CubeCountTooBig(CubeCount),
24
25    /// The requested cube dimensions are too large for the current runtime or hardware.
26    CubeDimTooBig(CubeDim),
27}
28
29impl From<AttentionAvailabilityError> for AttentionSetupError {
30    fn from(value: AttentionAvailabilityError) -> Self {
31        Self::Unavailable(value)
32    }
33}
34
35impl From<InvalidConfigError> for AttentionSetupError {
36    fn from(value: InvalidConfigError) -> Self {
37        Self::InvalidConfig(value)
38    }
39}
40
41impl From<LineSizeError> for AttentionSetupError {
42    fn from(value: LineSizeError) -> Self {
43        Self::LineSize(value)
44    }
45}
46
47impl Display for AttentionSetupError {
48    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
49        write!(f, "{self:?}")
50    }
51}
52
53impl Debug for AttentionSetupError {
54    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
55        match self {
56            AttentionSetupError::Unavailable(err) => {
57                writeln!(
58                    f,
59                    "Unable to launch attention because a required feature is unavailable: {err:?}"
60                )
61            }
62            AttentionSetupError::InvalidConfig(err) => {
63                writeln!(
64                    f,
65                    "Unable to launch attention because the config is invalid: {:?}",
66                    err.to_string()
67                )
68            }
69            AttentionSetupError::LineSize(err) => {
70                writeln!(
71                    f,
72                    "Unable to launch attention because could not find supported line size: {err:?}"
73                )
74            }
75            AttentionSetupError::MatmulSetup(matmul_setup_error) => {
76                writeln!(f, "{matmul_setup_error:?}")
77            }
78        }
79    }
80}
81
82impl Debug for AttentionAvailabilityError {
83    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
84        match self {
85            AttentionAvailabilityError::CubeCountTooBig(count) => {
86                writeln!(f, "Cube count too big {count:?}")
87            }
88            AttentionAvailabilityError::CubeDimTooBig(dim) => {
89                writeln!(f, "Cube dim too big {dim:?}")
90            }
91        }
92    }
93}
94
95/// Error that arises from invalid configurations
96pub type InvalidConfigError = Box<dyn Display>;
97
98/// Error that arises from invalid configurations
99pub struct FormattedConfigError {
100    func: Box<dyn Fn() -> String>,
101}
102
103impl FormattedConfigError {
104    #[allow(clippy::new_ret_no_self)]
105    pub fn new<F: Fn() -> String + 'static>(func: F) -> Box<dyn Display> {
106        Box::new(Self {
107            func: Box::new(func),
108        })
109    }
110}
111
112impl Display for FormattedConfigError {
113    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
114        let string = (self.func)();
115        write!(f, "{string}")
116    }
117}
118
119impl From<MatmulSetupError> for AttentionSetupError {
120    fn from(value: MatmulSetupError) -> Self {
121        Self::MatmulSetup(value)
122    }
123}
cubecl_attention/components/error.rs

cubecl_attention/components/
error.rs
