cubecl_core/post_processing/
expression_merge.rs

1use alloc::{vec, vec::Vec};
2use cubecl_ir::{CoopMma, GlobalState, Instruction, Operation, Operator, UnaryOperands, Value};
3use hashbrown::HashMap;
4
5use crate::post_processing::{
6    analysis_helper::GlobalAnalyses,
7    util::AtomicCounter,
8    visitor::{InstructionVisitor, Visitor},
9};
10
11/// Inline constants or simple reassignments that don't change the type. This simplifies the code
12/// and makes it easier to find optimizable expressions.
13#[derive(Default, Debug)]
14pub struct InlineAssignments {
15    substitutions: HashMap<Value, Value>,
16}
17
18impl InstructionVisitor for InlineAssignments {
19    fn visit_instruction(
20        &mut self,
21        mut inst: Instruction,
22        _global_state: &GlobalState,
23        analyses: &GlobalAnalyses,
24        changes: &AtomicCounter,
25    ) -> Vec<Instruction> {
26        let mut visitor = Visitor(());
27        visitor.visit_operation(&mut inst.operation, analyses, |_, val| {
28            if let Some(substitution) = self.substitutions.get(val) {
29                *val = *substitution;
30                changes.inc();
31            }
32        });
33
34        match &mut inst.operation {
35            Operation::Copy(input)
36            | Operation::Operator(Operator::Cast(UnaryOperands { input }))
37            | Operation::Operator(Operator::Reinterpret(UnaryOperands { input }))
38            | Operation::CoopMma(CoopMma::Cast { input })
39                if (input.is_immutable() || input.is_array_like() || input.ty.is_ptr())
40                    && (inst.out.unwrap().is_immutable()
41                        || inst.out.unwrap().is_array_like()
42                        || inst.out.unwrap().ty.is_ptr())
43                    && input.ty == inst.out.unwrap().ty =>
44            {
45                self.substitutions.insert(inst.out.unwrap(), *input);
46                changes.inc();
47                vec![]
48            }
49            _ => vec![inst],
50        }
51    }
52}
cubecl_core/post_processing/expression_merge.rs

cubecl_core/post_processing/
expression_merge.rs
