cubek-convolution 0.3.0

CubeK: Convolution Kernels
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90

use core::fmt::Debug;
use cubecl::server::LaunchError;
use cubek_matmul::definition::{MatmulAvailabilityError, MatmulSetupError};
use cubek_std::InvalidConfigError;
use cubek_tile::launch::Refusal;

#[allow(clippy::large_enum_variant)]
pub enum ConvSetupError {
    Matmul(MatmulSetupError),
    Groups(usize),
    /// The depthwise routine was handed a convolution that is not depthwise: it requires
    /// `groups == in_channels == out_channels`, one filter per channel.
    NotDepthwise {
        groups: usize,
        input_channels: usize,
        output_channels: usize,
        weight_channels: usize,
        weight_group_channels: usize,
    },
    /// A caller-provided convolution strategy cannot form valid launch geometry.
    InvalidConfig(InvalidConfigError),
    Unknown,
    Launch(LaunchError),
}

impl From<LaunchError> for ConvSetupError {
    fn from(value: LaunchError) -> Self {
        Self::Launch(value)
    }
}

impl Debug for ConvSetupError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ConvSetupError::Matmul(err) => {
                write!(f, "{err:?}")
            }
            ConvSetupError::Groups(groups) => {
                writeln!(
                    f,
                    "Unable to launch matmul because groups must be one, is actually {groups}",
                )
            }
            ConvSetupError::NotDepthwise {
                groups,
                input_channels,
                output_channels,
                weight_channels,
                weight_group_channels,
            } => writeln!(
                f,
                "Unable to launch the depthwise convolution because it needs one filter per \
                 channel, but groups is {groups}, input channels is {input_channels}, output \
                 channels is {output_channels}, weight channels is {weight_channels}, and weight \
                 channels per group is {weight_group_channels}",
            ),
            ConvSetupError::InvalidConfig(err) => {
                write!(f, "Invalid convolution config: {err}")
            }
            ConvSetupError::Unknown => write!(f, "Unknown"),
            ConvSetupError::Launch(err) => write!(f, "Launch error {err:?}"),
        }
    }
}

impl From<MatmulSetupError> for ConvSetupError {
    fn from(value: MatmulSetupError) -> Self {
        Self::Matmul(value)
    }
}

/// A launch the tile engine refused is a configuration this routine cannot run.
impl From<Refusal> for ConvSetupError {
    fn from(refusal: Refusal) -> Self {
        Self::InvalidConfig(Box::new(refusal))
    }
}

impl From<MatmulAvailabilityError> for ConvSetupError {
    fn from(value: MatmulAvailabilityError) -> Self {
        Self::Matmul(MatmulSetupError::Unavailable(value))
    }
}

#[allow(clippy::from_over_into)]
impl Into<String> for ConvSetupError {
    fn into(self) -> String {
        format!("{self:?}")
    }
}