burn_cubecl/kernel/conv/
base.rs

1use burn_backend::ops::ConvOptions;
2use burn_std::Shape;
3use cubek::convolution::{AcceleratedTileKind, components::ConvSetupError};
4
5#[cfg(feature = "autotune")]
6use crate::kernel::conv::{backward_weight::wgrad_autotune, dgrad_autotune};
7use crate::{
8    kernel::conv::{
9        backward_data::{fallback::conv_data_backward_fallback, implicit_gemm::*},
10        backward_weight::{fallback::conv_weight_backward_fallback, implicit_gemm::*},
11        forward::implicit_gemm::conv_gemm_simple_sync,
12    },
13    ops::{permute_nchw_to_nhwc, permute_nchw_to_nhwc_shape, permute_nhwc_to_nchw},
14    tensor::CubeTensor,
15};
16
17use super::conv_direct;
18#[cfg(feature = "autotune")]
19use super::forward::conv_autotune;
20
21/// The strategy to be used when launching a convolution kernel.
22pub enum ConvStrategy {
23    /// A simple direct convolution.
24    Direct,
25    #[cfg(feature = "autotune")]
26    /// Using autotune to choose the best kernel based on runtime information.
27    Autotune,
28    /// Implicit GEMM implementation of convolution. Lower memory usage but requires CMMA and
29    /// has constraints on tensor shape.
30    ImplicitGemm,
31}
32
33impl Default for ConvStrategy {
34    fn default() -> Self {
35        // if autotune is enabled, default to autotune
36        #[cfg(feature = "autotune")]
37        return ConvStrategy::Autotune;
38
39        // if autotune is disabled, default to the more memory-conservative algorithm
40        #[cfg(not(feature = "autotune"))]
41        ConvStrategy::Direct
42    }
43}
44
45/// Performs an N-dimensional convolution with the given strategy
46///
47/// * `input` - The input feature map
48/// * `weight` - The weights (filter) applied to each kernel
49/// * `bias` - The bias added to each channel
50/// * `options` - The options to use for the convolution
51/// * `strategy` - The convolution algorithm to use. Autotune will pick the fastest available option.
52pub fn conv_forward<const N: usize>(
53    input: CubeTensor,
54    weight: CubeTensor,
55    bias: Option<CubeTensor>,
56    options: ConvOptions<N>,
57    strategy: ConvStrategy,
58) -> Result<CubeTensor, ConvSetupError> {
59    let input = permute_nchw_to_nhwc(input);
60    let weight = permute_nchw_to_nhwc(weight);
61
62    let out = conv_forward_nhwc(input, weight, bias, options, strategy)?;
63
64    Ok(permute_nhwc_to_nchw(out))
65}
66
67/// Performs an N-dimensional convolution with the given strategy on NHWC inputs/outputs
68///
69/// * `input` - The input feature map
70/// * `weight` - The weights (filter) applied to each kernel
71/// * `bias` - The bias added to each channel
72/// * `options` - The options to use for the convolution
73/// * `strategy` - The convolution algorithm to use. Autotune will pick the fastest available option.
74pub fn conv_forward_nhwc<const N: usize>(
75    input: CubeTensor,
76    weight: CubeTensor,
77    bias: Option<CubeTensor>,
78    options: ConvOptions<N>,
79    strategy: ConvStrategy,
80) -> Result<CubeTensor, ConvSetupError> {
81    match strategy {
82        ConvStrategy::Direct => conv_direct::<N>(input, weight, bias, options),
83        #[cfg(feature = "autotune")]
84        ConvStrategy::Autotune => Ok(conv_autotune::<N>(input, weight, bias, options)),
85        ConvStrategy::ImplicitGemm => {
86            if options.groups != 1 {
87                conv_direct::<N>(input, weight, bias, options)
88            } else {
89                conv_gemm_simple_sync::<N>(input, weight, bias, options, AcceleratedTileKind::Cmma)
90            }
91        }
92    }
93}
94
95/// Performs an N-dimensional convolution backwards pass with regard to weight, with the given strategy
96///
97/// * `input` - The input feature map
98/// * `out_grad` - The output gradients
99/// * `weight_shape` - The shape of the weights/weight gradients
100/// * `options` - The options used for the convolution
101/// * `strategy` - The convolution algorithm to use. Autotune will pick the fastest available option.
102pub fn conv_weight_backward<const N: usize>(
103    input: CubeTensor,
104    out_grad: CubeTensor,
105    weight_shape: Shape,
106    options: ConvOptions<N>,
107    strategy: ConvStrategy,
108) -> Result<CubeTensor, ConvSetupError> {
109    let input = permute_nchw_to_nhwc(input);
110    let out_grad = permute_nchw_to_nhwc(out_grad);
111    let weight_shape = permute_nchw_to_nhwc_shape(weight_shape);
112
113    let weight_grad = match strategy {
114        ConvStrategy::Direct => {
115            conv_weight_backward_fallback::<N>(input, out_grad, weight_shape, options)
116        }
117        #[cfg(feature = "autotune")]
118        ConvStrategy::Autotune => Ok(wgrad_autotune::<N>(input, out_grad, weight_shape, options)),
119        ConvStrategy::ImplicitGemm => {
120            if options.groups != 1 {
121                conv_weight_backward_fallback::<N>(input, out_grad, weight_shape, options)
122            } else {
123                wgrad_gemm_simple_sync::<N>(
124                    input,
125                    out_grad,
126                    weight_shape,
127                    options,
128                    AcceleratedTileKind::Cmma,
129                )
130            }
131        }
132    }?;
133
134    Ok(permute_nhwc_to_nchw(weight_grad))
135}
136
137/// Performs an N-dimensional convolution backwards data pass with the given strategy
138///
139/// * `input` - The input feature map
140/// * `weight` - The weights (filter) applied to each kernel
141/// * `in_shape` - The shape of the input to the layer
142/// * `options` - The options to use for the convolution
143/// * `strategy` - The convolution algorithm to use. Autotune will pick the fastest available option.
144pub fn conv_data_backward<const N: usize>(
145    out_grad: CubeTensor,
146    weights: CubeTensor,
147    in_shape: Shape,
148    options: ConvOptions<N>,
149    strategy: ConvStrategy,
150) -> Result<CubeTensor, ConvSetupError> {
151    let out_grad = permute_nchw_to_nhwc(out_grad);
152    let weights = permute_nchw_to_nhwc(weights);
153    let in_shape = permute_nchw_to_nhwc_shape(in_shape);
154
155    let weight_grad = match strategy {
156        ConvStrategy::Direct => {
157            conv_data_backward_fallback::<N>(out_grad, weights, in_shape, options)?
158        }
159        #[cfg(feature = "autotune")]
160        ConvStrategy::Autotune => dgrad_autotune::<N>(out_grad, weights, in_shape, options),
161        ConvStrategy::ImplicitGemm => {
162            if options.groups != 1 || options.stride.iter().any(|&s| s != 1) {
163                conv_data_backward_fallback::<N>(out_grad, weights, in_shape, options)?
164            } else {
165                dgrad_gemm_simple_sync::<N>(
166                    out_grad,
167                    weights,
168                    in_shape,
169                    options,
170                    AcceleratedTileKind::Cmma,
171                )?
172            }
173        }
174    };
175
176    Ok(permute_nhwc_to_nchw(weight_grad))
177}
burn_cubecl/kernel/conv/base.rs

burn_cubecl/kernel/conv/
base.rs
