metrale_storage/
expert_arena.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `ExpertArena`: one pinned host allocation (`cuMemAllocHost`)
4//! divided into `num_slabs` slabs of `slots_per_slab` expert-record slots.
5//! `new` requires the allocation's device address to equal its host address,
6//! so a record written into a slot is read by device code at `slot_dev_va`
7//! without a copy. The expert tiers decide what fills a slot.
8//!
9//! Owner: metrale-storage experts.
10//! Invariants:
11//! - `dev_base` equals the allocation's host address.
12//! - `record_stride` is a non-zero multiple of 4096.
13
14use anyhow::{Context, Result, bail};
15use std::ffi::c_void;
16
17use crate::cuda_min::PinnedBuffer;
18
19/// 2026-09-25: Slabs of slots, each slot one expert record of `record_stride`
20/// bytes.
21pub struct ExpertArena {
22    pinned: PinnedBuffer,
23    /// 2026-09-25: Device address of the arena base; `new` checks that it
24    /// equals the host address.
25    dev_base: u64,
26    num_slabs: u32,
27    slots_per_slab: u32,
28    record_stride: usize,
29}
30
31impl ExpertArena {
32    /// 2026-09-25: Allocate `num_slabs * slots_per_slab * record_stride` pinned
33    /// bytes. Fails on a zero dimension, a stride that is not a multiple of
34    /// 4096, a size overflow, or a device address that differs from the host
35    /// address.
36    pub fn new(num_slabs: u32, slots_per_slab: u32, record_stride: usize) -> Result<Self> {
37        if num_slabs == 0 || slots_per_slab == 0 || record_stride == 0 {
38            bail!("ExpertArena: zero geometry ({num_slabs},{slots_per_slab},{record_stride})");
39        }
40        if !record_stride.is_multiple_of(4096) {
41            bail!("ExpertArena: record_stride {record_stride} must be a 4 KiB multiple (O_DIRECT)");
42        }
43        let total = (num_slabs as usize)
44            .checked_mul(slots_per_slab as usize)
45            .and_then(|v| v.checked_mul(record_stride))
46            .context("ExpertArena: size overflow")?;
47        let pinned = PinnedBuffer::new(total)?;
48        let dev_base = pinned.device_ptr()?;
49        let host_base = pinned.ptr as u64;
50        if dev_base != host_base {
51            bail!(
52                "ExpertArena: pinned host VA {host_base:#x} != device VA {dev_base:#x} \
53                 — host is not unified-addressing (UMA zero-copy unavailable)"
54            );
55        }
56        Ok(Self {
57            pinned,
58            dev_base,
59            num_slabs,
60            slots_per_slab,
61            record_stride,
62        })
63    }
64
65    pub fn num_slabs(&self) -> u32 {
66        self.num_slabs
67    }
68    pub fn slots_per_slab(&self) -> u32 {
69        self.slots_per_slab
70    }
71    pub fn record_stride(&self) -> usize {
72        self.record_stride
73    }
74
75    fn linear_slot(&self, slab: u32, slot: u32) -> Result<usize> {
76        if slab >= self.num_slabs || slot >= self.slots_per_slab {
77            bail!(
78                "ExpertArena: slot ({slab},{slot}) out of range ({},{})",
79                self.num_slabs,
80                self.slots_per_slab
81            );
82        }
83        Ok((slab as usize) * (self.slots_per_slab as usize) + (slot as usize))
84    }
85
86    /// 2026-09-25: Host pointer of a slot, where its record is written.
87    pub fn slot_host_ptr(&self, slab: u32, slot: u32) -> Result<*mut u8> {
88        let i = self.linear_slot(slab, slot)?;
89        // 2026-09-25: SAFETY: `linear_slot` checked i < num_slabs *
90        // slots_per_slab, so the offset is inside the allocation.
91        Ok(unsafe { (self.pinned.ptr as *mut u8).add(i * self.record_stride) })
92    }
93
94    /// 2026-09-25: Device address of a slot.
95    pub fn slot_dev_va(&self, slab: u32, slot: u32) -> Result<u64> {
96        let i = self.linear_slot(slab, slot)?;
97        Ok(self.dev_base + (i as u64) * (self.record_stride as u64))
98    }
99
100    /// 2026-09-25: Base of the pinned allocation; with `total_bytes` it
101    /// describes the whole arena.
102    pub fn base_ptr(&self) -> *mut c_void {
103        self.pinned.ptr
104    }
105    pub fn total_bytes(&self) -> usize {
106        self.pinned.bytes
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113
114    // 2026-09-25: `ExpertArena::new` allocates pinned CUDA memory, so the test
115    // needs a GPU.
116    #[test]
117    #[ignore = "requires GPU"]
118    fn arena_slots_are_strided_and_same_va() {
119        let _ctx = crate::cuda_min::CudaCtx::new(0).unwrap();
120        let stride = 8192;
121        let arena = ExpertArena::new(2, 3, stride).unwrap();
122        assert_eq!(arena.total_bytes(), 2 * 3 * stride);
123        let a = arena.slot_dev_va(0, 0).unwrap();
124        let b = arena.slot_dev_va(0, 1).unwrap();
125        let c = arena.slot_dev_va(1, 0).unwrap();
126        assert_eq!(b - a, stride as u64);
127        assert_eq!(c - a, 3 * stride as u64);
128        assert_eq!(
129            arena.slot_host_ptr(1, 2).unwrap() as u64,
130            arena.slot_dev_va(1, 2).unwrap()
131        );
132        assert!(arena.slot_dev_va(2, 0).is_err());
133    }
134}
metrale_storage/expert_arena.rs

metrale_storage/
expert_arena.rs
