cubecl-opt 0.11.0-pre.1

Compiler optimizations for CubeCL
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
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

use alloc::vec::Vec;
use cubecl_ir::{AddressSpace, CopyMemoryOperands, Instruction, Memory, Operation, Type, Value};
use hashbrown::{HashMap, HashSet};

use crate::{
    AtomicCounter, Function, GlobalState,
    analyses::{integer_range::val_id, pointer_source::PointerSource},
    visit_noop,
};

use super::OptimizerPass;

/// Transform consecutive `Index` and `IndexAssign` operations that don't modify the value into a
/// single `Copy` operation.
pub struct CopyTransform;

impl OptimizerPass for CopyTransform {
    fn apply_post_ssa(&mut self, func: &mut Function, state: &GlobalState, changes: AtomicCounter) {
        let ptr_source = func.analysis::<PointerSource>(state);
        for block in func.node_ids() {
            let mut reads = HashMap::new();
            let mut writes = HashMap::new();
            let ops = func[block].ops.clone();
            let indices = ops.borrow().indices().collect::<Vec<_>>();
            for idx in indices {
                let inst = ops.borrow()[idx].clone();
                match &inst.operation {
                    Operation::Memory(Memory::Load(ptr))
                        if can_copy(&ptr.ty)
                            && let Some(source) = ptr_source.get(ptr)
                            && source.value_ty == inst.ty()
                            && !is_reused(func, state, &inst.out) =>
                    {
                        if let Some(id) = val_id(&inst.out()) {
                            reads.insert(id, (idx, *ptr, source.root_ptr));
                        }
                    }
                    Operation::Memory(Memory::Store(op))
                        if can_copy(&op.ptr.ty)
                            && let Some(source) = ptr_source.get(&op.ptr) =>
                    {
                        if let Some(id) = val_id(&op.value) {
                            writes.insert(id, (idx, op.ptr, source.root_ptr));
                        }
                    }
                    _ => {}
                }
            }
            let read_ids: HashSet<_> = reads.keys().collect();
            let write_ids: HashSet<_> = writes.keys().collect();
            let copy_ids = read_ids.intersection(&write_ids);
            for id in copy_ids {
                let (read_idx, in_ptr, in_source) = reads[*id];
                let (write_idx, out_ptr, out_source) = writes[*id];
                let mut is_overwritten = false;
                for mut inst in
                    (read_idx..write_idx).filter_map(|idx| ops.borrow().get(idx).cloned())
                {
                    func.visit_instruction(state, &mut inst, visit_noop, |_, val| {
                        if *val == in_source || *val == out_source {
                            is_overwritten = true;
                        }
                    });
                }
                if is_overwritten {
                    continue;
                }

                ops.borrow_mut().remove(read_idx);
                let copy = Memory::CopyMemory(CopyMemoryOperands {
                    source: in_ptr,
                    target: out_ptr,
                    len: 1,
                });
                ops.borrow_mut()[write_idx] = Instruction::no_out(copy);
                changes.inc();
            }
        }
    }
}

/// Copy is only implemented in SPIR-V for global -> shared or shared -> global. So no point merging
/// locals.
fn can_copy(ty: &Type) -> bool {
    matches!(
        ty.address_space(),
        Some(AddressSpace::Global(..)) | Some(AddressSpace::Shared)
    )
}

fn is_reused(func: &mut Function, state: &GlobalState, val: &Option<Value>) -> bool {
    if let Some(val) = val.as_ref() {
        let count = AtomicCounter::new(0);
        func.visit_all(
            state,
            |_, other| {
                if other == val {
                    count.inc();
                }
            },
            |_, _| {},
        );
        count.get() > 1
    } else {
        false
    }
}