cubecl_ir/
runtime_properties.rs

1use pliron::{context::Context, r#type::TypedHandle};
2
3use crate::{
4    TypeHash, VectorSize,
5    types::matrix::{MatrixIdent, MatrixLayout, MatrixType},
6};
7
8/// Hacky solution for getting comptime properties into the scope.
9/// Allows querying certain target-specific properties at compile time, rather than at runtime.
10/// Review on how to better solve this and delegate to the compiler if possible.
11#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
12#[derive(Debug, Clone, PartialEq, Eq, TypeHash, Default)]
13pub struct TargetProperties {
14    pub mma: MmaProperties,
15}
16
17#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
18#[derive(Debug, Clone, PartialEq, Eq, TypeHash)]
19pub struct MmaProperties {
20    /// Size of registers in bits, used to calculate vector size
21    pub register_size_bits: usize,
22    /// Constant size of planes, for calculating lane indices in a matrix
23    pub const_plane_size: u32,
24    /// Layout of registers in Matrix A
25    pub register_layout_a: MatrixLayout,
26    /// Layout of registers in Matrix B
27    pub register_layout_b: MatrixLayout,
28    /// Layout of registers in Matrix C/D
29    pub register_layout_acc: MatrixLayout,
30
31    /// How many copies of each piece of data exist for matrix A
32    pub register_duplication_a: usize,
33    /// How many copies of each piece of data exist for matrix B
34    pub register_duplication_b: usize,
35    /// How many copies of each piece of data exist for matrix C/D
36    pub register_duplication_acc: usize,
37    #[cfg_attr(feature = "serde", serde(skip))]
38    pub contiguous_elements: ContiguousElements,
39}
40
41#[derive(Clone)]
42pub struct ContiguousElements {
43    #[allow(clippy::type_complexity)]
44    inner: alloc::sync::Arc<
45        dyn Fn(&Context, MatrixIdent, TypedHandle<MatrixType>) -> VectorSize + Send + Sync,
46    >,
47}
48
49impl ContiguousElements {
50    pub fn new(
51        func: impl Fn(&Context, MatrixIdent, TypedHandle<MatrixType>) -> VectorSize
52        + Send
53        + Sync
54        + 'static,
55    ) -> Self {
56        Self {
57            inner: alloc::sync::Arc::new(func),
58        }
59    }
60
61    pub fn apply(
62        &self,
63        ctx: &Context,
64        ident: MatrixIdent,
65        matrix: TypedHandle<MatrixType>,
66    ) -> VectorSize {
67        (self.inner)(ctx, ident, matrix)
68    }
69}
70
71impl Default for ContiguousElements {
72    fn default() -> Self {
73        Self {
74            inner: alloc::sync::Arc::new(|_, _, _| 2),
75        }
76    }
77}
78
79impl core::fmt::Debug for ContiguousElements {
80    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
81        f.debug_struct("ContiguousElements").finish()
82    }
83}
84
85impl Eq for ContiguousElements {}
86impl PartialEq for ContiguousElements {
87    fn eq(&self, other: &Self) -> bool {
88        alloc::sync::Arc::ptr_eq(&self.inner, &other.inner)
89    }
90}
91
92impl TypeHash for ContiguousElements {
93    fn write_hash(hasher: &mut impl core::hash::Hasher) {
94        hasher.write_i32(0);
95    }
96}
97
98impl Default for MmaProperties {
99    fn default() -> Self {
100        Self {
101            register_size_bits: 32,
102            const_plane_size: 32,
103            register_layout_a: MatrixLayout::RowMajor,
104            register_layout_b: MatrixLayout::ColMajor,
105            register_layout_acc: MatrixLayout::RowMajor,
106            register_duplication_a: 1,
107            register_duplication_b: 1,
108            register_duplication_acc: 1,
109            contiguous_elements: Default::default(),
110        }
111    }
112}
cubecl_ir/runtime_properties.rs

cubecl_ir/
runtime_properties.rs
