cubecl_opt/
transformers.rs

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

cubecl_opt/
transformers.rs
