cubecl_core/codegen/
metadata.rs

1//! Metadata helpers to easily get offsets etc.
2//!
3//! Conceptually, metadata is represented like this:
4//! ```ignore
5//! struct Metadata<const NUM_BUFS: usize, const NUM_EXT: usize> {
6//!     base: BaseMeta<NUM_BUFS>,
7//!     extended: ExtendedMeta<NUM_EXT>,
8//! }
9//!
10//! struct BaseMeta<const N: usize> {
11//!     buffer_lengths: [usize; N],
12//! }
13//!
14//! struct ExtendedMeta<const N: usize> {
15//!     shape_offsets: [usize; N],
16//!     stride_offsets: [usize; N],
17//!     shapes: [usize],
18//!     strides: [usize]
19//! }
20//! ```
21//! where `Vec` isn't an actual `Vec`, just a dynamically sized series of values.
22//!
23//! Ranks and lengths have a constant offset, while shapes/strides involve loading the tensor's
24//! offset, then adding `dim` to the offset to get each shape/stride.
25
26use alloc::vec::Vec;
27use bytemuck::Pod;
28use cubecl_ir::{
29    AddressType,
30    metadata::{METADATA_BASE_LEN, METADATA_EXT_LEN},
31};
32use cubecl_zspace::{Shape, Strides};
33use num_traits::NumCast;
34
35/// Builder for a serialized metadata struct
36///
37/// Inputs/Outputs must be added in the same order they're defined in the bind group
38#[derive(Default)]
39pub struct MetadataBuilder {
40    state_32: State<u32>,
41    state_64: State<u64>,
42}
43
44#[derive(Default)]
45struct State<T: Pod> {
46    buffer_lens: Vec<T>,
47    shapes: Vec<T>,
48    strides: Vec<T>,
49
50    offsets: Vec<usize>,
51}
52
53impl MetadataBuilder {
54    /// Add an array to a builder
55    pub fn register_buffer(&mut self, buffer_len: u64, address_type: AddressType) {
56        match address_type {
57            AddressType::U64 => {
58                self.state_64.buffer_lens.push(buffer_len);
59            }
60            AddressType::U32 => {
61                self.state_32.buffer_lens.push(buffer_len as u32);
62            }
63        }
64    }
65
66    /// Add a tensor to a builder
67    pub fn register_tensor(
68        &mut self,
69        buffer_len: u64,
70        shape: Shape,
71        strides: Strides,
72        address_type: AddressType,
73    ) {
74        match address_type {
75            AddressType::U64 => {
76                let state = &mut self.state_64;
77                state.buffer_lens.push(buffer_len);
78                state.offsets.push(state.shapes.len());
79                state.shapes.extend(shape.iter().map(|s| *s as u64));
80                state.strides.extend(strides.iter().map(|s| *s as u64));
81            }
82            AddressType::U32 => {
83                let state = &mut self.state_32;
84                state.buffer_lens.push(buffer_len as u32);
85                state.offsets.push(state.shapes.len());
86                state.shapes.extend(shape.iter().map(|s| *s as u32));
87                state.strides.extend(strides.iter().map(|s| *s as u32));
88            }
89        }
90    }
91
92    pub fn static_len(&self, address_type: AddressType) -> usize {
93        let (base, ext) = match address_type {
94            AddressType::U32 => (self.state_32.buffer_lens.len(), self.state_32.offsets.len()),
95            AddressType::U64 => (self.state_64.buffer_lens.len(), self.state_64.offsets.len()),
96        };
97        base * METADATA_BASE_LEN + ext * METADATA_EXT_LEN
98    }
99
100    pub fn dynamic_len(&self, address_type: AddressType) -> usize {
101        match address_type {
102            AddressType::U32 => self.state_32.shapes.len() + self.state_32.strides.len(),
103            AddressType::U64 => self.state_64.shapes.len() + self.state_64.strides.len(),
104        }
105    }
106
107    /// Build the final serialized metadata struct
108    pub fn finish(&mut self, address_type: AddressType, out: (&mut [u64], &mut [u64])) {
109        fn finish_inner<T: Pod + NumCast>(state: &mut State<T>, out: (&mut [u64], &mut [u64])) {
110            let mut sized = bytemuck::cast_slice_mut::<u64, u8>(out.0);
111            let mut dynamic = bytemuck::cast_slice_mut::<u64, u8>(out.1);
112
113            {
114                let buffer_lens = bytemuck::cast_slice::<T, u8>(&state.buffer_lens);
115
116                sized[..buffer_lens.len()].copy_from_slice(buffer_lens);
117                sized = &mut sized[buffer_lens.len()..];
118            }
119
120            state.buffer_lens.clear();
121
122            let strides_offset_base = state.shapes.len();
123
124            for offs in state.offsets.iter() {
125                let offset = [T::from(*offs).unwrap()];
126                let bytes = bytemuck::cast_slice(&offset);
127                sized[..bytes.len()].copy_from_slice(bytes);
128                sized = &mut sized[size_of::<T>()..];
129            }
130
131            for offs in state.offsets.drain(..) {
132                let offset = [T::from(strides_offset_base + offs).unwrap()];
133                let bytes = bytemuck::cast_slice(&offset);
134                sized[..bytes.len()].copy_from_slice(bytes);
135                sized = &mut sized[size_of::<T>()..];
136            }
137
138            {
139                let shapes = bytemuck::cast_slice::<T, u8>(&state.shapes);
140                let strides = bytemuck::cast_slice::<T, u8>(&state.strides);
141
142                dynamic[..shapes.len()].copy_from_slice(shapes);
143                dynamic = &mut dynamic[shapes.len()..];
144
145                dynamic[..strides.len()].copy_from_slice(strides);
146            }
147
148            state.shapes.clear();
149            state.strides.clear();
150        }
151
152        match address_type {
153            AddressType::U32 => finish_inner(&mut self.state_32, out),
154            AddressType::U64 => finish_inner(&mut self.state_64, out),
155        }
156    }
157}
cubecl_core/codegen/metadata.rs

cubecl_core/codegen/
metadata.rs
