cubecl_core/compute/
builder.rs

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

cubecl_core/compute/
builder.rs
