cubek_convolution/kernels/forward/selector/
select_kernel.rs

1use crate::{
2    components::{ConvSetupError, ConvolutionProblem, global::args::RuntimeArgs},
3    forward::args::{ConcreteArgs, ConcreteInputsFactory, ConcreteOutputFactory},
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    input: InputBinding,
23    weight: InputBinding,
24    bias: Option<InputBinding>,
25    out: TensorBinding,
26    problem: ConvolutionProblem,
27    vector_sizes: MatmulVectorSizes,
28    blueprint_strategy: &BlueprintStrategy<Args::Config, A>,
29    dtypes: &MatmulElems,
30) -> Result<(), ConvSetupError> {
31    // The bias is read into the accumulator through a stage of its own, which the
32    // blueprint's shared-memory budget has to charge.
33    let accumulator = match bias {
34        Some(_) => AccumulatorOperand::Present,
35        None => AccumulatorOperand::Absent,
36    };
37
38    let mut view_vector_sizes = vector_sizes;
39
40    if let InputBinding::Quantized { scheme, .. } = input {
41        view_vector_sizes.lhs *= scheme.num_quants();
42    }
43    if let InputBinding::Quantized { scheme, .. } = weight {
44        view_vector_sizes.rhs *= scheme.num_quants();
45    }
46
47    let device_settings = A::device_settings(client, view_vector_sizes);
48    let expand_info = A::expand_blueprint(
49        &problem.as_matmul_problem(accumulator),
50        &device_settings,
51        blueprint_strategy,
52    )?;
53
54    let problem = Args::adjust_problem(client, problem, &expand_info.blueprint, dtypes);
55    let launch_info = A::prepare(
56        &problem.as_matmul_problem(accumulator),
57        &device_settings,
58        expand_info,
59    )?;
60
61    let (input, runtime_args) = <InputArg<Args> as ConcreteInputsFactory<A>>::create(
62        input,
63        weight,
64        bias,
65        &launch_info.blueprint,
66        &problem,
67        dtypes,
68    );
69    let output = <OutputArg<Args> as ConcreteOutputFactory<A>>::create(
70        out,
71        &launch_info.blueprint,
72        &problem,
73        dtypes,
74    );
75
76    cubek_matmul::multi_level::launch_kernel::<Args, A>(
77        client,
78        input,
79        output,
80        runtime_args,
81        launch_info,
82    )
83    .map_err(ConvSetupError::Matmul)
84}
cubek_convolution/kernels/forward/selector/select_kernel.rs

cubek_convolution/kernels/forward/selector/
select_kernel.rs
