cubecl_core/codegen/
info.rs

1use alloc::{vec, vec::Vec};
2
3use cubecl_ir::{AddressType, StorageType};
4use cubecl_runtime::{kernel::ScalarKernelArg, server::MetadataBindingInfo};
5
6use crate::{Metadata, MetadataBuilder, ScalarBuilder};
7
8pub(crate) const INFO_ALIGN: usize = size_of::<u64>();
9
10/// Helper to calculate info struct fields
11#[derive(Clone, Debug, Default)]
12pub struct Info {
13    pub scalars: Vec<SizedInfoField>,
14    pub sized_meta: Option<SizedInfoField>,
15    pub has_dynamic_meta: bool,
16    pub dynamic_meta_offset: usize,
17    pub metadata: Metadata,
18}
19
20#[derive(Clone, Copy, Debug)]
21pub struct SizedInfoField {
22    pub ty: StorageType,
23    pub size: usize,
24    pub offset: usize,
25}
26
27impl SizedInfoField {
28    pub fn padded_size(&self) -> usize {
29        let padding_factor = INFO_ALIGN / self.ty.size();
30        self.size.next_multiple_of(padding_factor)
31    }
32}
33
34impl Info {
35    pub fn new(scalars: &[ScalarKernelArg], metadata: Metadata, address_type: StorageType) -> Self {
36        let mut scalar_fields = Vec::with_capacity(scalars.len());
37        let mut sized_meta = None;
38
39        let mut offset = 0;
40
41        for scalar in scalars {
42            scalar_fields.push(SizedInfoField {
43                ty: scalar.ty,
44                size: scalar.count,
45                offset,
46            });
47            offset += (scalar.ty.size() * scalar.count).next_multiple_of(INFO_ALIGN);
48        }
49
50        if metadata.static_len() > 0 {
51            let size = metadata.static_len() as usize;
52            sized_meta = Some(SizedInfoField {
53                ty: address_type,
54                size,
55                offset,
56            });
57            offset += (address_type.size() * size).next_multiple_of(INFO_ALIGN);
58        }
59
60        Info {
61            scalars: scalar_fields,
62            sized_meta,
63            has_dynamic_meta: metadata.num_extended_meta() > 0,
64            dynamic_meta_offset: offset,
65            metadata,
66        }
67    }
68
69    pub fn has_info(&self) -> bool {
70        !self.scalars.is_empty() || self.sized_meta.is_some()
71    }
72}
73
74#[derive(Default)]
75pub struct InfoBuilder {
76    pub scalars: ScalarBuilder,
77    pub metadata: MetadataBuilder,
78}
79
80impl InfoBuilder {
81    pub fn finish(&mut self, address_type: AddressType) -> MetadataBindingInfo {
82        let addr_packing = INFO_ALIGN / address_type.size();
83
84        let scalars_size = self.scalars.len_aligned();
85        let static_len = self.metadata.static_len(address_type);
86        let static_size = static_len.div_ceil(addr_packing);
87        let dynamic_len = self.metadata.dynamic_len(address_type);
88        let dynamic_size = dynamic_len.div_ceil(addr_packing);
89
90        let mut out = vec![0; scalars_size + static_size + dynamic_size];
91        self.scalars.finish(&mut out[..scalars_size]);
92        self.metadata
93            .finish(address_type, out[scalars_size..].split_at_mut(static_size));
94
95        MetadataBindingInfo {
96            data: out,
97            dynamic_metadata_offset: scalars_size + static_size,
98        }
99    }
100}
cubecl_core/codegen/info.rs

cubecl_core/codegen/
info.rs
