cubek_matmul/
routine.rs

1use cubecl::prelude::*;
2use cubecl::std::tensor::{MatrixBatchLayout, matrix_batch_layout};
3use cubek_std::InputBinding;
4use std::fmt::{Debug, Display};
5
6use crate::definition::{MatmulSetupError, MatmulVectorSizes};
7
8/// A launch-time config carried alongside the tensors. `()` for a plain matmul; the
9/// fusion path substitutes its own.
10pub trait RuntimeConfig: LaunchArg + CubeType<ExpandType: Clone> + Clone + Send + Sync {}
11impl<T: LaunchArg + CubeType<ExpandType: Clone> + Clone + Send + Sync> RuntimeConfig for T {}
12
13/// A stride-0 (broadcast) matrix dim, or any layout that isn't row/col-major,
14/// classifies as [`MatrixBatchLayout::HighlyPermuted`] and can't be consumed
15/// directly; materialize such an operand into a contiguous tensor. A broadcast
16/// *batch* dim is only `MildlyPermuted` and stays untouched (handled natively).
17///
18/// A storage-tiled operand is passed as it is: its strides order its tiles, not a matrix's rows
19/// and columns, so this classification does not read them, and a routine that reads storage
20/// tiles checks the layout it was handed itself.
21#[allow(clippy::result_large_err)]
22pub(crate) fn into_contiguous_if_highly_permuted(
23    client: &Client,
24    binding: InputBinding,
25) -> Result<InputBinding, MatmulSetupError> {
26    if binding.data().tiling.is_tiled() {
27        return Ok(binding);
28    }
29    match matrix_batch_layout(&binding.data().strides, binding.scheme()) {
30        MatrixBatchLayout::HighlyPermuted => Ok(binding.into_contiguous(client)?),
31        _ => Ok(binding),
32    }
33}
34
35/// The contract to solve a matmul
36pub trait Routine<RC: RuntimeConfig>: Sized {
37    type Strategy: Default + Display + Clone;
38    type Blueprint: Debug + Clone;
39}
40
41/// What the routine reads about the device before it can shape a blueprint.
42pub struct DeviceSettings {
43    pub client: Client,
44    pub plane_dim: u32,
45    pub vector_sizes: MatmulVectorSizes,
46    pub max_cube_count: (u32, u32, u32),
47}
48
49pub enum BlueprintStrategy<RC: RuntimeConfig, A: Routine<RC>> {
50    /// Use a predefined blueprint
51    Forced(A::Blueprint),
52    /// Allows to give limited blueprint information, and the rest is inferred from it
53    Inferred(A::Strategy),
54}
55
56impl<RC: RuntimeConfig, A: Routine<RC>> BlueprintStrategy<RC, A> {
57    pub fn maybe_forced_default(s: &Option<A::Blueprint>) -> Self {
58        s.as_ref()
59            .map(|s| Self::Forced(s.clone()))
60            .unwrap_or_default()
61    }
62    pub fn maybe_forced_or(s: &Option<A::Blueprint>, args: &A::Strategy) -> Self {
63        s.as_ref()
64            .map(|s| Self::Forced(s.clone()))
65            .unwrap_or_else(|| Self::Inferred(args.clone()))
66    }
67}
68
69impl<RC: RuntimeConfig, A: Routine<RC>> Default for BlueprintStrategy<RC, A> {
70    fn default() -> Self {
71        Self::Inferred(Default::default())
72    }
73}
74
75impl<RC: RuntimeConfig, A: Routine<RC>> Display for BlueprintStrategy<RC, A> {
76    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77        match self {
78            Self::Forced(_) => f.write_str("_forced"),
79            Self::Inferred(strategy) => write!(f, "{}", strategy),
80        }
81    }
82}
83
84impl<RC: RuntimeConfig, A: Routine<RC>> Clone for BlueprintStrategy<RC, A> {
85    fn clone(&self) -> Self {
86        match self {
87            Self::Forced(blueprint) => Self::Forced(blueprint.clone()),
88            Self::Inferred(strategy) => Self::Inferred(strategy.clone()),
89        }
90    }
91}
cubek_matmul/routine.rs

cubek_matmul/
routine.rs
