cubecl_opt/
transformers.rs

1use alloc::{boxed::Box, rc::Rc, vec::Vec};
2use cubecl_core::{CubeDim, post_processing::visitor::InstructionVisitor};
3use cubecl_ir::{Instruction, Processor, Scope};
4
5use crate::Optimizer;
6
7/// Build an optimizer with IR transformers
8#[derive(Default)]
9pub struct OptimizerBuilder {
10    transformers: Vec<Rc<dyn IrTransformer>>,
11    visitors: Vec<Box<dyn InstructionVisitor>>,
12    processors: Vec<Box<dyn Processor>>,
13}
14
15impl OptimizerBuilder {
16    /// Add an IR transformer to the optimizer
17    pub fn with_transformer(mut self, transformer: impl IrTransformer + 'static) -> Self {
18        self.transformers.push(Rc::new(transformer));
19        self
20    }
21
22    pub fn with_visitor(mut self, visitor: impl InstructionVisitor + 'static) -> Self {
23        self.visitors.push(Box::new(visitor));
24        self
25    }
26
27    pub fn with_processor(mut self, processor: impl Processor + 'static) -> Self {
28        self.processors.push(Box::new(processor));
29        self
30    }
31
32    /// Build and run optimizer on the scope
33    pub fn optimize(self, expand: Scope, cube_dim: CubeDim) -> Optimizer {
34        Optimizer::new(
35            expand,
36            cube_dim,
37            self.transformers,
38            self.visitors,
39            self.processors,
40        )
41    }
42}
43
44/// The action that should be performed on an instruction, returned by [`IrTransformer::maybe_transform`]
45pub enum TransformAction {
46    /// The transformer doesn't apply to this instruction
47    Ignore,
48    /// Replace this instruction with one or more other instructions
49    Replace(Vec<Instruction>),
50    /// Remove this instruction with no substitute (i.e. debug info)
51    Remove,
52}
53
54/// A transformer that can modify instructions before they get added to the control flow graph.
55pub trait IrTransformer: core::fmt::Debug {
56    /// Inspect an instruction and potentially transform it.
57    fn maybe_transform(&self, scope: &Scope, inst: &Instruction) -> TransformAction;
58}
cubecl_opt/transformers.rs

cubecl_opt/
transformers.rs
