Skip to main content

cubecl_core/codegen/
info.rs

1use alloc::vec;
2
3use cubecl_ir::{AddressType, metadata::INFO_ALIGN};
4use cubecl_runtime::server::MetadataBindingInfo;
5
6use crate::{MetadataBuilder, ScalarBuilder};
7
8#[derive(Default)]
9pub struct InfoBuilder {
10    pub scalars: ScalarBuilder,
11    pub metadata: MetadataBuilder,
12}
13
14impl InfoBuilder {
15    pub fn finish(&mut self, address_type: AddressType) -> MetadataBindingInfo {
16        let addr_packing = INFO_ALIGN / address_type.size();
17
18        let scalars_size = self.scalars.len_aligned();
19        let static_len = self.metadata.static_len(address_type);
20        let static_size = static_len.div_ceil(addr_packing);
21        let dynamic_len = self.metadata.dynamic_len(address_type);
22        let dynamic_size = dynamic_len.div_ceil(addr_packing);
23
24        let mut out = vec![0; scalars_size + static_size + dynamic_size];
25        self.scalars.finish(&mut out[..scalars_size]);
26        self.metadata
27            .finish(address_type, out[scalars_size..].split_at_mut(static_size));
28
29        MetadataBindingInfo {
30            data: out,
31            dynamic_metadata_offset: scalars_size + static_size,
32        }
33    }
34}