cubecl_ir/
arena.rs

1use alloc::vec::Vec;
2
3use bumpalo::Bump;
4use cubecl_macros_internal::TypeHash;
5
6/// Bump allocator that runs drop on non-trivially droppable values on reset.
7/// This is useful to prevent memory leaks without drastically affecting performance for the vast
8/// majority of values that do not need custom drops.
9///
10/// # SAFETY
11///
12/// Values inserted *must not* access external borrowed data in a custom `Drop` implementation.
13/// This shouldn't happen for normal types, but should be kept in mind regardless.
14/// Contrived Example:
15///
16/// ```no_run
17/// # use cubecl_ir::arena::DropBump;
18/// struct Wrapper<'a>(&'a mut String);
19///
20/// impl<'a> Drop for Wrapper<'a> {
21///     fn drop(&mut self) {
22///         self.0.push_str(" dropped"); // Accesses external borrow
23///     }
24/// }
25///
26/// let mut external = String::from("hello");
27///
28/// let mut pool = DropBump::new();
29/// let w = pool.alloc(Wrapper(&mut external));
30/// drop(external);
31/// pool.reset(); // Runs `drop_in_place::<Wrapper>`
32/// ```
33///
34/// So don't do that.
35///
36#[derive(Default, Debug, TypeHash)]
37pub struct DropBump {
38    bump: Bump,
39    drop_thunks: Vec<DropThunk>,
40}
41
42unsafe impl Send for DropBump {}
43
44#[derive(Debug, TypeHash)]
45struct DropThunk {
46    func: fn(*mut ()),
47    ptr: *mut (),
48}
49
50impl DropBump {
51    pub fn new() -> Self {
52        Default::default()
53    }
54
55    pub fn reset(&mut self) {
56        for drop in self.drop_thunks.drain(..).rev() {
57            (drop.func)(drop.ptr)
58        }
59        self.bump.reset();
60    }
61
62    pub fn alloc<T>(&mut self, val: T) -> &mut T {
63        let ret = self.bump.alloc(val);
64
65        if core::mem::needs_drop::<T>() {
66            self.drop_thunks.push(DropThunk {
67                func: drop_glue::<T>,
68                ptr: ret as *mut T as *mut (),
69            });
70        }
71
72        ret
73    }
74}
75
76fn drop_glue<T>(ptr: *mut ()) {
77    unsafe { core::ptr::drop_in_place::<T>(ptr as *mut T) };
78}
cubecl_ir/arena.rs

cubecl_ir/
arena.rs
