Skip to main content

ironaccelerator_core/
memory.rs

1//! Vendor-neutral memory abstractions. Backends provide concrete `Allocation`
2//! and `Stream` types; this module defines the trait surface and the
3//! categories of memory a driver can hand out.
4
5use crate::error::Result;
6use core::ptr::NonNull;
7
8#[cfg(not(feature = "std"))]
9use alloc::boxed::Box;
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
12#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
13#[non_exhaustive]
14pub enum MemoryKind {
15    /// Plain device-local memory.
16    Device,
17    /// Host-allocated, page-locked, zero-copy mappable.
18    Pinned,
19    /// Unified / managed memory.
20    Unified,
21    /// Constant memory (CUDA `__constant__`, equivalent on others).
22    Constant,
23    /// Shared / scratchpad (per-block / per-workgroup).
24    Shared,
25    /// Texture / surface memory.
26    Texture,
27}
28
29/// A handle to backend-allocated memory. Concrete backends provide a
30/// non-trait struct that derefs to a typed slice for the hot path.
31pub trait Allocation: Send + Sync {
32    fn kind(&self) -> MemoryKind;
33    fn len_bytes(&self) -> usize;
34    /// Raw device pointer. Backends that virtualise pointers (Metal, QNN)
35    /// return their handle cast to `*mut u8`.
36    fn as_ptr(&self) -> NonNull<u8>;
37}
38
39/// Pluggable allocator. Concrete backends register one default pool per
40/// device (e.g. CUDA stream-ordered async allocator) plus optional pools for
41/// pinned / unified memory.
42pub trait MemoryPool: Send + Sync {
43    fn kind(&self) -> MemoryKind;
44    /// Allocate `bytes` aligned to `align`. Speed-over-safety: the returned
45    /// allocation is *uninitialised*.
46    fn alloc(&self, bytes: usize, align: usize) -> Result<Box<dyn Allocation>>;
47    /// Optional: hint a deallocation watermark. No-op for backends without it.
48    fn trim(&self) {}
49}