cubek_std/
cube_dim_resource.rs

1use cubecl::prelude::*;
2
3use crate::{InvalidConfigError, PlaneFlowConfig};
4
5#[derive(Debug)]
6/// Number of compute primitives required by some component, specified as either units, planes,
7/// or a specialized plane-flow split.
8pub enum CubeDimResource {
9    Units(u32),
10    Planes(u32),
11    Specialized(PlaneFlowConfig),
12}
13
14impl CubeDimResource {
15    /// Ensures [CubeDimResource] is the Planes variant, converting
16    /// units using plane_dim, the number of units in a plane.
17    ///
18    /// Will fail if the number of units does not correspond to an exact number of planes
19    pub fn as_plane_resource(self, plane_dim: u32) -> Result<Self, InvalidConfigError> {
20        match self {
21            CubeDimResource::Units(units) => {
22                if units % plane_dim == 0 {
23                    Ok(CubeDimResource::Planes(units / plane_dim))
24                } else {
25                    Err(Box::new(format!(
26                        "Number of units {units:?} should be divisible by plane_dim {plane_dim:?}"
27                    )))
28                }
29            }
30            CubeDimResource::Planes(_) => Ok(self),
31            CubeDimResource::Specialized(spec) => {
32                Ok(CubeDimResource::Planes(spec.counts.total_count()))
33            }
34        }
35    }
36
37    /// Make a [CubeDim] from specified resources.
38    ///
39    /// Obtained CubeDim is always (plane_dim, number_of_planes, 1)
40    ///
41    /// Will fail if the number of units does not correspond to an exact number of planes
42    pub fn to_cube_dim(self, plane_dim: u32) -> Result<CubeDim, InvalidConfigError> {
43        match self {
44            CubeDimResource::Units(_) => self.as_plane_resource(plane_dim)?.to_cube_dim(plane_dim),
45            CubeDimResource::Planes(num_planes) => Ok(CubeDim::new_2d(plane_dim, num_planes)),
46            CubeDimResource::Specialized(_) => {
47                self.as_plane_resource(plane_dim)?.to_cube_dim(plane_dim)
48            }
49        }
50    }
51
52    /// Get the number of planes
53    ///
54    /// Will fail if the number of units does not correspond to an exact number of planes
55    pub fn num_planes(self, plane_dim: u32) -> Result<u32, InvalidConfigError> {
56        let plane_resources = self.as_plane_resource(plane_dim)?;
57        if let CubeDimResource::Planes(num_planes) = plane_resources {
58            Ok(num_planes)
59        } else {
60            unreachable!()
61        }
62    }
63
64    /// Recover the [PlaneFlowConfig] view of this resource. `Units`/`Planes` produce a
65    /// non-specialized config (all planes in the main flow).
66    pub fn as_specialized(self, plane_dim: u32) -> Result<PlaneFlowConfig, InvalidConfigError> {
67        match self {
68            CubeDimResource::Units(_) => {
69                self.as_plane_resource(plane_dim)?.as_specialized(plane_dim)
70            }
71            CubeDimResource::Planes(num_planes) => {
72                Ok(PlaneFlowConfig::new_unspecialized(num_planes))
73            }
74            CubeDimResource::Specialized(spec) => Ok(spec),
75        }
76    }
77}
cubek_std/cube_dim_resource.rs

cubek_std/
cube_dim_resource.rs
