cubecl_ir/
metadata.rs

1use core::ops::Deref;
2
3use alloc::vec::Vec;
4use pliron::context::Context;
5
6use crate::{ContextExt, ElemType};
7
8pub const INFO_ALIGN: usize = size_of::<u64>();
9
10// Metadata
11const BUFFER_LEN: usize = 0;
12pub const METADATA_BASE_LEN: usize = 1;
13
14// Extended Metadata
15const SHAPE_OFFSETS: usize = 0;
16const STRIDE_OFFSETS: usize = 1;
17pub const METADATA_EXT_LEN: usize = 2;
18
19/// Helper to calculate metadata offsets based on buffer count and position
20#[derive(Clone, Copy, Debug, Default, Hash)]
21pub struct Metadata {
22    num_meta: usize,
23    num_extended_meta: usize,
24}
25
26impl Metadata {
27    pub fn new(num_meta: usize, num_extended_meta: usize) -> Self {
28        Self {
29            num_meta,
30            num_extended_meta,
31        }
32    }
33
34    fn offset_of(&self, id: usize) -> usize {
35        self.num_meta * id
36    }
37
38    fn base_len(&self) -> usize {
39        self.num_meta * METADATA_BASE_LEN
40    }
41
42    pub fn static_len(&self) -> usize {
43        self.num_meta * METADATA_BASE_LEN + self.num_extended_meta * METADATA_EXT_LEN
44    }
45
46    pub fn num_meta(&self) -> usize {
47        self.num_meta
48    }
49
50    pub fn num_extended_meta(&self) -> usize {
51        self.num_extended_meta
52    }
53
54    fn offset_of_extended(&self, id: usize) -> usize {
55        self.base_len() + self.num_extended_meta * id
56    }
57
58    pub fn buffer_len_index(&self, buffer_idx: usize) -> usize {
59        self.offset_of(BUFFER_LEN) + buffer_idx
60    }
61
62    pub fn shape_offset_index(&self, buffer_idx: usize) -> usize {
63        self.offset_of_extended(SHAPE_OFFSETS) + buffer_idx
64    }
65
66    pub fn stride_offset_index(&self, buffer_idx: usize) -> usize {
67        self.offset_of_extended(STRIDE_OFFSETS) + buffer_idx
68    }
69}
70
71/// Helper to calculate info struct fields
72#[derive(Clone, Debug, Default, Hash)]
73pub struct Info {
74    pub scalars: Vec<SizedInfoField>,
75    pub sized_meta: Option<SizedInfoField>,
76    pub has_dynamic_meta: bool,
77    pub dynamic_meta_offset: usize,
78    pub metadata: Metadata,
79}
80
81impl Deref for Info {
82    type Target = Metadata;
83
84    fn deref(&self) -> &Self::Target {
85        &self.metadata
86    }
87}
88
89#[derive(Clone, Copy, Debug, Hash)]
90pub struct SizedInfoField {
91    pub ty: ElemType,
92    pub count: usize,
93    pub offset: usize,
94}
95
96impl SizedInfoField {
97    pub fn padded_size(&self, ctx: &Context) -> usize {
98        let padding_factor = INFO_ALIGN / self.ty.expand_size(ctx.address_type());
99        self.count.next_multiple_of(padding_factor)
100    }
101}
102
103impl Info {
104    pub fn has_info(&self) -> bool {
105        !self.scalars.is_empty() || self.sized_meta.is_some()
106    }
107}
cubecl_ir/metadata.rs

cubecl_ir/
metadata.rs
