cubek_convolution/kernels/backward_weight/
launch.rs

1use crate::{
2    backward_weight::args::ConcreteArgs,
3    components::{
4        ConvSetupError, ConvolutionOperation, ConvolutionProblem, Dimensionality,
5        global::args::RuntimeArgs,
6    },
7    kernels::backward_weight::selector::launch_kernel_concrete,
8    launch::ConvolutionArgs,
9    routines::Routine,
10};
11use cubecl::{client::Client, prelude::*};
12use cubek_matmul::{
13    definition::{AvailableVectorSizes, MatmulElems},
14    routine::BlueprintStrategy,
15};
16use cubek_std::{InputBinding, MatrixLayout};
17
18/// Backward-weight dispatch helper.
19///
20/// Called by `cubek_convolution::launch_ref` after the routine and
21/// blueprint-strategy have been resolved.
22#[allow(clippy::result_large_err, clippy::too_many_arguments)]
23pub(crate) fn launch_internal<const N_SPATIAL: usize, Rt: Routine>(
24    client: &Client,
25    input: InputBinding,
26    out_grad: InputBinding,
27    weight_grad: TensorBinding,
28    args: ConvolutionArgs<N_SPATIAL>,
29    blueprint_strategy: &BlueprintStrategy<RuntimeArgs, Rt::MatmulRoutine>,
30    dtypes: MatmulElems,
31) -> Result<(), ConvSetupError>
32where
33    Rt::Args: ConcreteArgs<Rt::MatmulRoutine>,
34{
35    let ConvolutionArgs {
36        stride,
37        padding,
38        dilation,
39    } = args;
40
41    let dimensionality = match N_SPATIAL {
42        1 => Dimensionality::Dim1,
43        2 => Dimensionality::Dim2,
44        3 => Dimensionality::Dim3,
45        other => unimplemented!("Unsupported dimensionality {other}"),
46    };
47
48    launch_with_routine::<Rt>(
49        client,
50        input,
51        out_grad,
52        weight_grad,
53        (&stride, &padding, &dilation),
54        dimensionality,
55        blueprint_strategy,
56        dtypes,
57    )
58}
59
60#[allow(clippy::too_many_arguments)]
61fn launch_with_routine<Rt: Routine>(
62    client: &Client,
63    input: InputBinding,
64    out_grad: InputBinding,
65    weight_grad: TensorBinding,
66    (stride, padding, dilation): (&[usize], &[usize], &[usize]),
67    dimensionality: Dimensionality,
68    blueprint_strategy: &BlueprintStrategy<RuntimeArgs, Rt::MatmulRoutine>,
69    dtypes: MatmulElems,
70) -> Result<(), ConvSetupError>
71where
72    Rt::Args: ConcreteArgs<Rt::MatmulRoutine>,
73{
74    let rank = input.data().shape.len();
75    let dim_c = rank - 1;
76
77    let n = input.shape()[0];
78    let c = input.shape()[dim_c];
79
80    let out_c = out_grad.shape()[dim_c];
81
82    let in_shape = &input.shape()[1..dim_c];
83    let kernel_shape = &weight_grad.shape[1..dim_c];
84    let out_shape = &out_grad.shape()[1..dim_c];
85
86    let op = ConvolutionOperation::BackwardWeight;
87
88    let input_data = Rt::correct_layout(client, input.clone().into_data(), dtypes.lhs_global, op)?;
89    let out_grad_data =
90        Rt::correct_layout(client, out_grad.clone().into_data(), dtypes.rhs_global, op)?;
91
92    let mut input = input.clone();
93    let mut out_grad = out_grad.clone();
94
95    *input.data_mut() = input_data;
96    *out_grad.data_mut() = out_grad_data;
97
98    let address_type = input
99        .required_address_type()
100        .max(out_grad.required_address_type())
101        .max(weight_grad.required_address_type(dtypes.acc_global.size()));
102
103    let problem = ConvolutionProblem {
104        m: out_c,
105        n: c * kernel_shape.iter().product::<usize>(),
106        k: n * out_shape.iter().product::<usize>(),
107        lhs_strides: input.data().strides.clone(),
108        rhs_strides: out_grad.data().strides.clone(),
109        lhs_layout: MatrixLayout::ColMajor,
110        rhs_layout: MatrixLayout::RowMajor,
111        kernel_size: kernel_shape.iter().map(|it| *it as u32).collect(),
112        stride: stride.iter().map(|it| *it as u32).collect(),
113        padding: padding.iter().map(|it| *it as i32).collect(),
114        dilation: dilation.iter().map(|it| *it as u32).collect(),
115
116        batches: n,
117        in_shape: in_shape.into(),
118        out_shape: out_shape.into(),
119        channels: c,
120        out_channels: out_c,
121
122        padded_channels: c,
123        operation: op,
124
125        dimensionality,
126        global_dtypes: dtypes.as_global_elems(),
127        address_type,
128    };
129
130    launch_kernel::<Rt>(
131        client,
132        input,
133        out_grad,
134        weight_grad,
135        problem,
136        blueprint_strategy,
137        dtypes,
138    )
139}
140
141#[allow(clippy::result_large_err, clippy::too_many_arguments)]
142pub fn launch_kernel<Rt: Routine>(
143    client: &Client,
144    input: InputBinding,
145    out_grad: InputBinding,
146    weight_grad: TensorBinding,
147    problem: ConvolutionProblem,
148    blueprint_strategy: &BlueprintStrategy<RuntimeArgs, Rt::MatmulRoutine>,
149    dtypes: MatmulElems,
150) -> Result<(), ConvSetupError>
151where
152    Rt::Args: ConcreteArgs<Rt::MatmulRoutine>,
153{
154    // Shape/strides are treated as k-major, with the last dim always being the contiguous one.
155    // So for the sake of selecting a vector size, the shape/strides are always row-major.
156    let vector_sizes = AvailableVectorSizes::from_type_sizes(
157        client,
158        input.data_elem_size(),
159        out_grad.data_elem_size(),
160        dtypes.acc_global.size(),
161    )
162    .filter_lhs_with_tensor(
163        &out_grad.data().strides,
164        &out_grad.data().shape,
165        MatrixLayout::RowMajor,
166    )
167    .filter_rhs_with_tensor(
168        &input.data().strides,
169        &input.data().shape,
170        MatrixLayout::RowMajor,
171    )
172    .filter_out_with_tensor(&weight_grad.strides, &weight_grad.shape);
173
174    let vector_sizes = Rt::filter_vector_sizes(vector_sizes).pick_max()?;
175
176    launch_kernel_concrete::<Rt::Args, Rt::MatmulRoutine>(
177        client,
178        input,
179        out_grad,
180        weight_grad,
181        problem,
182        vector_sizes,
183        blueprint_strategy,
184        &dtypes,
185    )
186}
cubek_convolution/kernels/backward_weight/launch.rs

cubek_convolution/kernels/backward_weight/
launch.rs
