Skip to main content

cubecl_core/post_processing/
util.rs

1use alloc::rc::Rc;
2use core::sync::atomic::{AtomicUsize, Ordering};
3
4/// An atomic counter with a simplified interface.
5#[derive(Clone, Debug, Default)]
6pub struct AtomicCounter {
7    inner: Rc<AtomicUsize>,
8}
9
10impl AtomicCounter {
11    /// Creates a new counter with `val` as its initial value.
12    pub fn new(val: usize) -> Self {
13        Self {
14            inner: Rc::new(AtomicUsize::new(val)),
15        }
16    }
17
18    /// Increments the counter and returns the last count.
19    pub fn inc(&self) -> usize {
20        self.inner.fetch_add(1, Ordering::SeqCst)
21    }
22
23    /// Gets the value of the counter without incrementing it.
24    pub fn get(&self) -> usize {
25        self.inner.load(Ordering::SeqCst)
26    }
27
28    pub fn get_and_reset(&self) -> usize {
29        self.inner.swap(0, Ordering::SeqCst)
30    }
31}