use std::collections::HashMap;
use crate::error::{CudaError, CudaResult};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct StreamOrderId(pub u64);
impl StreamOrderId {
pub const NULL: StreamOrderId = StreamOrderId(0);
#[inline]
pub fn raw(self) -> u64 {
self.0
}
}
impl From<u64> for StreamOrderId {
#[inline]
fn from(value: u64) -> Self {
StreamOrderId(value)
}
}
#[derive(Debug, Clone, Copy, Default)]
struct StreamClock {
submit: u64,
reached: u64,
}
impl StreamClock {
fn enqueue(&mut self) -> u64 {
let seq = self.submit;
self.submit = self.submit.saturating_add(1);
seq
}
fn advance_to_head(&mut self) {
self.reached = self.submit;
}
fn has_reached(&self, seq: u64) -> bool {
self.reached > seq
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct Block {
ptr: u64,
capacity: usize,
}
#[derive(Debug, Clone, Copy)]
struct PendingFree {
block: Block,
stream: StreamOrderId,
seq: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ModelAllocation {
pub ptr: u64,
pub size: usize,
pub capacity: usize,
pub stream: StreamOrderId,
pub ready_seq: u64,
}
#[derive(Debug, Clone, Copy)]
pub struct ModelLimits {
pub max_pool_size: usize,
pub release_threshold: usize,
}
#[derive(Debug)]
pub struct StreamOrderModel {
clocks: HashMap<StreamOrderId, StreamClock>,
free_list: Vec<Block>,
pending_frees: Vec<PendingFree>,
live: HashMap<u64, Block>,
next_addr: u64,
reserved: usize,
used: usize,
reserved_high: usize,
used_high: usize,
active: usize,
peak_active: usize,
limits: ModelLimits,
}
const VIRTUAL_BASE: u64 = 0x0000_7F00_0000_0000;
const GRANULARITY: usize = 512;
impl StreamOrderModel {
pub fn new(limits: ModelLimits) -> Self {
Self {
clocks: HashMap::new(),
free_list: Vec::new(),
pending_frees: Vec::new(),
live: HashMap::new(),
next_addr: VIRTUAL_BASE,
reserved: 0,
used: 0,
reserved_high: 0,
used_high: 0,
active: 0,
peak_active: 0,
limits,
}
}
fn align(size: usize) -> usize {
size.saturating_add(GRANULARITY - 1) / GRANULARITY * GRANULARITY
}
pub fn set_release_threshold(&mut self, threshold: usize) {
self.limits.release_threshold = threshold;
}
pub fn alloc(&mut self, size: usize, stream: StreamOrderId) -> CudaResult<ModelAllocation> {
if size == 0 {
return Err(CudaError::InvalidValue);
}
self.collect_ready_frees();
let want = Self::align(size);
let block = if let Some(idx) = self.pick_free_block(want) {
self.free_list.swap_remove(idx)
} else {
if self.limits.max_pool_size > 0
&& self.reserved.saturating_add(want) > self.limits.max_pool_size
{
return Err(CudaError::OutOfMemory);
}
let ptr = self.next_addr;
self.next_addr = self.next_addr.saturating_add(want as u64);
self.reserved = self.reserved.saturating_add(want);
if self.reserved > self.reserved_high {
self.reserved_high = self.reserved;
}
Block {
ptr,
capacity: want,
}
};
let ready_seq = self.clock_mut(stream).enqueue();
self.live.insert(block.ptr, block);
self.used = self.used.saturating_add(block.capacity);
if self.used > self.used_high {
self.used_high = self.used;
}
self.active = self.active.saturating_add(1);
if self.active > self.peak_active {
self.peak_active = self.active;
}
Ok(ModelAllocation {
ptr: block.ptr,
size,
capacity: block.capacity,
stream,
ready_seq,
})
}
fn pick_free_block(&self, want: usize) -> Option<usize> {
let mut first_fit: Option<usize> = None;
for (idx, block) in self.free_list.iter().enumerate() {
if block.capacity == want {
return Some(idx);
}
if block.capacity > want && first_fit.is_none() {
first_fit = Some(idx);
}
}
first_fit
}
pub fn free(&mut self, ptr: u64, stream: StreamOrderId) -> CudaResult<()> {
let block = self.live.remove(&ptr).ok_or(CudaError::InvalidValue)?;
self.used = self.used.saturating_sub(block.capacity);
self.active = self.active.saturating_sub(1);
let seq = self.clock_mut(stream).enqueue();
self.pending_frees.push(PendingFree { block, stream, seq });
self.collect_ready_frees();
Ok(())
}
pub fn is_live(&self, ptr: u64) -> bool {
self.live.contains_key(&ptr)
}
pub fn synchronize(&mut self, stream: StreamOrderId) {
self.clock_mut(stream).advance_to_head();
self.collect_ready_frees();
}
pub fn is_ready_same_stream(&self, alloc: &ModelAllocation) -> bool {
self.clocks
.get(&alloc.stream)
.is_some_and(|c| c.has_reached(alloc.ready_seq))
}
pub fn is_ready_cross_stream(
&self,
alloc: &ModelAllocation,
consumer: StreamOrderId,
wait_seq: u64,
) -> bool {
if consumer == alloc.stream {
return self.is_ready_same_stream(alloc);
}
wait_seq > alloc.ready_seq
}
pub fn record_event(&mut self, stream: StreamOrderId) -> u64 {
self.clock_mut(stream).submit
}
fn collect_ready_frees(&mut self) {
let mut still_pending = Vec::with_capacity(self.pending_frees.len());
let drained = std::mem::take(&mut self.pending_frees);
for pf in drained {
let reached = self
.clocks
.get(&pf.stream)
.is_some_and(|c| c.has_reached(pf.seq));
if reached {
self.free_list.push(pf.block);
} else {
still_pending.push(pf);
}
}
self.pending_frees = still_pending;
self.release_excess();
}
fn release_excess(&mut self) {
let pending_bytes: usize = self.pending_frees.iter().map(|p| p.block.capacity).sum();
let pinned = self.used.saturating_add(pending_bytes);
let keep_floor = self.limits.release_threshold.max(pinned);
while self.reserved > keep_floor {
let Some(block) = self.free_list.pop() else {
break;
};
self.reserved = self.reserved.saturating_sub(block.capacity);
}
}
pub fn trim_to(&mut self, min_bytes_to_keep: usize) {
self.collect_ready_frees();
let pending_bytes: usize = self.pending_frees.iter().map(|p| p.block.capacity).sum();
let pinned = self.used.saturating_add(pending_bytes);
let keep_floor = min_bytes_to_keep.max(pinned);
while self.reserved > keep_floor {
let Some(block) = self.free_list.pop() else {
break;
};
self.reserved = self.reserved.saturating_sub(block.capacity);
}
}
pub fn reset_peaks(&mut self) {
self.reserved_high = self.reserved;
self.used_high = self.used;
self.peak_active = self.active;
}
#[inline]
pub fn reserved(&self) -> usize {
self.reserved
}
#[inline]
pub fn used(&self) -> usize {
self.used
}
#[inline]
pub fn reserved_high(&self) -> usize {
self.reserved_high
}
#[inline]
pub fn used_high(&self) -> usize {
self.used_high
}
#[inline]
pub fn active(&self) -> usize {
self.active
}
#[inline]
pub fn peak_active(&self) -> usize {
self.peak_active
}
#[inline]
pub fn free_block_count(&self) -> usize {
self.free_list.len()
}
#[inline]
pub fn pending_free_count(&self) -> usize {
self.pending_frees.len()
}
fn clock_mut(&mut self, stream: StreamOrderId) -> &mut StreamClock {
self.clocks.entry(stream).or_default()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn limits(max: usize, release: usize) -> ModelLimits {
ModelLimits {
max_pool_size: max,
release_threshold: release,
}
}
#[test]
fn alloc_rounds_up_to_granularity() {
let mut m = StreamOrderModel::new(limits(0, 0));
let a = m.alloc(1, StreamOrderId::NULL).expect("alloc");
assert_eq!(a.size, 1);
assert_eq!(a.capacity, GRANULARITY);
assert_eq!(m.reserved(), GRANULARITY);
assert_eq!(m.used(), GRANULARITY);
assert_ne!(a.ptr, 0);
}
#[test]
fn zero_size_rejected() {
let mut m = StreamOrderModel::new(limits(0, 0));
assert_eq!(
m.alloc(0, StreamOrderId::NULL),
Err(CudaError::InvalidValue)
);
}
#[test]
fn free_of_foreign_pointer_rejected() {
let mut m = StreamOrderModel::new(limits(0, 0));
assert_eq!(
m.free(0xDEAD_BEEF, StreamOrderId::NULL),
Err(CudaError::InvalidValue)
);
}
#[test]
fn double_free_rejected() {
let mut m = StreamOrderModel::new(limits(0, 0));
let s = StreamOrderId(7);
let a = m.alloc(256, s).expect("alloc");
assert!(m.free(a.ptr, s).is_ok());
assert_eq!(m.free(a.ptr, s), Err(CudaError::InvalidValue));
}
#[test]
fn pending_free_holds_block_until_stream_advances() {
let mut m = StreamOrderModel::new(limits(0, usize::MAX));
let s = StreamOrderId(1);
let a = m.alloc(512, s).expect("alloc");
let _b = m.alloc(512, s).expect("alloc2");
m.free(a.ptr, s).expect("free");
assert_eq!(m.pending_free_count(), 1);
assert_eq!(m.free_block_count(), 0);
m.synchronize(s);
assert_eq!(m.pending_free_count(), 0);
assert_eq!(m.free_block_count(), 1);
}
#[test]
fn freed_block_is_reused_by_same_size_alloc() {
let mut m = StreamOrderModel::new(limits(0, usize::MAX));
let s = StreamOrderId(2);
let a = m.alloc(1024, s).expect("alloc");
let ptr_a = a.ptr;
let reserved_after_first = m.reserved();
m.free(a.ptr, s).expect("free");
m.synchronize(s); let b = m.alloc(1024, s).expect("alloc reuse");
assert_eq!(b.ptr, ptr_a, "freed block should be reused");
assert_eq!(m.reserved(), reserved_after_first, "no new reservation");
assert_eq!(m.free_block_count(), 0);
}
#[test]
fn same_stream_visibility_rule() {
let mut m = StreamOrderModel::new(limits(0, usize::MAX));
let s = StreamOrderId(3);
let a = m.alloc(64, s).expect("alloc");
assert!(!m.is_ready_same_stream(&a));
m.synchronize(s);
assert!(m.is_ready_same_stream(&a));
}
#[test]
fn cross_stream_requires_event_after_alloc() {
let mut m = StreamOrderModel::new(limits(0, usize::MAX));
let producer = StreamOrderId(10);
let consumer = StreamOrderId(20);
let a = m.alloc(128, producer).expect("alloc");
let early_wait = 0u64;
assert!(!m.is_ready_cross_stream(&a, consumer, early_wait));
let late_wait = m.record_event(producer);
assert!(late_wait > a.ready_seq);
assert!(m.is_ready_cross_stream(&a, consumer, late_wait));
}
#[test]
fn reserved_vs_used_accounting_consistent() {
let mut m = StreamOrderModel::new(limits(0, usize::MAX));
let s = StreamOrderId(4);
let a = m.alloc(1000, s).expect("a");
let b = m.alloc(2000, s).expect("b");
let total = a.capacity + b.capacity;
assert_eq!(m.used(), total);
assert_eq!(m.reserved(), total);
m.free(a.ptr, s).expect("free a");
assert_eq!(m.used(), b.capacity);
assert_eq!(m.reserved(), total);
m.synchronize(s);
assert_eq!(m.reserved(), total);
assert_eq!(m.free_block_count(), 1);
}
#[test]
fn trim_releases_free_blocks_above_threshold() {
let mut m = StreamOrderModel::new(limits(0, usize::MAX));
let s = StreamOrderId(5);
let a = m.alloc(4096, s).expect("a");
let reserved_full = m.reserved();
m.free(a.ptr, s).expect("free");
m.synchronize(s);
assert_eq!(m.free_block_count(), 1);
assert_eq!(m.reserved(), reserved_full);
m.trim_to(0);
assert_eq!(m.free_block_count(), 0);
assert_eq!(m.reserved(), 0);
}
#[test]
fn release_threshold_keeps_some_reserved() {
let mut m = StreamOrderModel::new(limits(0, 4096));
let s = StreamOrderId(6);
let a = m.alloc(4096, s).expect("a");
m.free(a.ptr, s).expect("free");
m.synchronize(s);
assert_eq!(m.reserved(), 4096);
assert_eq!(m.free_block_count(), 1);
}
#[test]
fn max_pool_size_enforced_on_fresh_blocks() {
let mut m = StreamOrderModel::new(limits(1024, 0));
let s = StreamOrderId::NULL;
assert!(m.alloc(1024, s).is_ok());
assert_eq!(m.alloc(1, s), Err(CudaError::OutOfMemory));
}
#[test]
fn larger_request_reuses_oversized_free_block() {
let mut m = StreamOrderModel::new(limits(0, usize::MAX));
let s = StreamOrderId(8);
let big = m.alloc(2048, s).expect("big");
let big_ptr = big.ptr;
m.free(big.ptr, s).expect("free big");
m.synchronize(s);
let small = m.alloc(1024, s).expect("small reuse");
assert_eq!(small.ptr, big_ptr, "oversized block reused");
assert_eq!(small.capacity, 2048, "capacity retained from block");
}
#[test]
fn peak_stats_track_and_reset() {
let mut m = StreamOrderModel::new(limits(0, usize::MAX));
let s = StreamOrderId(9);
let a = m.alloc(1024, s).expect("a");
let _b = m.alloc(2048, s).expect("b");
assert_eq!(m.peak_active(), 2);
assert_eq!(m.used_high(), a.capacity + 2048);
m.free(a.ptr, s).expect("free a");
m.synchronize(s);
m.reset_peaks();
assert_eq!(m.peak_active(), 1);
assert_eq!(m.used_high(), m.used());
}
}