cubecl_ir/
allocator.rs

1use alloc::rc::Rc;
2
3use cubecl_macros_internal::TypeHash;
4use portable_atomic::{AtomicU32, Ordering};
5
6use super::{Type, Value};
7
8/// An allocator for the [static single-assignment](https://en.wikipedia.org/wiki/Static_single-assignment_form)
9/// values of a kernel. For mutable variables, the value is the *root pointer* that serves as a
10/// handle into the place where the inner value is stored (`memref`, `Place`, `lvalue`, whatever your
11/// compiler of choice uses as a name). This means all values are immutable, and only the memory
12/// referenced by them may be mutated.
13#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
14#[derive(Clone, Debug, Default, TypeHash)]
15pub struct Allocator {
16    next_id: Rc<AtomicU32>,
17}
18
19impl PartialEq for Allocator {
20    fn eq(&self, other: &Self) -> bool {
21        Rc::ptr_eq(&self.next_id, &other.next_id)
22    }
23}
24impl Eq for Allocator {}
25
26impl Allocator {
27    pub fn clone_deep(&self) -> Self {
28        Allocator {
29            next_id: Rc::new(AtomicU32::new(self.next_id.load(Ordering::SeqCst))),
30        }
31    }
32
33    /// Create a new immutable value of type specified by `ty`.
34    pub fn create_value(&self, ty: Type) -> Value {
35        let id = self.new_local_index();
36        Value::new(id, ty)
37    }
38
39    pub fn new_local_index(&self) -> u32 {
40        self.next_id.fetch_add(1, Ordering::Release) + 1
41    }
42
43    pub fn current_local_index(&self) -> u32 {
44        self.next_id.load(Ordering::SeqCst)
45    }
46}
cubecl_ir/allocator.rs

cubecl_ir/
allocator.rs
