cubek_convolution/kernels/backward_data/selector/
select_kernel.rs

1use crate::{
2    backward_data::args::{ConcreteArgs, ConcreteInputsFactory, ConcreteOutputFactory},
3    components::global::args::RuntimeArgs,
4};
5use cubecl::{
6    prelude::TensorBinding,
7    {Runtime, client::ComputeClient},
8};
9use cubek_matmul::{
10    args::{InputArg, OutputArg},
11    definition::{MatmulElems, MatmulVectorSizes},
12    routines::{BatchMatmulRoutine, BlueprintStrategy},
13};
14use cubek_std::InputBinding;
15
16use crate::components::{ConvSetupError, ConvolutionProblem};
17
18/// Select which kernel to launch for the given Algorithm.
19///
20/// Only works for concrete tensor inputs and output.
21#[allow(clippy::result_large_err, clippy::too_many_arguments)]
22pub fn launch_kernel_concrete<
23    R: Runtime,
24    Args: ConcreteArgs<A>,
25    A: BatchMatmulRoutine<RuntimeArgs>,
26>(
27    client: &ComputeClient<R>,
28    out_grad: InputBinding<R>,
29    weights: InputBinding<R>,
30    in_grad: TensorBinding<R>,
31    problem: ConvolutionProblem,
32    vector_sizes: MatmulVectorSizes,
33    blueprint_strategy: &BlueprintStrategy<RuntimeArgs, A>,
34    dtypes: &MatmulElems,
35) -> Result<(), ConvSetupError> {
36    let mut view_vector_sizes = vector_sizes;
37
38    if let InputBinding::Quantized { scheme, .. } = out_grad {
39        view_vector_sizes.lhs *= scheme.num_quants();
40    }
41    if let InputBinding::Quantized { scheme, .. } = weights {
42        view_vector_sizes.rhs *= scheme.num_quants();
43    }
44
45    let device_settings = A::device_settings(client, view_vector_sizes);
46    let expand_info = A::expand_blueprint(
47        &problem.as_matmul_problem(),
48        &device_settings,
49        blueprint_strategy,
50    )?;
51
52    let problem = Args::adjust_problem(client, problem, &expand_info.blueprint, dtypes);
53    let launch_info = A::prepare(&problem.as_matmul_problem(), &device_settings, expand_info)?;
54
55    let (input, runtime_args) = <InputArg<Args> as ConcreteInputsFactory<A>>::create(
56        out_grad,
57        weights,
58        &launch_info.blueprint,
59        &problem,
60        dtypes,
61    );
62    let output = <OutputArg<Args> as ConcreteOutputFactory<A>>::create(
63        in_grad,
64        &launch_info.blueprint,
65        &problem,
66    );
67
68    let result = cubek_matmul::strategy::launch_kernel::<Args, R, A>(
69        client,
70        input,
71        output,
72        runtime_args,
73        launch_info,
74    );
75
76    result.map_err(ConvSetupError::Matmul)
77}
cubek_convolution/kernels/backward_data/selector/select_kernel.rs

cubek_convolution/kernels/backward_data/selector/
select_kernel.rs
