cubecl_opt/
transformers.rs

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

cubecl_opt/
transformers.rs
