cubek_convolution/routines/
base.rs

1use crate::components::{ConvolutionOperation, global::args::RuntimeArgs};
2use cubecl::{
3    prelude::*,
4    std::tensor::{into_contiguous_pitched, is_contiguous_pitched},
5};
6use cubek_matmul::{
7    definition::AvailableVectorSizes,
8    multi_level::{BatchMatmulRoutine as MatmulRoutine, args::MatmulArgs, definition::Blueprint},
9};
10use std::fmt::Display;
11
12/// Specifications for a convolution routine.
13///
14/// A `Routine` is the convolution-side counterpart of `cubek_matmul::routine::Routine`:
15/// it pairs a per-operation matmul routine with the metadata needed to wire the
16/// kernel up (input args, optional layout fixups, vector-size filtering).
17///
18/// `Blueprint` and `Strategy` are surfaced as direct associated types so callers
19/// don't have to reach through `MatmulRoutine` to bound them.
20pub trait Routine {
21    type Blueprint: Blueprint;
22    type Strategy: Default + Display + Clone;
23
24    type MatmulRoutine: MatmulRoutine<RuntimeArgs, Blueprint = Self::Blueprint, Strategy = Self::Strategy>;
25    type Args: MatmulArgs<Config = RuntimeArgs>;
26
27    /// Whether to select specialized load flow in tests. Should replace with something cleaner
28    /// eventually, but this is nice and simple.
29    const IS_SPECIALIZED: bool = false;
30
31    fn correct_layout(
32        client: &Client,
33        handle: TensorBinding,
34        dtype: ElemType,
35        operation: ConvolutionOperation,
36    ) -> Result<TensorBinding, LaunchError>;
37
38    fn filter_vector_sizes(vector_sizes: AvailableVectorSizes) -> AvailableVectorSizes {
39        vector_sizes
40    }
41}
42
43pub(crate) fn contiguous_pitched_layout(
44    client: &Client,
45    binding: TensorBinding,
46    dtype: ElemType,
47) -> Result<TensorBinding, LaunchError> {
48    let binding = if has_valid_layout(&binding) {
49        binding
50    } else {
51        into_contiguous_pitched(client, binding, dtype).binding()
52    };
53    Ok(binding)
54}
55
56fn has_valid_layout(binding: &TensorBinding) -> bool {
57    let rank = binding.shape.len();
58    let dim_c = rank - 1;
59    binding.strides[dim_c] == 1
60}
61
62const TMA_STRIDE_ALIGN: usize = 16;
63
64pub(crate) fn into_tensor_handle_tma(
65    client: &Client,
66    handle: TensorBinding,
67    dtype: ElemType,
68    operation: ConvolutionOperation,
69) -> Result<TensorBinding, LaunchError> {
70    let binding = if has_valid_layout_tma(&handle, dtype, operation) {
71        handle
72    } else {
73        into_contiguous_pitched(client, handle, dtype).binding()
74    };
75    Ok(binding)
76}
77
78pub(crate) fn has_valid_layout_tma(
79    binding: &TensorBinding,
80    dtype: ElemType,
81    operation: ConvolutionOperation,
82) -> bool {
83    let stride_align = TMA_STRIDE_ALIGN / dtype.size();
84    let rank = binding.shape.len();
85    let dim_c = rank - 1;
86
87    let aligned = binding.strides[..dim_c]
88        .iter()
89        .all(|stride| stride % stride_align == 0);
90
91    let valid_layout = binding.strides[dim_c] == 1;
92
93    let is_valid_wgrad = if operation == ConvolutionOperation::BackwardWeight {
94        is_contiguous_pitched(&binding.shape, &binding.strides)
95    } else {
96        true
97    };
98
99    valid_layout && aligned && is_valid_wgrad
100}
cubek_convolution/routines/base.rs

cubek_convolution/routines/
base.rs
