cubecl_core/post_processing/
analysis_helper.rs

1use alloc::{rc::Rc, vec, vec::Vec};
2use core::cell::{RefCell, RefMut};
3use cubecl_runtime::kernel::Visibility;
4use derive_more::{Deref, DerefMut};
5
6use cubecl_ir::{
7    AddressSpace, GlobalState, Id, Instruction, Memory, Operation, Scope, Value, ValueKind,
8};
9use hashbrown::{HashMap, HashSet};
10
11use crate::post_processing::{
12    util::AtomicCounter,
13    visitor::{InstructionVisitor, Visitor},
14};
15
16#[derive(Debug, Clone, Default)]
17pub struct GlobalAnalyses {
18    ptr_source: Rc<RefCell<PointerSource>>,
19    used_values: Rc<RefCell<UsedValues>>,
20}
21
22impl GlobalAnalyses {
23    pub fn recalculate_pointer_source(&self, scope: &Scope) {
24        *self.ptr_source.borrow_mut() = PointerSource::new(scope);
25    }
26
27    pub fn recalculate_used_values(&self, scope: &Scope) {
28        let mut used_values = UsedValues::default();
29        used_values.visit_scope(scope, self, &AtomicCounter::new(0));
30        *self.used_values.borrow_mut() = used_values;
31    }
32
33    pub fn ptr_source(&self) -> RefMut<'_, PointerSource> {
34        self.ptr_source.borrow_mut()
35    }
36
37    pub fn used_values(&self) -> RefMut<'_, UsedValues> {
38        self.used_values.borrow_mut()
39    }
40}
41
42#[derive(Default, Debug, Deref, DerefMut)]
43pub struct UsedValues {
44    used: HashSet<Value>,
45}
46
47impl InstructionVisitor for UsedValues {
48    fn visit_instruction(
49        &mut self,
50        mut inst: Instruction,
51        _global_state: &GlobalState,
52        analyses: &GlobalAnalyses,
53        _changes: &AtomicCounter,
54    ) -> Vec<Instruction> {
55        let mut visitor = Visitor(self);
56        visitor.visit_operation(&mut inst.operation, analyses, |this, val| {
57            this.used.insert(*val);
58        });
59        vec![inst]
60    }
61}
62
63#[derive(Debug, Default, Deref, DerefMut)]
64pub struct PointerSource {
65    sources: HashMap<ValueKind, Value>,
66}
67
68impl PointerSource {
69    pub fn new(scope: &Scope) -> Self {
70        let mut this = PointerSource::default();
71        this.visit_scope(scope, &GlobalAnalyses::default(), &AtomicCounter::new(0));
72        this
73    }
74}
75
76impl InstructionVisitor for PointerSource {
77    fn visit_instruction(
78        &mut self,
79        inst: Instruction,
80        _global_state: &GlobalState,
81        _analyses: &GlobalAnalyses,
82        _changes: &AtomicCounter,
83    ) -> Vec<Instruction> {
84        match &inst.operation {
85            Operation::Copy(val) if val.ty.is_ptr() && inst.out().ty.is_ptr() => {
86                if let Some(source) = self.sources.get(&val.kind) {
87                    self.sources.insert(inst.out().kind, *source);
88                }
89            }
90            Operation::Memory(Memory::Index(index_operands)) => {
91                self.sources.insert(inst.out().kind, index_operands.list);
92            }
93            _ => {}
94        }
95        vec![inst]
96    }
97}
98
99#[derive(Debug, Deref, Default)]
100pub struct BufferVisibility {
101    buffers: Vec<Visibility>,
102}
103
104impl From<BufferVisibility> for Vec<Visibility> {
105    fn from(value: BufferVisibility) -> Self {
106        value.buffers
107    }
108}
109
110impl BufferVisibility {
111    pub fn new(scope: &Scope, analyses: &GlobalAnalyses) -> Self {
112        let mut this = BufferVisibility::default();
113        this.visit_scope(scope, analyses, &AtomicCounter::new(0));
114        this
115    }
116}
117
118impl InstructionVisitor for BufferVisibility {
119    fn visit_instruction(
120        &mut self,
121        mut inst: Instruction,
122        _global_state: &GlobalState,
123        analyses: &GlobalAnalyses,
124        _changes: &AtomicCounter,
125    ) -> Vec<Instruction> {
126        let mut visitor = Visitor(self);
127
128        visitor.visit_instruction(
129            &mut inst,
130            analyses,
131            |this, val| {
132                if let Some(id) = global_buffer_id(val) {
133                    this.set_readable(id as usize);
134                }
135            },
136            |this, val| {
137                if let Some(id) = global_buffer_id(val) {
138                    this.set_writable(id as usize);
139                }
140            },
141        );
142
143        vec![inst]
144    }
145}
146
147impl BufferVisibility {
148    // There for consistency and in case we make it more granular like it is in SPIR-V
149    fn set_readable(&mut self, id: usize) {
150        if self.buffers.len() <= id {
151            self.buffers.resize(id + 1, Visibility::Read);
152        }
153    }
154
155    fn set_writable(&mut self, id: usize) {
156        if self.buffers.len() <= id {
157            self.buffers.resize(id + 1, Visibility::Read);
158        }
159        self.buffers[id] = Visibility::ReadWrite;
160    }
161}
162
163fn global_buffer_id(variable: &Value) -> Option<Id> {
164    match variable.address_space() {
165        AddressSpace::Global(id) => Some(id),
166        _ => None,
167    }
168}
cubecl_core/post_processing/analysis_helper.rs

cubecl_core/post_processing/
analysis_helper.rs
