cubecl-core 0.11.0-pre.1

CubeCL core create
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

use alloc::{vec, vec::Vec};
use cubecl_ir::{CoopMma, GlobalState, Instruction, Operation, Operator, UnaryOperands, Value};
use hashbrown::HashMap;

use crate::post_processing::{
    analysis_helper::GlobalAnalyses,
    util::AtomicCounter,
    visitor::{InstructionVisitor, Visitor},
};

/// Inline constants or simple reassignments that don't change the type. This simplifies the code
/// and makes it easier to find optimizable expressions.
#[derive(Default, Debug)]
pub struct InlineAssignments {
    substitutions: HashMap<Value, Value>,
}

impl InstructionVisitor for InlineAssignments {
    fn visit_instruction(
        &mut self,
        mut inst: Instruction,
        _global_state: &GlobalState,
        analyses: &GlobalAnalyses,
        changes: &AtomicCounter,
    ) -> Vec<Instruction> {
        let mut visitor = Visitor(());
        visitor.visit_operation(&mut inst.operation, analyses, |_, val| {
            if let Some(substitution) = self.substitutions.get(val) {
                *val = *substitution;
                changes.inc();
            }
        });

        match &mut inst.operation {
            Operation::Copy(input)
            | Operation::Operator(Operator::Cast(UnaryOperands { input }))
            | Operation::Operator(Operator::Reinterpret(UnaryOperands { input }))
            | Operation::CoopMma(CoopMma::Cast { input })
                if (input.is_immutable() || input.is_array_like() || input.ty.is_ptr())
                    && (inst.out.unwrap().is_immutable()
                        || inst.out.unwrap().is_array_like()
                        || inst.out.unwrap().ty.is_ptr())
                    && input.ty == inst.out.unwrap().ty =>
            {
                self.substitutions.insert(inst.out.unwrap(), *input);
                changes.inc();
                vec![]
            }
            _ => vec![inst],
        }
    }
}