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        let info = self.create_info();
89        KernelIntegrator::new(KernelExpansion {
90            scope: self.scope,
91            info,
92        })
93        .integrate(self.settings)
94    }
95
96    fn create_info(&self) -> Info {
97        let address_type = self.settings.address_type;
98        let metadata = Metadata::new(self.buffer_idx, self.ext_meta_idx);
99        let mut scalar_fields = Vec::with_capacity(self.scalars.len());
100        let mut sized_meta = None;
101
102        let mut offset = 0;
103
104        for (&ty, &count) in self.scalars.iter() {
105            scalar_fields.push(SizedInfoField { ty, count, offset });
106            offset += (ty.expand_size(address_type) * count).next_multiple_of(INFO_ALIGN);
107        }
108
109        if metadata.static_len() > 0 {
110            let size = metadata.static_len();
111            sized_meta = Some(SizedInfoField {
112                ty: address_type.unsigned_type(),
113                count: size,
114                offset,
115            });
116            offset += (address_type.size() * size).next_multiple_of(INFO_ALIGN);
117        }
118
119        Info {
120            scalars: scalar_fields,
121            sized_meta,
122            has_dynamic_meta: metadata.num_extended_meta() > 0,
123            dynamic_meta_offset: offset,
124            metadata,
125        }
126    }
127
128    pub fn new(mut settings: KernelSettings) -> Self {
129        let debug = DEBUG.load(Ordering::Relaxed);
130        let debug = if debug == -1 {
131            let val = match CubeClRuntimeConfig::get().compilation.logger.level {
132                CompilationLogLevel::Full => 1,
133                _ => 0,
134            };
135
136            DEBUG.store(val, Ordering::Relaxed);
137            val == 1
138        } else {
139            debug == 1
140        };
141        settings.debug_symbols |= debug;
142
143        Self {
144            scope: Scope::root(settings.clone()),
145            scalars: Default::default(),
146            settings,
147            buffer_idx: 0,
148            ext_meta_idx: 0,
149        }
150    }
151}
cubecl_core/compute/builder.rs

cubecl_core/compute/
builder.rs
