cubecl_core/codegen/
integrator.rs

1use alloc::{string::ToString, vec::Vec};
2use cubecl_ir::{Id, Scope, StorageType, Value};
3use cubecl_runtime::{
4    kernel::{KernelArg, KernelDefinition, KernelOptions, ScalarKernelArg},
5    server::CubeDim,
6};
7
8use crate::prelude::AddressType;
9
10/// The kernel integrator allows you to create a [kernel definition](KernelDefinition) based on
11/// [kernel expansion](KernelExpansion) and [kernel settings](KernelSettings).
12#[derive(Clone, Debug)]
13pub struct KernelIntegrator {
14    expansion: KernelExpansion,
15    buffer_bindings: Vec<KernelArg>,
16    scalar_bindings: Vec<ScalarKernelArg>,
17    tensor_maps: Vec<KernelArg>,
18}
19
20/// The information necessary to compile a [kernel definition](KernelDefinition).
21#[derive(Clone, Debug)]
22pub struct KernelExpansion {
23    pub buffers: Vec<BufferInfo>,
24    pub scalars: Vec<ScalarInfo>,
25    pub tensor_maps: Vec<BufferInfo>,
26    pub scope: Scope,
27}
28
29#[derive(Clone, Debug, PartialEq, Eq)]
30pub struct KernelSettings {
31    pub cube_dim: CubeDim,
32    pub address_type: AddressType,
33    pub options: KernelOptions,
34}
35
36impl Default for KernelSettings {
37    fn default() -> Self {
38        Self {
39            cube_dim: CubeDim::new_1d(1),
40            address_type: AddressType::U32,
41            options: Default::default(),
42        }
43    }
44}
45
46impl KernelSettings {
47    /// Set cube dimension.
48    pub fn cube_dim(mut self, cube_dim: CubeDim) -> Self {
49        self.cube_dim = cube_dim;
50        self
51    }
52
53    /// Set address type.
54    pub fn address_type(mut self, ty: AddressType) -> Self {
55        self.address_type = ty;
56        self
57    }
58
59    /// Set kernel name.
60    pub fn kernel_name<S: AsRef<str>>(mut self, name: S) -> Self {
61        self.options.kernel_name = name.as_ref().to_string();
62        self
63    }
64
65    /// Activate debug symbols
66    pub fn debug_symbols(mut self) -> Self {
67        self.options.debug_symbols = true;
68        self
69    }
70
71    /// Set cluster dim
72    pub fn cluster_dim(mut self, cluster_dim: CubeDim) -> Self {
73        self.options.cluster_dim = Some(cluster_dim);
74        self
75    }
76}
77
78/// Information related to a buffer binding.
79#[derive(Clone, Debug)]
80pub struct BufferInfo {
81    pub id: Id,
82    pub value: Value,
83    /// Whether this input has extended metadata (rank, shape, strides)
84    pub has_extended_meta: bool,
85}
86
87/// Information related to a scalar input.
88#[derive(Clone, Debug)]
89pub struct ScalarInfo {
90    pub ty: StorageType,
91    pub count: usize,
92}
93
94impl KernelIntegrator {
95    /// Starts a new compilation.
96    pub fn new(info: KernelExpansion) -> Self {
97        Self {
98            expansion: info,
99            buffer_bindings: Default::default(),
100            scalar_bindings: Default::default(),
101            tensor_maps: Default::default(),
102        }
103    }
104
105    /// Performs the compilation with the provided [settings](KernelSettings).
106    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip(self)))]
107    pub fn integrate(mut self, settings: KernelSettings) -> KernelDefinition {
108        self.register_buffers();
109        self.register_scalars();
110        self.register_tensor_maps();
111
112        self.scalar_bindings.sort_by_key(|binding| binding.ty);
113
114        KernelDefinition {
115            buffers: self.buffer_bindings,
116            tensor_maps: self.tensor_maps,
117            scalars: self.scalar_bindings,
118            cube_dim: settings.cube_dim,
119            body: self.expansion.scope,
120            options: settings.options,
121        }
122    }
123
124    fn register_buffers(&mut self) {
125        for buffer in self.expansion.buffers.drain(..) {
126            self.buffer_bindings.push(KernelArg {
127                id: buffer.id,
128                value: buffer.value,
129                has_extended_meta: buffer.has_extended_meta,
130            });
131        }
132    }
133
134    fn register_scalars(&mut self) {
135        for scalar in self.expansion.scalars.drain(..) {
136            self.scalar_bindings.push(ScalarKernelArg {
137                ty: scalar.ty,
138                count: scalar.count,
139            });
140        }
141    }
142
143    fn register_tensor_maps(&mut self) {
144        for buffer in self.expansion.tensor_maps.drain(..) {
145            self.tensor_maps.push(KernelArg {
146                id: buffer.id,
147                value: buffer.value,
148                has_extended_meta: buffer.has_extended_meta,
149            });
150        }
151    }
152}
cubecl_core/codegen/integrator.rs

cubecl_core/codegen/
integrator.rs
