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 buffer in memory shared with the other devices of the runtime, and return the
55    /// [element](Value) to be used for kernel expansion.
56    pub fn peer_buffer(&mut self, value_ty: TypeHandle) -> Value {
57        let id = self.inc_buffer_id();
58        self.scope.peer_global(id, value_ty)
59    }
60
61    /// Register a tensor and return the [element](Value) to be used for kernel expansion.
62    pub fn tensor(&mut self, value_ty: TypeHandle) -> Value {
63        let id = self.inc_buffer_id();
64        let ext_id = self.inc_ext_meta_id();
65        self.scope.global(id, Some(ext_id), value_ty)
66    }
67
68    /// Register a tensor map and return the [element](Value) to be used for kernel expansion.
69    pub fn tensor_map(&mut self) -> Value {
70        let id = self.inc_buffer_id();
71        let ext_id = self.inc_ext_meta_id();
72        self.scope.tensor_map(id, ext_id)
73    }
74
75    /// Register an output that uses the same resource as the input as the given position.
76    pub fn inplace(&mut self, position: usize) -> Value {
77        self.scope.kernel_arg(position)
78    }
79
80    pub fn runtime_properties(&mut self, properties: TargetProperties) {
81        self.scope.state_mut().target_properties = properties;
82    }
83
84    pub fn device_properties(&mut self, properties: Arc<DeviceProperties>) {
85        self.scope.device_properties(properties);
86    }
87
88    /// Build the [kernel definition](KernelDefinition).
89    pub fn build(self) -> KernelDefinition {
90        for warning in self.scope.pop_warnings() {
91            log::warn!("[{}] {warning}", self.settings.kernel_name);
92        }
93        if let Some(debug) = self.scope.debug_state() {
94            debug.finish();
95        }
96        let info = self.create_info();
97        KernelIntegrator::new(KernelExpansion {
98            scope: self.scope,
99            info,
100        })
101        .integrate(self.settings)
102    }
103
104    fn create_info(&self) -> Info {
105        let address_type = self.settings.address_type;
106        let metadata = Metadata::new(self.buffer_idx, self.ext_meta_idx);
107        let mut scalar_fields = Vec::with_capacity(self.scalars.len());
108        let mut sized_meta = None;
109
110        let mut offset = 0;
111
112        for (&ty, &count) in self.scalars.iter() {
113            scalar_fields.push(SizedInfoField { ty, count, offset });
114            offset += (ty.expand_size(address_type) * count).next_multiple_of(INFO_ALIGN);
115        }
116
117        if metadata.static_len() > 0 {
118            let size = metadata.static_len();
119            sized_meta = Some(SizedInfoField {
120                ty: address_type.unsigned_type(),
121                count: size,
122                offset,
123            });
124            offset += (address_type.size() * size).next_multiple_of(INFO_ALIGN);
125        }
126
127        Info {
128            scalars: scalar_fields,
129            sized_meta,
130            has_dynamic_meta: metadata.num_extended_meta() > 0,
131            dynamic_meta_offset: offset,
132            metadata,
133        }
134    }
135
136    pub fn new(mut settings: KernelSettings) -> Self {
137        settings.debug_info = CubeClRuntimeConfig::get()
138            .compilation
139            .resolve_debug_info(settings.debug_info);
140
141        Self {
142            scope: Scope::root(settings.clone()),
143            scalars: Default::default(),
144            settings,
145            buffer_idx: 0,
146            ext_meta_idx: 0,
147        }
148    }
149}
cubecl_core/compute/builder.rs

cubecl_core/compute/
builder.rs
