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::rc::Rc<dyn Fn(&Context, MatrixIdent, TypedHandle<MatrixType>) -> VectorSize>,
45}
46
47impl ContiguousElements {
48    pub fn new(
49        func: impl Fn(&Context, MatrixIdent, TypedHandle<MatrixType>) -> VectorSize + 'static,
50    ) -> Self {
51        Self {
52            inner: alloc::rc::Rc::new(func),
53        }
54    }
55
56    pub fn apply(
57        &self,
58        ctx: &Context,
59        ident: MatrixIdent,
60        matrix: TypedHandle<MatrixType>,
61    ) -> VectorSize {
62        (self.inner)(ctx, ident, matrix)
63    }
64}
65
66impl Default for ContiguousElements {
67    fn default() -> Self {
68        Self {
69            inner: alloc::rc::Rc::new(|_, _, _| 2),
70        }
71    }
72}
73
74impl core::fmt::Debug for ContiguousElements {
75    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
76        f.debug_struct("ContiguousElements").finish()
77    }
78}
79
80impl Eq for ContiguousElements {}
81impl PartialEq for ContiguousElements {
82    fn eq(&self, other: &Self) -> bool {
83        alloc::rc::Rc::ptr_eq(&self.inner, &other.inner)
84    }
85}
86
87impl TypeHash for ContiguousElements {
88    fn write_hash(hasher: &mut impl core::hash::Hasher) {
89        hasher.write_i32(0);
90    }
91}
92
93impl Default for MmaProperties {
94    fn default() -> Self {
95        Self {
96            register_size_bits: 32,
97            const_plane_size: 32,
98            register_layout_a: MatrixLayout::RowMajor,
99            register_layout_b: MatrixLayout::ColMajor,
100            register_layout_acc: MatrixLayout::RowMajor,
101            register_duplication_a: 1,
102            register_duplication_b: 1,
103            register_duplication_acc: 1,
104            contiguous_elements: Default::default(),
105        }
106    }
107}
cubecl_ir/runtime_properties.rs

cubecl_ir/
runtime_properties.rs
