cubecl_spirv/
lookups.rs

1use core::ops::{Deref, DerefMut};
2use std::collections::VecDeque;
3
4use cubecl_core::{
5    ir::{self, Builtin, Id, ValueKind},
6    prelude::KernelDefinition,
7};
8use cubecl_opt::NodeIndex;
9use hashbrown::{HashMap, HashSet};
10use rspirv::{
11    dr,
12    spirv::{
13        self, BuiltIn, CooperativeMatrixLayout, CooperativeMatrixUse, MemoryAccess, Scope,
14        StorageClass, Word,
15    },
16};
17
18use crate::{
19    SpirvCompiler, SpirvTarget,
20    item::{Elem, Item},
21    value::ConstVal,
22};
23
24#[derive(Clone, Debug, Default)]
25pub struct CompilerState {
26    pub extra_funcs: HashMap<Id, FuncDefinition>,
27
28    pub scalar_bindings: HashMap<ir::StorageType, u32>,
29    pub params: Word,
30    pub info: Option<Buffer>,
31    pub cube_dims: Vec<Word>,
32    pub cube_size: Word,
33
34    // For break, continue
35    pub loops: VecDeque<Loop>,
36
37    pub debug_types: HashSet<Word>,
38
39    /// Base lookups, used to query global parameters like buffers and shared memory for function
40    /// definitions
41    pub base_lookups: LookupTables,
42    /// Lookups for the current function
43    pub lookups: LookupTables,
44}
45
46impl Deref for CompilerState {
47    type Target = LookupTables;
48
49    fn deref(&self) -> &Self::Target {
50        &self.lookups
51    }
52}
53
54impl DerefMut for CompilerState {
55    fn deref_mut(&mut self) -> &mut Self::Target {
56        &mut self.lookups
57    }
58}
59
60#[derive(Clone, Debug, Default)]
61pub struct LookupTables {
62    pub buffers: Vec<Buffer>,
63    pub shared: HashMap<Id, SharedVal>,
64
65    pub globals: HashMap<Builtin, Word>,
66    pub loaded_builtins: HashMap<BuiltIn, Word>,
67    pub used_builtins: HashMap<BuiltIn, (Word, Item)>,
68
69    pub constants: HashMap<(ConstVal, Item), Word>,
70    pub values: HashMap<Id, Word>,
71    pub labels: HashMap<NodeIndex, Word>,
72    pub end_labels: HashMap<NodeIndex, Word>,
73}
74
75#[derive(Clone, Debug)]
76pub struct FuncDefinition {
77    pub type_id: Word,
78    pub id: Word,
79}
80
81#[derive(Clone, Debug)]
82pub struct SharedVal {
83    pub id: Word,
84    pub val_id: Word,
85    pub ptr_ty_id: Word,
86    pub item: Item,
87    pub offset: u32,
88    pub align: u32,
89}
90
91#[derive(Debug, Clone, Copy, PartialEq)]
92#[allow(missing_docs)]
93pub struct Matrix {
94    pub id: Word,
95    pub ident: CooperativeMatrixUse,
96    pub scope: Scope,
97    pub m: u32,
98    pub n: u32,
99    pub k: u32,
100    pub elem: Elem,
101    pub layout: Option<CooperativeMatrixLayout>,
102}
103
104#[derive(Clone, Debug)]
105pub struct Loop {
106    pub header: Word,
107    pub continue_target: Word,
108    pub post: Word,
109}
110
111#[derive(Debug, Clone, Copy, PartialEq)]
112pub struct Buffer {
113    pub id: Word,
114    pub struct_ty_id: Word,
115    pub struct_ptr_ty_id: Word,
116    pub arr_ty_id: Word,
117    pub arr_ptr_ty_id: Word,
118    pub storage_class: StorageClass,
119}
120
121impl<T: SpirvTarget> SpirvCompiler<T> {
122    pub fn init_base_state(&mut self, kernel: &mut KernelDefinition) {
123        let cube_dims = [kernel.cube_dim.x, kernel.cube_dim.y, kernel.cube_dim.z];
124        self.state.cube_dims = cube_dims.iter().map(|dim| self.const_u32(*dim)).collect();
125        self.state.cube_size = self.const_u32(cube_dims.iter().product());
126
127        let mut target = self.target.clone();
128
129        let max_vector_size = self.compilation_options.vulkan.max_vector_size;
130        for binding in &mut kernel.buffers {
131            // This is safe when combined with the unroll transform that adjusts all indices.
132            // Must not be used alone
133            if binding.value.ty.vector_size() > max_vector_size {
134                binding.value.ty = binding.value.ty.with_vector_size(max_vector_size);
135            }
136        }
137
138        let opt = self.opt.clone();
139        let mut visibility = opt.global_state.buffer_visibility.borrow_mut();
140        // Just in case not all buffers were accessed when tracking reads/writes
141        visibility.resize(kernel.num_global_buffers(), Default::default());
142        self.state.base_lookups.buffers =
143            target.generate_params(self, &kernel.buffers, &visibility);
144
145        let shared_liveness = self.shared_liveness.clone();
146        for alloc in shared_liveness.allocations.values() {
147            let smem_id = self.id();
148            let smem_ptr_ty_id = self.id();
149
150            let smem_val_id = if self.compilation_options.vulkan.supports_explicit_smem {
151                self.id()
152            } else {
153                smem_id
154            };
155
156            let item = self.compile_type(alloc.smem.value_ty);
157            self.state.base_lookups.shared.insert(
158                alloc.id,
159                SharedVal {
160                    id: smem_id,
161                    val_id: smem_val_id,
162                    ptr_ty_id: smem_ptr_ty_id,
163                    item,
164                    offset: alloc.offset as u32,
165                    align: alloc.smem.alignment as u32,
166                },
167            );
168        }
169
170        self.state.lookups = self.state.base_lookups.clone();
171    }
172
173    pub fn init_kernel_state(&mut self, kernel: KernelDefinition) {
174        self.state.scalar_bindings = kernel
175            .scalars
176            .into_iter()
177            .enumerate()
178            .map(|(i, arg)| (arg.ty, i as u32))
179            .collect();
180        self.state.lookups = self.state.base_lookups.clone();
181    }
182
183    fn dedup_const(&mut self, inst: &dr::Instruction) -> Option<Word> {
184        self.module_ref()
185            .types_global_values
186            .iter()
187            .find(|it| {
188                it.class == inst.class
189                    && it.result_type == inst.result_type
190                    && it.operands == inst.operands
191            })
192            .and_then(|it| it.result_id)
193    }
194
195    pub fn dedup_constant_bit32(&mut self, ty: Word, val: u32) -> Word {
196        let inst = dr::Instruction::new(
197            spirv::Op::Constant,
198            Some(ty),
199            None,
200            vec![dr::Operand::LiteralBit32(val)],
201        );
202        if let Some(id) = self.dedup_const(&inst) {
203            id
204        } else {
205            self.constant_bit32(ty, val)
206        }
207    }
208
209    pub fn dedup_constant_bit64(&mut self, ty: Word, val: u64) -> Word {
210        let inst = dr::Instruction::new(
211            spirv::Op::Constant,
212            Some(ty),
213            None,
214            vec![dr::Operand::LiteralBit64(val)],
215        );
216        if let Some(id) = self.dedup_const(&inst) {
217            id
218        } else {
219            self.constant_bit64(ty, val)
220        }
221    }
222
223    pub fn const_u32(&mut self, value: u32) -> Word {
224        let ty = Item::Scalar(Elem::Int(32, false));
225        let ty_id = ty.id(self);
226        self.dedup_constant_bit32(ty_id, value)
227    }
228
229    pub fn insert_builtin(
230        &mut self,
231        builtin: BuiltIn,
232        insert: impl FnOnce(&mut Self) -> Word,
233    ) -> Word {
234        if let Some(id) = self.state.loaded_builtins.get(&builtin) {
235            *id
236        } else {
237            let id = self.insert_in_setup(insert);
238            self.state.loaded_builtins.insert(builtin, id);
239            id
240        }
241    }
242
243    pub fn insert_global(
244        &mut self,
245        builtin: Builtin,
246        insert: impl FnOnce(&mut Self) -> Word,
247    ) -> Word {
248        if let Some(id) = self.state.globals.get(&builtin) {
249            *id
250        } else {
251            let id = self.insert_in_setup(insert);
252            self.state.globals.insert(builtin, id);
253            id
254        }
255    }
256
257    pub fn insert_in_setup(&mut self, insert: impl FnOnce(&mut Self) -> Word) -> Word {
258        let current_block = self.selected_block();
259        let setup = self.setup_block;
260        self.select_block(Some(setup)).unwrap();
261        let id = insert(self);
262        self.select_block(current_block).unwrap();
263        id
264    }
265
266    pub fn insert_in_root(&mut self, insert: impl FnOnce(&mut Self) -> Word) -> Word {
267        let current_block = self.selected_block();
268        self.select_block(None).unwrap();
269        let id = insert(self);
270        self.select_block(current_block).unwrap();
271        id
272    }
273
274    pub fn get_value(&mut self, id: Id) -> Word {
275        if let Some(existing) = self.state.values.get(&id) {
276            *existing
277        } else {
278            let word = self.id();
279            self.state.values.insert(id, word);
280            self.debug_val_name(word, id);
281            word
282        }
283    }
284
285    pub fn insert_value(&mut self, id: Id, word: Word) {
286        self.state.values.insert(id, word);
287    }
288
289    pub fn label(&mut self, block: NodeIndex) -> Word {
290        if let Some(existing) = self.state.labels.get(&block) {
291            *existing
292        } else {
293            let word = self.id();
294            self.debug_name(word, format!("bb{}", block.index()));
295            self.state.labels.insert(block, word);
296            word
297        }
298    }
299
300    pub fn end_label(&mut self, block: NodeIndex) -> Word {
301        if let Some(existing) = self.state.end_labels.get(&block) {
302            *existing
303        } else {
304            let word = self.label(block);
305            self.state.end_labels.insert(block, word);
306            word
307        }
308    }
309
310    pub fn init_function_param(&mut self, param: ir::Value, param_id: Word) {
311        let item = self.compile_type(param.ty);
312        match param.kind {
313            ValueKind::Constant(value) => {
314                let const_val = (value, item.clone()).into();
315                self.state.constants.insert((const_val, item), param_id);
316            }
317            ValueKind::Value { id } => {
318                self.state.values.insert(id, param_id);
319            }
320        }
321    }
322
323    pub fn end_function_and_reset_lookups(&mut self) {
324        self.builder.end_function().unwrap();
325        self.state.lookups = self.state.base_lookups.clone();
326    }
327
328    pub fn global_scalar(&mut self, id: Id, ty: ir::StorageType) -> Word {
329        self.insert_in_setup(|b| {
330            let field_id = b.const_u32(b.state.scalar_bindings[&ty]);
331            let offset = b.const_u32(id);
332            let item = b.compile_type(ir::Type::new(ty));
333            let align = item.size();
334            let ty_id = item.id(b);
335            let storage_class = T::info_storage_class(b);
336            let ptr_ty = Item::Pointer(storage_class, Box::new(item)).id(b);
337            let info = b.state.info.unwrap().id;
338            let access = b
339                .in_bounds_access_chain(ptr_ty, None, info, [field_id, offset])
340                .unwrap();
341            b.load(
342                ty_id,
343                None,
344                access,
345                Some(MemoryAccess::ALIGNED),
346                [align.into()],
347            )
348            .unwrap()
349        })
350    }
351}
cubecl_spirv/lookups.rs

cubecl_spirv/
lookups.rs
