cubecl_core/compute/
builder.rs

1use alloc::{sync::Arc, vec::Vec};
2use derive_more::Deref;
3use pliron::r#type::TypeHandle;
4
5use crate::{KernelExpansion, KernelIntegrator, prelude::KernelDefinition};
6use alloc::collections::BTreeMap;
7use cubecl_ir::{
8    DeviceProperties, ElemType, Scope, TargetProperties,
9    metadata::{INFO_ALIGN, Info, Metadata, SizedInfoField},
10    pliron::value::Value,
11    settings::KernelSettings,
12};
13use cubecl_runtime::config::{CubeClRuntimeConfig, RuntimeConfig};
14
15/// Prepare a kernel to create a [`KernelDefinition`].
16#[derive(Deref)]
17pub struct KernelBuilder {
18    /// Cube [scope](Scope).
19    #[deref]
20    pub scope: Scope,
21    scalars: BTreeMap<ElemType, usize>,
22    buffer_idx: usize,
23    ext_meta_idx: usize,
24    settings: KernelSettings,
25}
26
27impl KernelBuilder {
28    /// Register a scalar and return the [element](Value) to be used for kernel expansion.
29    pub fn scalar(&mut self, storage: ElemType) -> usize {
30        let current_id = self.scalars.entry(storage).or_default();
31        let id = *current_id;
32        *current_id += 1;
33        id
34    }
35
36    fn inc_buffer_id(&mut self) -> usize {
37        let id = self.buffer_idx;
38        self.buffer_idx += 1;
39        id
40    }
41
42    fn inc_ext_meta_id(&mut self) -> usize {
43        let id = self.ext_meta_idx;
44        self.ext_meta_idx += 1;
45        id
46    }
47
48    /// Register a buffer and return the [element](Value) to be used for kernel expansion.
49    pub fn buffer(&mut self, value_ty: TypeHandle) -> Value {
50        let id = self.inc_buffer_id();
51        self.scope.global(id, None, value_ty)
52    }
53
54    /// Register a tensor and return the [element](Value) to be used for kernel expansion.
55    pub fn tensor(&mut self, value_ty: TypeHandle) -> Value {
56        let id = self.inc_buffer_id();
57        let ext_id = self.inc_ext_meta_id();
58        self.scope.global(id, Some(ext_id), value_ty)
59    }
60
61    /// Register a tensor map and return the [element](Value) to be used for kernel expansion.
62    pub fn tensor_map(&mut self) -> Value {
63        let id = self.inc_buffer_id();
64        let ext_id = self.inc_ext_meta_id();
65        self.scope.tensor_map(id, ext_id)
66    }
67
68    /// Register an output that uses the same resource as the input as the given position.
69    pub fn inplace(&mut self, position: usize) -> Value {
70        self.scope.kernel_arg(position)
71    }
72
73    pub fn runtime_properties(&mut self, properties: TargetProperties) {
74        self.scope.state_mut().target_properties = properties;
75    }
76
77    pub fn device_properties(&mut self, properties: Arc<DeviceProperties>) {
78        self.scope.device_properties(properties);
79    }
80
81    /// Build the [kernel definition](KernelDefinition).
82    pub fn build(self) -> KernelDefinition {
83        for warning in self.scope.pop_warnings() {
84            log::warn!("[{}] {warning}", self.settings.kernel_name);
85        }
86        if let Some(debug) = self.scope.debug_state() {
87            debug.finish();
88        }
89        let info = self.create_info();
90        KernelIntegrator::new(KernelExpansion {
91            scope: self.scope,
92            info,
93        })
94        .integrate(self.settings)
95    }
96
97    fn create_info(&self) -> Info {
98        let address_type = self.settings.address_type;
99        let metadata = Metadata::new(self.buffer_idx, self.ext_meta_idx);
100        let mut scalar_fields = Vec::with_capacity(self.scalars.len());
101        let mut sized_meta = None;
102
103        let mut offset = 0;
104
105        for (&ty, &count) in self.scalars.iter() {
106            scalar_fields.push(SizedInfoField { ty, count, offset });
107            offset += (ty.expand_size(address_type) * count).next_multiple_of(INFO_ALIGN);
108        }
109
110        if metadata.static_len() > 0 {
111            let size = metadata.static_len();
112            sized_meta = Some(SizedInfoField {
113                ty: address_type.unsigned_type(),
114                count: size,
115                offset,
116            });
117            offset += (address_type.size() * size).next_multiple_of(INFO_ALIGN);
118        }
119
120        Info {
121            scalars: scalar_fields,
122            sized_meta,
123            has_dynamic_meta: metadata.num_extended_meta() > 0,
124            dynamic_meta_offset: offset,
125            metadata,
126        }
127    }
128
129    pub fn new(mut settings: KernelSettings) -> Self {
130        settings.debug_info = CubeClRuntimeConfig::get()
131            .compilation
132            .resolve_debug_info(settings.debug_info);
133
134        Self {
135            scope: Scope::root(settings.clone()),
136            scalars: Default::default(),
137            settings,
138            buffer_idx: 0,
139            ext_meta_idx: 0,
140        }
141    }
142}
cubecl_core/compute/builder.rs

cubecl_core/compute/
builder.rs
