cubek-convolution 0.3.0

CubeK: Convolution Kernels
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84

use crate::{
    components::{ConvSetupError, ConvolutionProblem, global::args::RuntimeArgs},
    forward::args::{ConcreteArgs, ConcreteInputsFactory, ConcreteOutputFactory},
};
use cubecl::{client::Client, prelude::TensorBinding};
use cubek_matmul::{
    definition::{AccumulatorOperand, MatmulElems, MatmulVectorSizes},
    multi_level::{
        BatchMatmulRoutine,
        args::{InputArg, OutputArg},
    },
    routine::BlueprintStrategy,
};
use cubek_std::InputBinding;

/// Select which kernel to launch for the given Algorithm.
///
/// Only works for concrete tensor inputs and output.
#[allow(clippy::result_large_err, clippy::too_many_arguments)]
pub fn launch_kernel_concrete<Args: ConcreteArgs<A>, A: BatchMatmulRoutine<RuntimeArgs>>(
    client: &Client,
    input: InputBinding,
    weight: InputBinding,
    bias: Option<InputBinding>,
    out: TensorBinding,
    problem: ConvolutionProblem,
    vector_sizes: MatmulVectorSizes,
    blueprint_strategy: &BlueprintStrategy<Args::Config, A>,
    dtypes: &MatmulElems,
) -> Result<(), ConvSetupError> {
    // The bias is read into the accumulator through a stage of its own, which the
    // blueprint's shared-memory budget has to charge.
    let accumulator = match bias {
        Some(_) => AccumulatorOperand::Present,
        None => AccumulatorOperand::Absent,
    };

    let mut view_vector_sizes = vector_sizes;

    if let InputBinding::Quantized { scheme, .. } = input {
        view_vector_sizes.lhs *= scheme.num_quants();
    }
    if let InputBinding::Quantized { scheme, .. } = weight {
        view_vector_sizes.rhs *= scheme.num_quants();
    }

    let device_settings = A::device_settings(client, view_vector_sizes);
    let expand_info = A::expand_blueprint(
        &problem.as_matmul_problem(accumulator),
        &device_settings,
        blueprint_strategy,
    )?;

    let problem = Args::adjust_problem(client, problem, &expand_info.blueprint, dtypes);
    let launch_info = A::prepare(
        &problem.as_matmul_problem(accumulator),
        &device_settings,
        expand_info,
    )?;

    let (input, runtime_args) = <InputArg<Args> as ConcreteInputsFactory<A>>::create(
        input,
        weight,
        bias,
        &launch_info.blueprint,
        &problem,
        dtypes,
    );
    let output = <OutputArg<Args> as ConcreteOutputFactory<A>>::create(
        out,
        &launch_info.blueprint,
        &problem,
        dtypes,
    );

    cubek_matmul::multi_level::launch_kernel::<Args, A>(
        client,
        input,
        output,
        runtime_args,
        launch_info,
    )
    .map_err(ConvSetupError::Matmul)
}