cubek_convolution/kernels/backward_data/selector/
select_kernel.rs

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

cubek_convolution/kernels/backward_data/selector/
select_kernel.rs
