cubecl_opt/
transformers.rs

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

cubecl_opt/
transformers.rs
