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

use cubecl::prelude::TensorBinding;
use cubek_std::InputBinding;

use crate::components::ConvolutionOperation;

/// Spatial convolution arguments (stride / beginning padding / dilation per spatial dim).
///
/// End padding is represented by the spatial extent of the output binding.
#[derive(Clone, Debug)]
pub struct ConvolutionArgs<const N_SPATIAL: usize> {
    pub stride: [usize; N_SPATIAL],
    /// Padding at the beginning of each spatial dimension.
    pub padding: [usize; N_SPATIAL],
    pub dilation: [usize; N_SPATIAL],
}

#[allow(clippy::large_enum_variant)]
/// Per-operation tensor bindings supplied to `launch_ref`.
///
/// Each variant carries exactly the bindings the corresponding operation needs.
/// The discriminant maps 1:1 to `ConvolutionOperation`.
pub enum ConvolutionInputs {
    Forward {
        input: InputBinding,
        weight: InputBinding,
        bias: Option<InputBinding>,
        out: TensorBinding,
    },
    BackwardData {
        out_grad: InputBinding,
        weights: InputBinding,
        in_grad: TensorBinding,
    },
    BackwardWeight {
        input: InputBinding,
        out_grad: InputBinding,
        weight_grad: TensorBinding,
    },
}

impl ConvolutionInputs {
    pub fn operation(&self) -> ConvolutionOperation {
        match self {
            ConvolutionInputs::Forward { .. } => ConvolutionOperation::Forward,
            ConvolutionInputs::BackwardData { .. } => ConvolutionOperation::BackwardData,
            ConvolutionInputs::BackwardWeight { .. } => ConvolutionOperation::BackwardWeight,
        }
    }
}