cubecl_opt/
lib.rs

1//! # `CubeCL` Optimizer
2//!
3//! A library that parses `CubeCL` IR into a
4//! [control flow graph](https://en.wikipedia.org/wiki/Control-flow_graph), transforms it to
5//! [static single-assignment form](https://en.wikipedia.org/wiki/Static_single-assignment_form)
6//! and runs various optimizations on it.
7//! The order of operations is as follows:
8//!
9//! 1. Parse root scope recursively into a [control flow graph](https://en.wikipedia.org/wiki/Control-flow_graph)
10//! 2. Run optimizations that must be done before SSA transformation
11//! 3. Analyze variable liveness
12//! 4. Transform the graph to [pruned SSA](https://en.wikipedia.org/wiki/Static_single-assignment_form#Pruned_SSA) form
13//! 5. Run post-SSA optimizations and analyses in a loop until no more improvements are found
14//! 6. Speed
15//!
16//! The output is represented as a [`petgraph`] graph of [`BasicBlock`]s terminated by [`ControlFlow`].
17//! This can then be compiled into actual executable code by walking the graph and generating all
18//! phi nodes, instructions and branches.
19//!
20//! # Representing [`PhiInstruction`] in non-SSA languages
21//!
22//! Phi instructions can be simulated by generating a mutable variable for each phi, then assigning
23//! `value` to it in each relevant `block`.
24//!
25
26#![no_std]
27#![allow(unknown_lints, unnecessary_transmutes)]
28
29extern crate alloc;
30
31#[cfg(any(feature = "std", test))]
32extern crate std;
33
34use core::{
35    cell::RefCell,
36    ops::{Deref, DerefMut},
37};
38
39use alloc::{boxed::Box, collections::vec_deque::VecDeque, rc::Rc, vec, vec::Vec};
40use analyses::{AnalysisCache, dominance::DomFrontiers, writes::LocalStores};
41use cubecl_core::{
42    CubeDim,
43    post_processing::{
44        analysis_helper::GlobalAnalyses,
45        constant_prop::{ConstEval, ConstOperandSimplify},
46        disaggregate::DisaggregateVisitor,
47        visitor::InstructionVisitor,
48    },
49};
50use cubecl_ir::{
51    self as ir, AddressSpace, Allocator, Branch, Id, Instruction, Operation, Processor, Scope,
52    Type, Value,
53};
54use gvn::GvnPass;
55use hashbrown::HashMap;
56use passes::{
57    EliminateConstBranches, EliminateDeadBlocks, EliminateDeadPhi, EliminateUnusedVariables,
58    EmptyBranchToSelect, InlineCopies, MergeBlocks, MergeSameExpressions, OptimizerPass,
59    ReduceStrength,
60};
61use petgraph::{Direction, prelude::StableDiGraph, visit::EdgeRef};
62
63mod analyses;
64mod block;
65mod control_flow;
66mod debug;
67mod gvn;
68mod instructions;
69mod passes;
70mod phi_frontiers;
71mod transformers;
72mod version;
73
74pub(crate) use cubecl_core::post_processing::util::AtomicCounter;
75
76pub use analyses::uniformity::Uniformity;
77pub use block::*;
78pub use control_flow::*;
79pub use petgraph::graph::{EdgeIndex, NodeIndex};
80pub use transformers::*;
81pub use version::PhiInstruction;
82
83pub use crate::analyses::liveness::MemoryLiveness;
84pub use crate::analyses::liveness::shared::SharedLiveness;
85use crate::{
86    analyses::{
87        dominance::Dominators,
88        liveness::{Captures, Liveness},
89        pointer_source::PointerSource,
90    },
91    passes::{CopyTransform, DisaggregateArray},
92};
93
94#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
95pub struct MemoryBlock {
96    pub address_space: AddressSpace,
97    pub value_ty: Type,
98    pub alignment: usize,
99    /// The root pointer value returned from the allocation or passed into the kernel. All other
100    /// pointers into the same value are derived from this.
101    pub root_ptr: Value,
102}
103
104impl MemoryBlock {
105    /// The byte size of this shared memory
106    pub fn size(&self) -> usize {
107        self.value_ty.size()
108    }
109}
110
111#[derive(Default, Debug, Clone)]
112pub struct Function {
113    /// Explicit parameters passed to the function, i.e. the inputs to a closure
114    pub explicit_params: Vec<Value>,
115    /// Implicit parameters passed to the function, i.e. kernel args, closure captures
116    pub implicit_params: Vec<Value>,
117    pub memories: HashMap<Id, MemoryBlock>,
118    pub graph: StableDiGraph<BasicBlock, u32>,
119    pub root: NodeIndex,
120    /// The single return block
121    pub ret: NodeIndex,
122    /// The return value, if any
123    pub return_value: Option<Value>,
124
125    /// Analyses with persistent state
126    analysis_cache: Rc<AnalysisCache>,
127    /// The current block while parsing
128    current_block: Option<NodeIndex>,
129    /// The current loop's break target
130    loop_break: VecDeque<NodeIndex>,
131}
132
133impl Deref for Function {
134    type Target = StableDiGraph<BasicBlock, u32>;
135
136    fn deref(&self) -> &Self::Target {
137        &self.graph
138    }
139}
140
141impl DerefMut for Function {
142    fn deref_mut(&mut self) -> &mut Self::Target {
143        &mut self.graph
144    }
145}
146
147/// An optimizer that applies various analyses and optimization passes to the IR.
148#[derive(Debug, Clone, Default)]
149pub struct Optimizer {
150    pub main: Function,
151    /// The overall program state
152    pub global_state: GlobalState,
153}
154
155#[derive(Debug, Clone)]
156pub struct GlobalState {
157    /// Allocator for kernel
158    pub allocator: Allocator,
159    /// Root scope to allocate variables on
160    pub root_scope: Scope,
161    pub buffer_visibility: RefCell<Vec<BufferVisibility>>,
162    pub extra_functions: HashMap<Id, Function>,
163    /// The `CubeDim` used for range analysis
164    #[allow(dead_code)]
165    pub(crate) cube_dim: CubeDim,
166    pub(crate) transformers: Vec<Rc<dyn IrTransformer>>,
167    pub(crate) processors: Rc<Vec<Box<dyn Processor>>>,
168    pub(crate) visitors: Rc<RefCell<Vec<Box<dyn InstructionVisitor>>>>,
169}
170
171#[derive(Debug, Clone, Default)]
172pub struct BufferVisibility {
173    /// Whether the buffer is ever read from
174    pub readable: bool,
175    /// Whether the buffer is ever written to
176    pub writable: bool,
177}
178
179// Needed for WGPU server
180unsafe impl Send for Optimizer {}
181unsafe impl Sync for Optimizer {}
182
183impl Default for GlobalState {
184    fn default() -> Self {
185        Self {
186            allocator: Default::default(),
187            root_scope: Scope::root(false),
188            buffer_visibility: Default::default(),
189            extra_functions: Default::default(),
190            cube_dim: CubeDim::new_1d(1),
191            transformers: Default::default(),
192            processors: Default::default(),
193            visitors: Default::default(),
194        }
195    }
196}
197
198impl Optimizer {
199    /// Create a new optimizer with the scope, `CubeDim` and execution mode passed into the compiler.
200    /// Parses the scope and runs several optimization and analysis loops.
201    pub fn new(
202        expand: Scope,
203        cube_dim: CubeDim,
204        transformers: Vec<Rc<dyn IrTransformer>>,
205        visitors: Vec<Box<dyn InstructionVisitor>>,
206        processors: Vec<Box<dyn Processor>>,
207    ) -> Self {
208        let extra_funcs = expand.state().functions.clone();
209        let mut global_state = GlobalState {
210            allocator: expand.state().allocator.clone(),
211            root_scope: expand.clone(),
212            buffer_visibility: Default::default(),
213            cube_dim,
214            transformers,
215            processors: Rc::new(processors),
216            visitors: Rc::new(RefCell::new(visitors)),
217            extra_functions: Default::default(),
218        };
219        for (id, func) in extra_funcs.into_iter() {
220            let mut function = Function {
221                explicit_params: func.explicit_params,
222                return_value: func.scope.return_value,
223                ..Default::default()
224            };
225            function.run_opt(&global_state, func.scope);
226            global_state.extra_functions.insert(id, function);
227        }
228        let mut root_func = Function::default();
229        root_func.run_opt(&global_state, expand);
230
231        Self {
232            global_state,
233            main: root_func,
234        }
235    }
236
237    /// Create a new optimizer with the scope, `CubeDim` and execution mode passed into the compiler.
238    /// Parses the scope and runs several optimization and analysis loops.
239    pub fn shared_only(expand: Scope, cube_dim: CubeDim) -> Self {
240        let extra_funcs = expand.state().functions.clone();
241        let disaggregate: Box<dyn InstructionVisitor> = Box::new(DisaggregateVisitor::default());
242        let mut global_state = GlobalState {
243            allocator: expand.state().allocator.clone(),
244            root_scope: expand.clone(),
245            buffer_visibility: Default::default(),
246            cube_dim,
247            transformers: Vec::new(),
248            processors: Rc::new(vec![]),
249            visitors: Rc::new(RefCell::new(vec![disaggregate])),
250            extra_functions: Default::default(),
251        };
252        for (id, func) in extra_funcs.into_iter() {
253            let mut function = Function {
254                explicit_params: func.explicit_params,
255                ..Default::default()
256            };
257            function.run_shared_only(&global_state, func.scope);
258            global_state.extra_functions.insert(id, function);
259        }
260        let mut root_func = Function::default();
261        root_func.run_shared_only(&global_state, expand);
262
263        Self {
264            global_state,
265            main: root_func,
266        }
267    }
268
269    /// The entry block of the program
270    pub fn entry(&self) -> NodeIndex {
271        self.main.root
272    }
273}
274
275impl GlobalState {
276    fn set_buffer_readable(&self, id: Id) {
277        let mut buffer_vis = self.buffer_visibility.borrow_mut();
278        let idx = id as usize;
279        if idx >= buffer_vis.len() {
280            buffer_vis.resize(idx + 1, Default::default());
281        }
282        buffer_vis[idx].readable = true;
283    }
284
285    fn set_buffer_writable(&self, id: Id) {
286        let mut buffer_vis = self.buffer_visibility.borrow_mut();
287        let idx = id as usize;
288        if idx >= buffer_vis.len() {
289            buffer_vis.resize(idx + 1, Default::default());
290        }
291        buffer_vis[idx].writable = true;
292    }
293}
294
295pub fn global_buffer_id(value: &ir::Value) -> Option<Id> {
296    match value.address_space() {
297        AddressSpace::Global(id) => Some(id),
298        _ => None,
299    }
300}
301
302impl Function {
303    fn parse_graph(&mut self, state: &GlobalState, scope: Scope) {
304        let entry = self.add_node(BasicBlock::default());
305        self.root = entry;
306        self.current_block = Some(entry);
307        self.ret = self.add_node(BasicBlock::default());
308        *self[self.ret].control_flow.borrow_mut() = ControlFlow::Return {
309            value: self.return_value,
310        };
311
312        self.parse_scope(state, scope);
313        if let Some(current_block) = self.current_block {
314            let ret = self.ret;
315            self.add_edge(current_block, ret, 0);
316        }
317        // Analyses shouldn't have run at this point, but just in case they have, invalidate
318        // all analyses that depend on the graph
319        self.invalidate_structure();
320    }
321
322    /// Recursively parse a scope into the graph
323    pub fn parse_scope(&mut self, state: &GlobalState, scope: Scope) -> bool {
324        let global_analyses = GlobalAnalyses::default();
325        global_analyses.recalculate_pointer_source(&scope);
326        global_analyses.recalculate_used_values(&scope);
327        for visitor in state.visitors.borrow_mut().iter_mut() {
328            visitor.visit_scope(&scope, &global_analyses, &AtomicCounter::new(0));
329        }
330        let processed = scope.process(state.processors.iter().map(|it| &**it));
331
332        for global_arg in processed.global_state.borrow().global_args.iter() {
333            let address_space = global_arg.address_space();
334            // Skip tensor maps
335            if matches!(address_space, AddressSpace::Local) {
336                continue;
337            }
338            self.memories.insert(
339                global_arg.id(),
340                MemoryBlock {
341                    address_space,
342                    value_ty: global_arg.ty.unwrap_ptr(),
343                    alignment: global_arg.ty.align(),
344                    root_ptr: *global_arg,
345                },
346            );
347        }
348
349        let is_break = processed.instructions.contains(&Branch::Break.into());
350
351        for mut instruction in processed.instructions {
352            let mut removed = false;
353            for transform in state.transformers.iter() {
354                match transform.maybe_transform(&scope, &instruction) {
355                    TransformAction::Ignore => {}
356                    TransformAction::Replace(replacement) => {
357                        self.current_block_mut()
358                            .ops
359                            .borrow_mut()
360                            .extend(replacement);
361                        removed = true;
362                        break;
363                    }
364                    TransformAction::Remove => {
365                        removed = true;
366                        break;
367                    }
368                }
369            }
370            if removed {
371                continue;
372            }
373            match &mut instruction.operation {
374                Operation::DeclareVariable {
375                    value_ty,
376                    addr_space,
377                    alignment,
378                } => {
379                    let out = instruction.out.unwrap();
380                    self.memories.insert(
381                        out.id(),
382                        MemoryBlock {
383                            address_space: *addr_space,
384                            value_ty: *value_ty,
385                            alignment: *alignment,
386                            root_ptr: out,
387                        },
388                    );
389                    self.current_block_mut().ops.borrow_mut().push(instruction);
390                }
391                Operation::Branch(branch) => match self.parse_control_flow(state, branch.clone()) {
392                    ControlFlowAction::None => {}
393                    ControlFlowAction::AbortBlock => {
394                        break;
395                    }
396                },
397                _ => {
398                    self.current_block_mut().ops.borrow_mut().push(instruction);
399                }
400            }
401        }
402
403        is_break
404    }
405
406    /// Mutable reference to the current basic block
407    pub(crate) fn current_block_mut(&mut self) -> &mut BasicBlock {
408        let current_block = self.current_block.unwrap();
409        &mut self[current_block]
410    }
411
412    /// List of predecessor IDs of the `block`
413    pub fn predecessors(&self, block: NodeIndex) -> Vec<NodeIndex> {
414        self.edges_directed(block, Direction::Incoming)
415            .map(|it| it.source())
416            .filter(|it| !self.is_unreachable(*it))
417            .collect()
418    }
419
420    /// List of successor IDs of the `block`
421    pub fn successors(&self, block: NodeIndex) -> Vec<NodeIndex> {
422        self.edges_directed(block, Direction::Outgoing)
423            .map(|it| it.target())
424            .collect()
425    }
426
427    /// Return the breadth-first list of nodes along the dominator tree.
428    /// This is useful for generating the blocks in a human-readable-ish order that follows the
429    /// Vulkan spec (dominators before dominated).
430    pub fn breadth_first_dominators(&self) -> Vec<NodeIndex> {
431        self.analysis_cache
432            .try_get::<Dominators>()
433            .expect("Dominators should be present")
434            .breadth_first_nodes()
435    }
436
437    /// Reference to the [`BasicBlock`] with ID `block`
438    #[track_caller]
439    pub fn block(&self, block: NodeIndex) -> &BasicBlock {
440        &self[block]
441    }
442
443    /// Reference to the [`BasicBlock`] with ID `block`
444    #[track_caller]
445    pub fn block_mut(&mut self, block: NodeIndex) -> &mut BasicBlock {
446        &mut self[block]
447    }
448
449    pub fn is_unreachable(&self, block: NodeIndex) -> bool {
450        let control_flow = self[block].control_flow.borrow();
451        matches!(*control_flow, ControlFlow::Unreachable)
452    }
453
454    /// A set of node indices for all blocks in the program
455    pub fn node_ids(&self) -> Vec<NodeIndex> {
456        self.node_indices().collect()
457    }
458
459    fn transform_ssa_and_merge_composites(&mut self, state: &GlobalState) {
460        self.ssa_transform(state);
461
462        let mut done = false;
463        while !done {
464            let changes = AtomicCounter::new(0);
465            if changes.get() > 0 {
466                self.ssa_transform(state);
467            } else {
468                done = true;
469            }
470        }
471    }
472
473    fn ssa_transform(&mut self, state: &GlobalState) {
474        InlineCopies.apply_pre_ssa(self, state, AtomicCounter::new(0));
475        self.place_phi_nodes(state);
476        self.version_program(state);
477        self.invalidate_analysis::<LocalStores>();
478        self.invalidate_analysis::<DomFrontiers>();
479    }
480
481    /// Run all optimizations
482    fn run_opt(&mut self, state: &GlobalState, scope: Scope) {
483        self.parse_graph(state, scope);
484        self.split_critical_edges();
485
486        self.transform_ssa_and_merge_composites(state);
487        self.analysis::<PointerSource>(state);
488        self.apply_post_ssa_passes(state);
489
490        // Special expensive passes that should only run once.
491        // Need more optimization rounds in between.
492
493        // Disaggregate arrays, remove the resulting pointer copies, then re-run mem2reg on the new
494        // variables
495        let arrays_prop = AtomicCounter::new(0);
496        log::debug!("Applying {}", DisaggregateArray.name());
497        DisaggregateArray.apply_post_ssa(self, state, arrays_prop.clone());
498        if arrays_prop.get() > 0 {
499            InlineCopies.apply_post_ssa(self, state, AtomicCounter::new(0));
500            self.invalidate_analysis::<Liveness>();
501            self.transform_ssa_and_merge_composites(state);
502            self.apply_post_ssa_passes(state);
503        }
504
505        let gvn_count = AtomicCounter::new(0);
506        log::debug!("Applying {}", GvnPass.name());
507        GvnPass.apply_post_ssa(self, state, gvn_count.clone());
508        log::debug!("Applying {}", ReduceStrength.name());
509        ReduceStrength.apply_post_ssa(self, state, gvn_count.clone());
510        log::debug!("Applying {}", CopyTransform.name());
511        CopyTransform.apply_post_ssa(self, state, gvn_count.clone());
512
513        if gvn_count.get() > 0 {
514            self.apply_post_ssa_passes(state);
515        }
516
517        self.split_free();
518        self.analysis::<SharedLiveness>(state);
519
520        log::debug!("Applying {}", MergeBlocks.name());
521        MergeBlocks.apply_post_ssa(self, state, AtomicCounter::new(0));
522
523        log::debug!("Collecting captures");
524        let captures = self.analysis::<Captures>(state);
525        self.implicit_params = captures
526            .at_block(self.root)
527            .iter()
528            .copied()
529            .filter(|param| !self.explicit_params.contains(param))
530            .collect();
531
532        self.update_buffer_vis(state);
533        self.analysis::<Dominators>(state);
534    }
535
536    /// Run only the shared memory analysis
537    fn run_shared_only(&mut self, state: &GlobalState, scope: Scope) {
538        self.parse_graph(state, scope);
539        self.split_critical_edges();
540        self.transform_ssa_and_merge_composites(state);
541        self.split_free();
542        self.analysis::<PointerSource>(state);
543        self.analysis::<SharedLiveness>(state);
544        self.update_buffer_vis(state);
545    }
546
547    fn update_buffer_vis(&mut self, state: &GlobalState) {
548        self.visit_all(
549            state,
550            |_, val| {
551                if let Some(id) = global_buffer_id(val) {
552                    state.set_buffer_readable(id);
553                }
554            },
555            |_, val| {
556                if let Some(id) = global_buffer_id(val) {
557                    state.set_buffer_writable(id);
558                }
559            },
560        );
561    }
562
563    fn apply_post_ssa_passes(&mut self, state: &GlobalState) {
564        // Passes that run regardless of execution mode
565        let mut passes: Vec<Box<dyn OptimizerPass>> = vec![
566            Box::new(InlineCopies),
567            Box::new(EliminateUnusedVariables),
568            Box::new(ConstOperandSimplify),
569            Box::new(MergeSameExpressions),
570            Box::new(ConstEval),
571            Box::new(EliminateConstBranches),
572            Box::new(EmptyBranchToSelect),
573            Box::new(EliminateDeadBlocks),
574            Box::new(EliminateDeadPhi),
575        ];
576
577        log::debug!("Applying post-SSA passes");
578        loop {
579            let counter = AtomicCounter::default();
580            for pass in &mut passes {
581                log::debug!("Applying {}", pass.name());
582                pass.apply_post_ssa(self, state, counter.clone());
583            }
584
585            if counter.get() == 0 {
586                break;
587            }
588        }
589    }
590
591    pub(crate) fn ret(&mut self) -> NodeIndex {
592        if self[self.ret].block_use.contains(&BlockUse::Merge) {
593            let ret = self.ret;
594            let new_ret = self.add_node(BasicBlock::default());
595            self.add_edge(new_ret, ret, 0);
596            self.ret = new_ret;
597            self.invalidate_structure();
598            new_ret
599        } else {
600            self.ret
601        }
602    }
603
604    pub fn all_params(&self) -> impl Iterator<Item = Value> {
605        self.explicit_params
606            .iter()
607            .copied()
608            .chain(self.implicit_params.iter().copied())
609    }
610
611    pub fn create_local_mut(&mut self, state: &GlobalState, value_ty: Type) -> Value {
612        let ty = Type::Pointer(value_ty.intern(), AddressSpace::Local);
613        let val = state.allocator.create_value(ty);
614        let root = self.root;
615        self[root].ops.borrow_mut().push(Instruction::new(
616            Operation::DeclareVariable {
617                value_ty,
618                addr_space: AddressSpace::Local,
619                alignment: value_ty.align(),
620            },
621            val,
622        ));
623        self.memories.insert(
624            val.id(),
625            MemoryBlock {
626                address_space: AddressSpace::Local,
627                value_ty,
628                alignment: value_ty.align(),
629                root_ptr: val,
630            },
631        );
632        val
633    }
634}
635
636/// A visitor that does nothing.
637pub fn visit_noop(_opt: &mut Function, _var: &mut Value) {}
638
639#[cfg(test)]
640mod test {
641    use alloc::vec;
642    use cubecl_core as cubecl;
643    use cubecl_core::prelude::*;
644    use cubecl_ir::{ElemType, Type, UIntKind};
645
646    use crate::Optimizer;
647
648    #[allow(unused)]
649    #[cube(launch)]
650    fn pre_kernel(x: u32, cond: u32, out: &mut [u32]) {
651        let mut y = 0;
652        let mut z = 0;
653        if cond == 0 {
654            y = x + 4;
655        }
656        z = x + 4;
657        out[0] = y;
658        out[1] = z;
659    }
660
661    #[test_log::test]
662    #[ignore = "no good way to assert opt is applied"]
663    fn test_pre() {
664        let ctx = Scope::root(false);
665        let u32 = Type::scalar(ElemType::UInt(UIntKind::U32));
666        let x = ctx.create_value(u32).into();
667        let cond = ctx.create_value(u32).into();
668        let mut arr = ctx.global(0, u32).into();
669
670        pre_kernel::expand(&ctx, x, cond, &mut arr);
671        let opt = Optimizer::new(ctx, CubeDim::new_1d(1), vec![], vec![], vec![]);
672        std::println!("{opt}")
673    }
674}
cubecl_opt/lib.rs

cubecl_opt/
lib.rs
