cubek_convolution/kernels/backward_data/
launch.rs

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

cubek_convolution/kernels/backward_data/
launch.rs
