cubecl_spirv/
compiler.rs

1use crate::{
2    SpirvKernel,
3    debug::DebugInfo,
4    item::Item,
5    lookups::CompilerState,
6    target::{GLCompute, SpirvTarget},
7    transformers::{BitwiseTransform, ErfTransform, HypotTransform, RhypotTransform},
8};
9use cubecl_common::backtrace::BackTrace;
10use cubecl_core::{
11    Compiler, CubeDim, Info, Metadata, WgpuCompilationOptions,
12    ir::{self as core, ElemType, Id, InstructionModes, StorageType, UIntKind, features::EnumSet},
13    post_processing::{
14        checked_io::CheckedIoVisitor, disaggregate::DisaggregateVisitor,
15        saturating::SaturatingArithmeticProcessor, unroll::UnrollVisitor,
16    },
17    prelude::{FastMath, KernelDefinition, Visibility},
18    server::ExecutionMode,
19};
20use cubecl_opt::{
21    BasicBlock, Function, NodeIndex, Optimizer, OptimizerBuilder, SharedLiveness, Uniformity,
22};
23use cubecl_runtime::{
24    compiler::CompilationError,
25    config::{CubeClRuntimeConfig, RuntimeConfig, compilation::CompilationLogLevel},
26};
27use rspirv::{
28    binary::Assemble,
29    dr::{Builder, InsertPoint, Instruction, Module, Operand},
30    spirv::{BuiltIn, Capability, Decoration, FPFastMathMode, Op, StorageClass, Word},
31};
32use std::{
33    collections::HashSet,
34    fmt::Debug,
35    mem::take,
36    ops::{Deref, DerefMut},
37    rc::Rc,
38    sync::Arc,
39};
40
41pub struct SpirvCompiler<Target: SpirvTarget = GLCompute> {
42    pub target: Target,
43    pub(crate) builder: Builder,
44
45    pub cube_dim: CubeDim,
46    pub mode: ExecutionMode,
47    pub addr_type: StorageType,
48    pub debug_symbols: bool,
49    global_invocation_id: Word,
50    num_workgroups: Word,
51    pub setup_block: usize,
52    pub opt: Rc<Optimizer>,
53    pub uniformity: Rc<Uniformity>,
54    pub shared_liveness: Rc<SharedLiveness>,
55    pub current_func: Option<Id>,
56    pub current_block: Option<NodeIndex>,
57
58    pub capabilities: HashSet<Capability>,
59    pub state: CompilerState,
60    pub ext_meta_pos: Vec<u32>,
61    pub info: Info,
62    pub debug_info: Option<DebugInfo>,
63    pub compilation_options: WgpuCompilationOptions,
64}
65
66unsafe impl<T: SpirvTarget> Send for SpirvCompiler<T> {}
67unsafe impl<T: SpirvTarget> Sync for SpirvCompiler<T> {}
68
69impl<T: SpirvTarget> Clone for SpirvCompiler<T> {
70    fn clone(&self) -> Self {
71        Self {
72            target: self.target.clone(),
73            builder: Builder::new_from_module(self.module_ref().clone()),
74            cube_dim: self.cube_dim,
75            mode: self.mode,
76            addr_type: self.addr_type,
77            global_invocation_id: self.global_invocation_id,
78            num_workgroups: self.num_workgroups,
79            setup_block: self.setup_block,
80            opt: self.opt.clone(),
81            uniformity: self.uniformity.clone(),
82            shared_liveness: self.shared_liveness.clone(),
83            current_func: self.current_func,
84            current_block: self.current_block,
85            capabilities: self.capabilities.clone(),
86            state: self.state.clone(),
87            debug_symbols: self.debug_symbols,
88            info: self.info.clone(),
89            debug_info: self.debug_info.clone(),
90            ext_meta_pos: self.ext_meta_pos.clone(),
91            compilation_options: self.compilation_options,
92        }
93    }
94}
95
96fn debug_symbols_activated() -> bool {
97    matches!(
98        CubeClRuntimeConfig::get().compilation.logger.level,
99        CompilationLogLevel::Full
100    )
101}
102
103impl<T: SpirvTarget> Default for SpirvCompiler<T> {
104    fn default() -> Self {
105        Self {
106            target: Default::default(),
107            builder: Builder::new(),
108            cube_dim: CubeDim::new_single(),
109            mode: Default::default(),
110            addr_type: ElemType::UInt(UIntKind::U32).into(),
111            global_invocation_id: Default::default(),
112            num_workgroups: Default::default(),
113            capabilities: Default::default(),
114            state: Default::default(),
115            setup_block: Default::default(),
116            opt: Default::default(),
117            uniformity: Default::default(),
118            shared_liveness: Default::default(),
119            current_func: Default::default(),
120            current_block: Default::default(),
121            debug_symbols: debug_symbols_activated(),
122            info: Default::default(),
123            debug_info: Default::default(),
124            ext_meta_pos: Default::default(),
125            compilation_options: Default::default(),
126        }
127    }
128}
129
130impl<T: SpirvTarget> Deref for SpirvCompiler<T> {
131    type Target = Builder;
132
133    fn deref(&self) -> &Self::Target {
134        &self.builder
135    }
136}
137
138impl<T: SpirvTarget> DerefMut for SpirvCompiler<T> {
139    fn deref_mut(&mut self) -> &mut Self::Target {
140        &mut self.builder
141    }
142}
143
144impl<T: SpirvTarget> Compiler for SpirvCompiler<T> {
145    type Representation = SpirvKernel;
146    type CompilationOptions = WgpuCompilationOptions;
147
148    fn compile(
149        &mut self,
150        value: KernelDefinition,
151        compilation_options: &Self::CompilationOptions,
152        mode: ExecutionMode,
153        addr_type: StorageType,
154    ) -> Result<Self::Representation, CompilationError> {
155        let errors = value.body.pop_errors();
156        if !errors.is_empty() {
157            let mut reason = "Can't compile spirv kernel".to_string();
158            for error in errors {
159                reason += error.as_str();
160                reason += "\n";
161            }
162
163            return Err(CompilationError::Validation {
164                reason,
165                backtrace: BackTrace::capture(),
166            });
167        }
168
169        let bindings = value.buffers.clone();
170        let mut ext_meta_pos = Vec::new();
171        let mut num_ext = 0;
172
173        let mut all_meta: Vec<_> = value
174            .buffers
175            .iter()
176            .chain(value.tensor_maps.iter())
177            .map(|buf| (buf.id, buf.has_extended_meta))
178            .collect();
179        all_meta.sort_by_key(|(id, _)| *id);
180
181        let num_meta = all_meta.len();
182
183        for (_, has_extended_meta) in all_meta.iter() {
184            ext_meta_pos.push(num_ext);
185            if *has_extended_meta {
186                num_ext += 1;
187            }
188        }
189
190        let metadata = Metadata::new(num_meta as u32, num_ext);
191
192        self.cube_dim = value.cube_dim;
193        self.mode = mode;
194        self.addr_type = addr_type;
195        self.info = Info::new(&value.scalars, metadata, addr_type);
196        self.compilation_options = *compilation_options;
197        self.ext_meta_pos = ext_meta_pos;
198
199        let (module, optimizer, shared_size) = self.compile_kernel(value);
200        let info_visibility = match T::info_storage_class(self) {
201            StorageClass::Uniform => Visibility::Uniform,
202            _ => Visibility::Read,
203        };
204        let immediate_size = match T::params_storage_class(self, bindings.len()) {
205            StorageClass::PushConstant => Some((bindings.len() + 1) * size_of::<u64>()),
206            _ => None,
207        };
208
209        let visibility = self.opt.global_state.buffer_visibility.borrow();
210        let bindings = visibility
211            .iter()
212            .map(|vis| match vis.writable {
213                true => Visibility::ReadWrite,
214                false => Visibility::Read,
215            })
216            .collect();
217
218        Ok(SpirvKernel {
219            assembled_module: module.assemble(),
220            module: Some(Arc::new(module)),
221            optimizer: Some(Arc::new(optimizer)),
222            bindings,
223            shared_size,
224            immediate_size,
225            info_visibility,
226        })
227    }
228
229    fn elem_size(&self, elem: core::ElemType) -> usize {
230        elem.size()
231    }
232
233    fn extension(&self) -> &'static str {
234        "spv"
235    }
236}
237
238impl<Target: SpirvTarget> Debug for SpirvCompiler<Target> {
239    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
240        write!(f, "spirv<{:?}>", self.target)
241    }
242}
243
244impl<Target: SpirvTarget> SpirvCompiler<Target> {
245    pub fn compile_kernel(&mut self, mut kernel: KernelDefinition) -> (Module, Optimizer, usize) {
246        let options = kernel.options.clone();
247
248        self.debug_symbols = debug_symbols_activated() || options.debug_symbols;
249
250        let version = self.compilation_options.vulkan.max_spirv_version;
251        self.set_version(version.0, version.1);
252
253        let mut target = self.target.clone();
254
255        let mut opt = OptimizerBuilder::default()
256            .with_transformer(ErfTransform)
257            .with_transformer(BitwiseTransform::new(
258                self.compilation_options.vulkan.supports_arbitrary_bitwise,
259            ))
260            .with_transformer(HypotTransform)
261            .with_transformer(RhypotTransform)
262            .with_visitor(CheckedIoVisitor::new(
263                self.mode,
264                kernel.options.kernel_name.clone(),
265            ))
266            .with_visitor(DisaggregateVisitor::default())
267            .with_visitor(UnrollVisitor::new(
268                self.compilation_options.vulkan.max_vector_size,
269            ))
270            .with_processor(SaturatingArithmeticProcessor::new(true))
271            .optimize(kernel.body.clone(), kernel.cube_dim);
272
273        self.uniformity = opt.main.analysis::<Uniformity>(&opt.global_state);
274        self.shared_liveness = opt.main.analysis::<SharedLiveness>(&opt.global_state);
275        self.opt = Rc::new(opt);
276
277        self.init_debug();
278        self.init_base_state(&mut kernel);
279
280        let cube_dims = vec![kernel.cube_dim.x, kernel.cube_dim.y, kernel.cube_dim.z];
281
282        target.set_kernel_name(options.kernel_name.clone());
283
284        let opt = self.opt.clone();
285        for (id, func) in opt.global_state.extra_functions.iter() {
286            let def = self.declare_function(func);
287            self.current_func = Some(*id);
288
289            let blocks = func.breadth_first_dominators();
290            for block in blocks {
291                self.compile_block(block);
292            }
293            self.end_function_and_reset_lookups();
294
295            self.state.extra_funcs.insert(*id, def);
296            self.current_func = None;
297        }
298
299        self.init_kernel_state(kernel);
300
301        let (main, debug_setup) = self.declare_main(&options.kernel_name);
302
303        let setup = self.id();
304        self.debug_name(setup, "setup");
305
306        let entry = self.opt.entry();
307        let body = self.label(entry);
308
309        let setup_block = self.setup(setup, debug_setup);
310        self.setup_block = setup_block;
311
312        let shared_size = self.declare_shared_memories();
313
314        let blocks = opt.main.breadth_first_dominators();
315        for block in blocks {
316            self.compile_block(block);
317        }
318
319        self.select_block(Some(setup_block)).unwrap();
320        self.branch(body).unwrap();
321
322        // Don't reset the state here, need to keep used builtins around
323        self.end_function().unwrap();
324
325        let builtins = self
326            .state
327            .used_builtins
328            .clone()
329            .into_iter()
330            .map(|(builtin, (id, item))| {
331                let ty = Item::Pointer(StorageClass::Input, Box::new(item)).id(self);
332                self.variable(ty, Some(id), StorageClass::Input, None);
333                self.decorate(id, Decoration::BuiltIn, vec![builtin.into()]);
334                id
335            })
336            .collect::<Vec<_>>();
337
338        target.set_modes(self, main, builtins, cube_dims);
339
340        let module = take(&mut self.builder).module();
341        (module, self.opt.as_ref().clone(), shared_size)
342    }
343
344    fn setup(&mut self, label: Word, debug_setup: impl Fn(&mut Self)) -> usize {
345        self.begin_block(Some(label)).unwrap();
346
347        Target::load_params(self);
348
349        debug_setup(self);
350
351        let setup_block = self.selected_block().unwrap();
352        self.select_block(None).unwrap();
353        setup_block
354    }
355
356    #[track_caller]
357    pub fn current_block(&self) -> BasicBlock {
358        self.current_func()
359            .block(self.current_block.unwrap())
360            .clone()
361    }
362
363    pub fn current_func(&self) -> &Function {
364        self.current_func
365            .map(|func| &self.opt.global_state.extra_functions[&func])
366            .unwrap_or(&self.opt.main)
367    }
368
369    pub fn builtin(&mut self, builtin: BuiltIn, item: Item) -> Word {
370        if let Some(existing) = self.state.used_builtins.get(&builtin) {
371            existing.0
372        } else {
373            let id = self.id();
374            self.state.used_builtins.insert(builtin, (id, item));
375            id
376        }
377    }
378
379    pub fn compile_block(&mut self, block: NodeIndex) {
380        self.current_block = Some(block);
381
382        let label = self.label(block);
383        self.begin_block(Some(label)).unwrap();
384        let block_id = self.selected_block().unwrap();
385
386        self.debug_start_block();
387
388        let operations = self.current_block().ops.borrow().clone();
389        for (_, operation) in operations {
390            self.compile_operation(operation);
391        }
392
393        let control_flow = self.current_block().control_flow.borrow().clone();
394        self.compile_control_flow(control_flow);
395
396        let current = self.selected_block();
397        self.select_block(Some(block_id)).unwrap();
398        let phi = { self.current_func().block(block).phi_nodes.borrow().clone() };
399        for phi in phi {
400            let out = self.compile_value(phi.out);
401            let ty = out.item().id(self);
402            let out_id = self.write_id(&out);
403            let entries: Vec<_> = phi
404                .entries
405                .into_iter()
406                .map(|it| {
407                    let label = self.end_label(it.block);
408                    let value = self.compile_value(it.value);
409                    let value = self.read(&value);
410                    (value, label)
411                })
412                .collect();
413            self.insert_phi(InsertPoint::Begin, ty, Some(out_id), entries)
414                .unwrap();
415        }
416        self.select_block(current).unwrap();
417    }
418
419    // Declare variable in the first block of the function
420    pub fn declare_function_variable(&mut self, ty: Word, init: Option<Word>) -> Word {
421        let setup = self.setup_block;
422        let id = self.id();
423        let mut val = Instruction::new(
424            Op::Variable,
425            Some(ty),
426            Some(id),
427            vec![Operand::StorageClass(StorageClass::Function)],
428        );
429        if let Some(init) = init {
430            val.operands.push(Operand::IdRef(init));
431        }
432        let current_block = self.selected_block();
433        self.select_block(Some(setup)).unwrap();
434        self.insert_into_block(InsertPoint::Begin, val).unwrap();
435        self.select_block(current_block).unwrap();
436        id
437    }
438
439    fn declare_shared_memories(&mut self) -> usize {
440        if self.compilation_options.vulkan.supports_explicit_smem {
441            self.declare_shared_memories_explicit() as usize
442        } else {
443            self.declare_shared_memories_implicit() as usize
444        }
445    }
446
447    /// When using `VK_KHR_workgroup_memory_explicit_layout`, all shared memory is declared as a
448    /// `Block`. This means they are all pointers into the same chunk of memory, with different
449    /// offsets and sizes. Unlike C++, this shared block is declared implicitly, not explicitly.
450    /// Alignment and total size is calculated by the driver.
451    fn declare_shared_memories_explicit(&mut self) -> u32 {
452        let mut shared_size = 0;
453
454        let shared = self.state.shared.clone();
455        if shared.is_empty() {
456            return shared_size;
457        }
458
459        self.capabilities
460            .insert(Capability::WorkgroupMemoryExplicitLayoutKHR);
461
462        for (index, memory) in shared {
463            let memory_size = memory.item.size();
464            let value_size = memory.item.value_type().size();
465            shared_size = shared_size.max(memory.offset + memory_size);
466
467            // It's safe to assume that if 8-bit/16-bit types are supported, they're supported for
468            // explicit layout as well.
469            match value_size {
470                1 => {
471                    self.capabilities
472                        .insert(Capability::WorkgroupMemoryExplicitLayout8BitAccessKHR);
473                }
474                2 => {
475                    self.capabilities
476                        .insert(Capability::WorkgroupMemoryExplicitLayout16BitAccessKHR);
477                }
478                _ => {}
479            }
480
481            let item_id = memory.item.id(self);
482            let block_id = self.id();
483
484            if let Item::Array(_, _) = memory.item {
485                self.decorate(item_id, Decoration::ArrayStride, [value_size.into()]);
486            }
487
488            self.type_struct_id(Some(block_id), [item_id]);
489
490            self.decorate(block_id, Decoration::Block, []);
491            self.member_decorate(block_id, 0, Decoration::Offset, [memory.offset.into()]);
492
493            let block_ptr_ty = self.type_pointer(None, StorageClass::Workgroup, block_id);
494            let ptr_ty =
495                self.type_pointer(Some(memory.ptr_ty_id), StorageClass::Workgroup, item_id);
496
497            self.debug_shared(memory.id, index);
498            self.variable(
499                block_ptr_ty,
500                Some(memory.val_id),
501                StorageClass::Workgroup,
502                None,
503            );
504            self.decorate(memory.val_id, Decoration::Aliased, []);
505
506            self.insert_in_setup(|b| {
507                let zero = b.const_u32(0);
508                b.in_bounds_access_chain(ptr_ty, Some(memory.id), memory.val_id, [zero])
509                    .unwrap()
510            });
511        }
512
513        shared_size
514    }
515
516    fn declare_shared_memories_implicit(&mut self) -> u32 {
517        let mut shared_size = 0;
518
519        let shared = self.state.shared.clone();
520        for (index, memory) in shared {
521            shared_size += memory.item.size();
522
523            let ty_id = memory.item.id(self);
524            let ptr_ty = self.type_pointer(Some(memory.ptr_ty_id), StorageClass::Workgroup, ty_id);
525
526            self.debug_shared(memory.id, index);
527            self.variable(ptr_ty, Some(memory.id), StorageClass::Workgroup, None);
528        }
529        shared_size
530    }
531
532    pub fn declare_math_mode(&mut self, modes: InstructionModes, out_id: Word) {
533        if !self.compilation_options.vulkan.supports_fp_fast_math || modes.fp_math_mode.is_empty() {
534            return;
535        }
536        let mode = convert_math_mode(modes.fp_math_mode);
537        self.capabilities.insert(Capability::FloatControls2);
538        self.decorate(
539            out_id,
540            Decoration::FPFastMathMode,
541            [Operand::FPFastMathMode(mode)],
542        );
543    }
544
545    pub fn is_uniform_block(&self) -> bool {
546        self.uniformity
547            .is_block_uniform(self.current_block.unwrap())
548    }
549}
550
551pub(crate) fn convert_math_mode(math_mode: EnumSet<FastMath>) -> FPFastMathMode {
552    let mut flags = FPFastMathMode::NONE;
553
554    for mode in math_mode.iter() {
555        match mode {
556            FastMath::NotNaN => flags |= FPFastMathMode::NOT_NAN,
557            FastMath::NotInf => flags |= FPFastMathMode::NOT_INF,
558            FastMath::UnsignedZero => flags |= FPFastMathMode::NSZ,
559            FastMath::AllowReciprocal => flags |= FPFastMathMode::ALLOW_RECIP,
560            FastMath::AllowContraction => flags |= FPFastMathMode::ALLOW_CONTRACT,
561            FastMath::AllowReassociation => flags |= FPFastMathMode::ALLOW_REASSOC,
562            FastMath::AllowTransform => {
563                flags |= FPFastMathMode::ALLOW_CONTRACT
564                    | FPFastMathMode::ALLOW_REASSOC
565                    | FPFastMathMode::ALLOW_TRANSFORM
566            }
567            _ => {}
568        }
569    }
570
571    flags
572}
cubecl_spirv/compiler.rs

cubecl_spirv/
compiler.rs
