cubek_convolution/components/
problem.rs

1use cubecl::{
2    ir::AddressType,
3    zspace::{Shape, Strides, shape},
4};
5use cubek_matmul::definition::{AccumulatorOperand, MatmulGlobalElems, MatmulProblem};
6use cubek_std::MatrixLayout;
7
8#[derive(Clone, Debug, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
9pub enum ConvolutionOperation {
10    Forward,
11    BackwardData,
12    BackwardWeight,
13    ForwardTransposed,
14}
15
16#[derive(Clone, Debug)]
17/// Description of a matmul problem to solve, regardless of actual data
18pub struct ConvolutionProblem {
19    pub m: usize,
20    pub n: usize,
21    pub k: usize,
22
23    pub lhs_strides: Strides,
24    pub rhs_strides: Strides,
25
26    pub lhs_layout: MatrixLayout,
27    pub rhs_layout: MatrixLayout,
28
29    pub kernel_size: Vec<u32>,
30    pub stride: Vec<u32>,
31    /// Padding at the beginning of each spatial dimension.
32    pub padding: Vec<i32>,
33    pub dilation: Vec<u32>,
34
35    pub batches: usize,
36    pub channels: usize,
37    pub out_channels: usize,
38    pub in_shape: Shape,
39    pub out_shape: Shape,
40
41    /// Channels after applying loader-specific padding
42    pub padded_channels: usize,
43    pub operation: ConvolutionOperation,
44
45    pub dimensionality: Dimensionality,
46
47    pub global_dtypes: MatmulGlobalElems,
48    /// Address type, defined as the max of each handle's `required_address_type`
49    pub address_type: AddressType,
50}
51
52impl ConvolutionProblem {
53    /// The bias is not part of the convolution problem: it arrives at launch, so
54    /// the caller states whether one is there. The matmul selector charges its
55    /// stage against the shared-memory budget.
56    pub fn as_matmul_problem(&self, accumulator: AccumulatorOperand) -> MatmulProblem {
57        let rank = self.lhs_strides.len();
58
59        // Strides are expected to be in row major (m, n) format so for matmul checks we need to
60        // convert them to that format, with all other dims treated as batch dims so they're still
61        // checked.
62        // lhs already has the right format, but rhs needs special handling.
63        // (h, w, c, n)
64        let lhs_strides = match self.lhs_layout {
65            MatrixLayout::RowMajor => self.lhs_strides.clone(),
66            MatrixLayout::ColMajor => {
67                let mut lhs_strides: Strides = self.lhs_strides[1..rank].into();
68                lhs_strides.push(self.lhs_strides[0]);
69                lhs_strides
70            }
71        };
72        let rhs_strides = match self.rhs_layout {
73            MatrixLayout::RowMajor => self.rhs_strides.clone(),
74            MatrixLayout::ColMajor => {
75                let mut rhs_strides: Strides = self.rhs_strides[1..rank].into();
76                rhs_strides.push(self.rhs_strides[0]);
77                rhs_strides
78            }
79        };
80
81        MatmulProblem {
82            m: self.m,
83            n: self.n,
84            k: self.k,
85            lhs_batches: shape![],
86            rhs_batches: shape![],
87            out_batches: shape![],
88            lhs_strides,
89            rhs_strides,
90            lhs_layout: self.lhs_layout,
91            rhs_layout: self.rhs_layout,
92            lhs_shape: shape![self.m, self.k],
93            rhs_shape: shape![self.k, self.n],
94            out_shape: shape![self.m, self.n],
95            out_strides: MatrixLayout::RowMajor.to_strides(&[self.m, self.n]),
96            out_layout: MatrixLayout::RowMajor,
97            lhs_scheme: None,
98            rhs_scheme: None,
99            global_dtypes: self.global_dtypes.clone(),
100            // Convolution retains its existing precision policy even when
101            // ordinary matmul disables automatic TF32 conversion.
102            allow_tf32: true,
103            address_type: self.address_type,
104            accumulator,
105        }
106    }
107
108    pub fn should_check_channel(&self) -> bool {
109        // Data gradients contract over the forward output channels. Comparing
110        // C_in misses the tail when it happens to equal padded C_out (8 -> 7,
111        // for example), allowing the im2col/weight readers to consume padding.
112        let channels = match self.operation {
113            ConvolutionOperation::BackwardData | ConvolutionOperation::ForwardTransposed => {
114                self.out_channels
115            }
116            _ => self.channels,
117        };
118        channels != self.padded_channels
119    }
120
121    pub fn should_check_spatial_bounds(&self) -> bool {
122        spatial_bounds_required(
123            self.operation,
124            &self.kernel_size,
125            &self.stride,
126            &self.padding,
127            &self.dilation,
128            &self.in_shape,
129            &self.out_shape,
130        )
131    }
132}
133
134fn spatial_bounds_required(
135    operation: ConvolutionOperation,
136    kernel_size: &[u32],
137    stride: &[u32],
138    padding: &[i32],
139    dilation: &[u32],
140    in_shape: &[usize],
141    out_shape: &[usize],
142) -> bool {
143    (0..kernel_size.len()).any(|dim| {
144        let kernel_extent = (kernel_size[dim] as i64 - 1) * dilation[dim] as i64;
145        let padding = padding[dim] as i64;
146
147        match operation {
148            ConvolutionOperation::Forward | ConvolutionOperation::BackwardWeight => {
149                let first = -padding;
150                let last =
151                    (out_shape[dim] as i64 - 1) * stride[dim] as i64 + kernel_extent - padding;
152                first < 0 || last >= in_shape[dim] as i64
153            }
154            ConvolutionOperation::ForwardTransposed | ConvolutionOperation::BackwardData => {
155                let first_numerator = padding - kernel_extent;
156                let last_numerator = in_shape[dim] as i64 - 1 + padding;
157                // A transposed convolution only has a source coordinate when the numerator is
158                // divisible by the stride. Non-unit strides therefore require the checked path
159                // even when every quotient falls within the source tensor.
160                stride[dim] != 1
161                    || first_numerator < 0
162                    || last_numerator >= out_shape[dim] as i64 * stride[dim] as i64
163            }
164        }
165    })
166}
167
168#[cfg(test)]
169mod tests {
170    use super::{ConvolutionOperation, spatial_bounds_required};
171
172    #[test]
173    fn forward_checks_bounds_for_end_only_padding() {
174        assert!(spatial_bounds_required(
175            ConvolutionOperation::Forward,
176            &[3],
177            &[1],
178            &[0],
179            &[1],
180            &[5],
181            &[5],
182        ));
183    }
184
185    #[test]
186    fn forward_skips_bounds_for_exact_unpadded_geometry() {
187        assert!(!spatial_bounds_required(
188            ConvolutionOperation::Forward,
189            &[3],
190            &[1],
191            &[0],
192            &[1],
193            &[5],
194            &[3],
195        ));
196    }
197
198    #[test]
199    fn backward_data_checks_kernel_overhang_without_begin_padding() {
200        assert!(spatial_bounds_required(
201            ConvolutionOperation::BackwardData,
202            &[3],
203            &[1],
204            &[0],
205            &[1],
206            &[5],
207            &[3],
208        ));
209    }
210
211    #[test]
212    fn backward_data_skips_bounds_for_pointwise_geometry() {
213        assert!(!spatial_bounds_required(
214            ConvolutionOperation::BackwardData,
215            &[1],
216            &[1],
217            &[0],
218            &[1],
219            &[5],
220            &[5],
221        ));
222    }
223
224    #[test]
225    fn backward_data_checks_stride_divisibility() {
226        assert!(spatial_bounds_required(
227            ConvolutionOperation::BackwardData,
228            &[1],
229            &[2],
230            &[0],
231            &[1],
232            &[5],
233            &[3],
234        ));
235    }
236}
237
238/// Spatial dimensionality of an operation
239#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug)]
240pub enum Dimensionality {
241    Dim1,
242    Dim2,
243    Dim3,
244}
245
246impl Dimensionality {
247    pub fn num_dims(&self) -> usize {
248        match self {
249            Dimensionality::Dim1 => 1,
250            Dimensionality::Dim2 => 2,
251            Dimensionality::Dim3 => 3,
252        }
253    }
254}
cubek_convolution/components/problem.rs

cubek_convolution/components/
problem.rs
