cubek_convolution/kernels/backward_weight/
launch.rs

1use crate::components::{ConvolutionProblem, Dimensionality};
2use crate::{
3    AcceleratedTileKind, ReadingStrategy, algorithm::Algorithm,
4    components::global::args::RuntimeArgs,
5};
6use crate::{
7    ConvolutionArgs, Strategy, backward_weight::args::ConcreteArgs,
8    components::ConvolutionOperation, kernels::algorithm::simple::*,
9};
10use crate::{
11    components::ConvSetupError, kernels::backward_weight::selector::launch_kernel_concrete,
12};
13use cubecl::{Runtime, client::ComputeClient, prelude::*, std::tensor::TensorHandle};
14use cubek_matmul::launch::{MatmulInputHandle, MatmulInputHandleRef};
15use cubek_matmul::{
16    components::tile::{cmma::CmmaMatmul, io::Strided, mma::MmaMatmul},
17    definition,
18};
19use cubek_matmul::{
20    definition::{AvailableLineSizes, MatmulElems, MatrixLayout},
21    routines::BlueprintStrategy,
22};
23use derive_new::new;
24
25macro_rules! with_tile_kind {
26    ($kind: expr, $T: ident, $launch: expr) => {
27        match $kind {
28            AcceleratedTileKind::Cmma => {
29                type $T = CmmaMatmul<Option<Strided>>;
30                ($launch)()
31            }
32            AcceleratedTileKind::Mma => {
33                type $T = MmaMatmul<Strided, Strided, Option<Strided>>;
34                ($launch)()
35            }
36        }
37    };
38}
39
40#[allow(clippy::result_large_err, clippy::too_many_arguments)]
41pub fn launch<R: Runtime, const N_SPATIAL: usize>(
42    strategy: &Strategy,
43    client: &ComputeClient<R>,
44    input: MatmulInputHandle<R>,
45    out_grad: MatmulInputHandle<R>,
46    weight_grad: TensorHandle<R>,
47    args: ConvolutionArgs<N_SPATIAL>,
48    dtypes: MatmulElems,
49) -> Result<(), ConvSetupError> {
50    launch_ref(
51        strategy,
52        client,
53        &input.as_ref(),
54        &out_grad.as_ref(),
55        &weight_grad.as_ref(),
56        args,
57        dtypes,
58    )
59}
60
61/// Perform an n-dimensional convolution using the implicit GEMM (im2col) algorithm, using cubecl
62/// tiling matmul components, using the specified algorithm.
63///
64/// * `input` - The input feature map, layout should be [batches, depth, height, width, in_channels]
65/// * `weight` - The weights (filter) applied to each kernel, layout should be [out_channels, kernel_d, kernel_h, kernel_w, in_channels]
66/// * `out` - The output feature map, layout should be [batches, out_depth, out_height, out_width, out_channels]
67/// * `bias` - The bias added to each out channel
68/// * `options` - The options to use for the convolution
69#[allow(clippy::result_large_err, clippy::too_many_arguments)]
70pub fn launch_ref<R: Runtime, const N_SPATIAL: usize>(
71    strategy: &Strategy,
72    client: &ComputeClient<R>,
73    input: &MatmulInputHandleRef<'_, R>,
74    out_grad: &MatmulInputHandleRef<'_, R>,
75    weight_grad: &TensorHandleRef<'_, R>,
76    args: ConvolutionArgs<N_SPATIAL>,
77    dtypes: MatmulElems,
78) -> Result<(), ConvSetupError> {
79    let backprop = BackwardsWeight::new(client, input, out_grad, weight_grad, args, dtypes);
80
81    match strategy {
82        Strategy::Simple {
83            read_strategy,
84            tile_kind,
85        } => with_tile_kind!(tile_kind, Accelerated, || match read_strategy {
86            ReadingStrategy::Cyclic => backprop.launch::<SimpleSyncCyclicConv<Accelerated>>(),
87            ReadingStrategy::Strided => backprop.launch::<SimpleSyncStridedConv<Accelerated>>(),
88            ReadingStrategy::Tilewise => backprop.launch::<SimpleSyncTilewiseConv<Accelerated>>(),
89            ReadingStrategy::AsyncCyclic => backprop.launch::<SimpleAsyncCyclicConv<Accelerated>>(),
90            ReadingStrategy::AsyncStrided =>
91                backprop.launch::<SimpleAsyncStridedConv<Accelerated>>(),
92            ReadingStrategy::Tma => backprop.launch::<SimpleAsyncTmaConv<Accelerated>>(),
93        }),
94    }
95}
96
97#[derive(new)]
98struct BackwardsWeight<'a, R: Runtime, const N_SPATIAL: usize> {
99    client: &'a ComputeClient<R>,
100    input: &'a MatmulInputHandleRef<'a, R>,
101    out_grad: &'a MatmulInputHandleRef<'a, R>,
102    weight_grad: &'a TensorHandleRef<'a, R>,
103    args: ConvolutionArgs<N_SPATIAL>,
104    dtypes: MatmulElems,
105}
106
107impl<'a, R: Runtime, const N_SPATIAL: usize> BackwardsWeight<'a, R, N_SPATIAL> {
108    fn launch<Alg: Algorithm>(self) -> Result<(), ConvSetupError>
109    where
110        Alg::Args: ConcreteArgs<Alg::Routine>,
111    {
112        let ConvolutionArgs {
113            stride,
114            padding,
115            dilation,
116        } = self.args;
117
118        let dimensionality = match N_SPATIAL {
119            1 => Dimensionality::Dim1,
120            2 => Dimensionality::Dim2,
121            3 => Dimensionality::Dim3,
122            other => unimplemented!("Unsupported dimensionality {other}"),
123        };
124
125        launch_with_algorithm::<R, Alg>(
126            self.client,
127            self.input,
128            self.out_grad,
129            self.weight_grad,
130            (&stride, &padding, &dilation),
131            dimensionality,
132            self.dtypes,
133        )
134    }
135}
136
137#[allow(clippy::too_many_arguments)]
138fn launch_with_algorithm<R: Runtime, Alg: Algorithm>(
139    client: &ComputeClient<R>,
140    input: &MatmulInputHandleRef<'_, R>,
141    out_grad: &MatmulInputHandleRef<'_, R>,
142    weight_grad: &TensorHandleRef<'_, R>,
143    (stride, padding, dilation): (&[usize], &[usize], &[usize]),
144    dimensionality: Dimensionality,
145    dtypes: MatmulElems,
146) -> Result<(), ConvSetupError>
147where
148    Alg::Args: ConcreteArgs<Alg::Routine>,
149{
150    let rank = input.data().shape.len();
151    let dim_c = rank - 1;
152
153    let n = input.shape()[0];
154    let c = input.shape()[dim_c];
155
156    let out_c = out_grad.shape()[dim_c];
157
158    let in_shape = &input.shape()[1..dim_c];
159    let kernel_shape = &weight_grad.shape[1..dim_c];
160    let out_shape = &out_grad.shape()[1..dim_c];
161
162    let op = ConvolutionOperation::BackwardWeight;
163
164    let input_data = Alg::into_tensor_handle(client, input.data(), dtypes.lhs_global, op)?;
165    let out_grad_data = Alg::into_tensor_handle(client, out_grad.data(), dtypes.rhs_global, op)?;
166
167    let mut input = *input;
168    let mut out_grad = *out_grad;
169
170    *input.data_mut() = input_data.as_ref();
171    *out_grad.data_mut() = out_grad_data.as_ref();
172
173    let address_type = input
174        .required_address_type()
175        .max(out_grad.required_address_type())
176        .max(weight_grad.required_address_type());
177
178    let problem = ConvolutionProblem {
179        m: out_c,
180        n: c * kernel_shape.iter().product::<usize>(),
181        k: n * out_shape.iter().product::<usize>(),
182        lhs_strides: input.data().strides.into(),
183        rhs_strides: out_grad.data().strides.into(),
184        lhs_layout: definition::MatrixLayout::ColMajor,
185        rhs_layout: definition::MatrixLayout::RowMajor,
186        kernel_size: kernel_shape.iter().map(|it| *it as u32).collect(),
187        stride: stride.iter().map(|it| *it as u32).collect(),
188        padding: padding.iter().map(|it| *it as i32).collect(),
189        dilation: dilation.iter().map(|it| *it as u32).collect(),
190
191        batches: n,
192        in_shape: in_shape.into(),
193        out_shape: out_shape.into(),
194        channels: c,
195        out_channels: out_c,
196
197        padded_channels: c,
198        operation: op,
199
200        dimensionality,
201        global_dtypes: dtypes.as_global_elems(),
202        address_type,
203    };
204
205    launch_kernel::<R, Alg>(
206        client,
207        &input,
208        &out_grad,
209        weight_grad,
210        problem,
211        &BlueprintStrategy::Inferred(Default::default()),
212        dtypes,
213    )
214}
215
216#[allow(clippy::result_large_err, clippy::too_many_arguments)]
217pub fn launch_kernel<R: Runtime, Alg: Algorithm>(
218    client: &ComputeClient<R>,
219    input: &MatmulInputHandleRef<'_, R>,
220    out_grad: &MatmulInputHandleRef<'_, R>,
221    weight_grad: &TensorHandleRef<'_, R>,
222    problem: ConvolutionProblem,
223    blueprint_strategy: &BlueprintStrategy<RuntimeArgs, Alg::Routine>,
224    dtypes: MatmulElems,
225) -> Result<(), ConvSetupError>
226where
227    Alg::Args: ConcreteArgs<Alg::Routine>,
228{
229    // Shape/strides are treated as k-major, with the last dim always being the contiguous one.
230    // So for the sake of selecting a line size, the shape/strides are always row-major.
231    let line_sizes = AvailableLineSizes::from_type_sizes(
232        client,
233        input.data().elem_size,
234        out_grad.data().elem_size,
235        weight_grad.elem_size,
236    )
237    .filter_lhs_with_tensor(
238        out_grad.data().strides,
239        out_grad.data().shape,
240        MatrixLayout::RowMajor,
241    )
242    .filter_rhs_with_tensor(
243        input.data().strides,
244        input.data().shape,
245        MatrixLayout::RowMajor,
246    )
247    .filter_out_with_tensor(weight_grad.strides, weight_grad.shape);
248
249    let line_sizes = Alg::filter_line_sizes(line_sizes).pick_max()?;
250
251    launch_kernel_concrete::<R, Alg::Args, Alg::Routine>(
252        client,
253        input,
254        out_grad,
255        weight_grad,
256        problem,
257        line_sizes,
258        blueprint_strategy,
259        &dtypes,
260    )
261}
cubek_convolution/kernels/backward_weight/launch.rs

cubek_convolution/kernels/backward_weight/
launch.rs
