cubecl_ir/
address.rs

1use pliron::derive::format;
2
3use crate::{ElemType, GlobalState, IntKind, UIntKind};
4
5/// The type used for addressing storage types in a kernel.
6/// This is the type `usize` maps to when used in a kernel, with `isize` being mapped to the signed
7/// equivalent.
8#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
9#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, Default, PartialOrd, Ord)]
10#[format]
11pub enum AddressType {
12    // Discriminants are explicit to ensure correct ordering
13    #[default]
14    U32 = 0,
15    U64 = 1,
16}
17
18impl core::fmt::Display for AddressType {
19    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
20        match self {
21            AddressType::U32 => f.write_str("u32"),
22            AddressType::U64 => f.write_str("u64"),
23        }
24    }
25}
26
27impl AddressType {
28    /// Pick an address type based on the number of elements in a buffer.
29    pub fn from_len(num_elems: usize) -> Self {
30        if num_elems > u32::MAX as usize {
31            AddressType::U64
32        } else {
33            AddressType::U32
34        }
35    }
36
37    /// Pick an address type based on the number of elements in a buffer, for a kernel that requires
38    /// signed indices.
39    pub fn from_len_signed(num_elems: usize) -> Self {
40        if num_elems > i32::MAX as usize {
41            AddressType::U64
42        } else {
43            AddressType::U32
44        }
45    }
46
47    pub fn register(&self, state: &mut GlobalState) {
48        state.register_type::<usize>(self.unsigned_type());
49        state.register_type::<isize>(self.signed_type());
50    }
51
52    pub fn unsigned_type(&self) -> ElemType {
53        match self {
54            AddressType::U32 => UIntKind::U32.into(),
55            AddressType::U64 => UIntKind::U64.into(),
56        }
57    }
58
59    pub fn signed_type(&self) -> ElemType {
60        match self {
61            AddressType::U32 => IntKind::I32.into(),
62            AddressType::U64 => IntKind::I64.into(),
63        }
64    }
65
66    pub fn size(&self) -> usize {
67        match self {
68            AddressType::U32 => size_of::<u32>(),
69            AddressType::U64 => size_of::<u64>(),
70        }
71    }
72
73    pub fn size_bits(&self) -> usize {
74        self.size() * 8
75    }
76}
cubecl_ir/address.rs

cubecl_ir/
address.rs
