External Memory | NumKong for Rust

Views wrap raw pointers without ownership, owned containers accept custom allocators, and the scalar trait API works on any &[T] regardless of how the memory was allocated.

VectorView::from_raw_parts and TensorView::from_raw_parts wrap device-accessible or externally allocated memory. The mutable counterparts VectorSpan::from_raw_parts and TensorSpan::from_raw_parts work the same way with *mut T.

use numkong::{VectorView, TensorView};

let embeddings_ptr: *const f32 = /* from CUDA, mmap, or FFI */;
let embeddings = unsafe {
    VectorView::from_raw_parts(embeddings_ptr, 1024, std::mem::size_of::<f32>() as isize)
};

let shape = [32, 64];
let strides = [64 * 4, 4]; // row-major f32
let matrix = unsafe { TensorView::<f32>::from_raw_parts(embeddings_ptr, &shape, &strides) }.unwrap();

Owned containers accept any core::alloc::Allocator. A CUDA unified memory allocator looks like this:

use core::alloc::{AllocError, Allocator, Layout};
use core::ptr::NonNull;
use numkong::Vector;

struct CudaAllocator;

unsafe impl Allocator for CudaAllocator {
    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
        let raw = unsafe { cuda_malloc_managed(layout.size()) };
        let base = NonNull::new(raw).ok_or(AllocError)?;
        Ok(NonNull::slice_from_raw_parts(base, layout.size()))
    }
    unsafe fn deallocate(&self, block: NonNull<u8>, _layout: Layout) {
        cuda_free(block.as_ptr());
    }
}

let queries = Vector::<f32, CudaAllocator>::zeros_in(1024, CudaAllocator).unwrap();

core also implements the trait for &A, so an allocator that is not Clone — a bump arena holding a cursor, say — goes in by reference and the container borrows it:

let arena = BumpArena::new();
let queries = Vector::<f32, _>::zeros_in(1024, &arena).unwrap();
let keys = Vector::<f32, _>::zeros_in(1024, &arena).unwrap();

The trait-based scalar API works on any &[T] — Vec, mmap, arena, or pinned buffer:

use numkong::Dot;

let weights: &[f32] = /* any contiguous slice */;
let similarity = f32::dot(weights, weights).unwrap();

API

StructUnifiedAllocatorMemory the host and the devices of one capability group both address, for the packs and outputs of their verbs: managed memory on CUDA and ROCm, a shared buffer on Metal, and the heap on the CPU.2

capabilities.rs:804 ↗

pub struct UnifiedAllocator<'stream>

Every block lands on the device of the stream it borrows, and a block aligned less strictly than its layout asks is refused.

Examples

use numkong::{Capabilities, Stream, Tensor, UnifiedAllocator};

let capabilities = match Capabilities::cuda_count_devices()? {
    0 => Capabilities::cpu_enabled(),
    _ => Capabilities::cuda_enabled(0)?,
};
let stream = Stream::default(capabilities);
let ones = Tensor::<f32, UnifiedAllocator>::ones_in(&[4, 4], UnifiedAllocator::new(&stream))?;
stream.synchronize()?;
assert_eq!(ones.as_slice()[15], 1.0);

Implements

Clone, Copy, Debug, Allocator

newUnified memory on the device of stream, the host heap for a CPU stream.

capabilities.rs:810 ↗

pub const fn new(stream: &'stream Stream) -> Self

capabilitiesThe capabilities whose memory every block comes from.

capabilities.rs:813 ↗

pub fn capabilities(&self) -> Capabilities

StructGlobalThe default allocator, forwarding to the heap registered with #[global_allocator].

tensor.rs:154 ↗

pub struct Global;

The Allocator trait itself is core's, so an arena, a pool, or a pinned-memory allocator written against it — including by reference, through impl Allocator for &A — drops into Tensor and the packed containers with no adapter. This type only names the default, and stays separate from alloc::alloc::Global so the heapless build still has one.

Without the alloc feature there is no heap to forward to and every request fails. That is a supported configuration rather than a broken one: the borrowed API — the scalar traits over &[T], the *_into verbs writing into caller-provided spans, and the from_raw_parts views — never allocates, so only the owning containers become unusable.

Implements

Debug, Clone, Copy, Default, Allocator

Constants

DEFAULT_ALIGNMENTByte boundary every owned allocation and pack starts on, the C library's NUMKONG_DEFAULT_ALIGNMENT for the same target.

tensor.rs:108 ↗

pub const DEFAULT_ALIGNMENT: usize = 128;