Chapter 6: Memory Allocators & SIMD

Part 1: Custom Memory Allocators

Understanding the Global Allocator API

Rust's global allocator allows you to replace the default memory allocator with a custom implementation. This is crucial for embedded systems, real-time applications, and performance-critical code.

ℹ️ Why Custom Allocators?

Reduce fragmentation
Improve cache locality
Deterministic allocation times
Memory-constrained environments
Track memory usage

Simple Bump Allocator

📝 Example: Bump Allocator

use core::alloc::{GlobalAlloc, Layout};
use core::cell::UnsafeCell;
use core::ptr;

pub struct BumpAllocator {
    heap: UnsafeCell<[u8; 64 * 1024]>, // 64 KB heap
    next: UnsafeCell,
}

unsafe impl Sync for BumpAllocator {}

impl BumpAllocator {
    pub const fn new() -> Self {
        BumpAllocator {
            heap: UnsafeCell::new([0; 64 * 1024]),
            next: UnsafeCell::new(0),
        }
    }

    fn align_up(addr: usize, align: usize) -> usize {
        (addr + align - 1) & !(align - 1)
    }
}

unsafe impl GlobalAlloc for BumpAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let heap = &*self.heap.get();
        let next = &mut *self.next.get();

        let alloc_start = Self::align_up(*next, layout.align());
        let alloc_end = alloc_start.saturating_add(layout.size());

        if alloc_end > heap.len() {
            ptr::null_mut()
        } else {
            *next = alloc_end;
            heap.as_ptr().add(alloc_start) as *mut u8
        }
    }

    unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {
        // Bump allocator doesn't support deallocation
    }
}

#[global_allocator]
static ALLOCATOR: BumpAllocator = BumpAllocator::new();

Pool Allocator

📝 Example: Fixed-Size Block Allocator

use core::alloc::{GlobalAlloc, Layout};
use core::mem;
use core::ptr::{self, NonNull};

const BLOCK_SIZES: &[usize] = &[8, 16, 32, 64, 128, 256, 512, 1024, 2048];

struct ListNode {
    next: Option<&'static mut ListNode>,
}

pub struct PoolAllocator {
    list_heads: [Option<&'static mut ListNode>; BLOCK_SIZES.len()],
    fallback_allocator: linked_list_allocator::Heap,
}

impl PoolAllocator {
    pub const fn new() -> Self {
        const EMPTY: Option<&'static mut ListNode> = None;
        PoolAllocator {
            list_heads: [EMPTY; BLOCK_SIZES.len()],
            fallback_allocator: linked_list_allocator::Heap::empty(),
        }
    }

    pub unsafe fn init(&mut self, heap_start: usize, heap_size: usize) {
        self.fallback_allocator.init(heap_start, heap_size);
    }

    fn fallback_alloc(&mut self, layout: Layout) -> *mut u8 {
        match self.fallback_allocator.allocate_first_fit(layout) {
            Ok(ptr) => ptr.as_ptr(),
            Err(_) => ptr::null_mut(),
        }
    }
}

fn list_index(layout: &Layout) -> Option {
    let required_block_size = layout.size().max(layout.align());
    BLOCK_SIZES.iter().position(|&s| s >= required_block_size)
}

unsafe impl GlobalAlloc for PoolAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let allocator = &mut *(self as *const Self as *mut Self);

        match list_index(&layout) {
            Some(index) => {
                match allocator.list_heads[index].take() {
                    Some(node) => {
                        allocator.list_heads[index] = node.next.take();
                        node as *mut ListNode as *mut u8
                    }
                    None => {
                        let block_size = BLOCK_SIZES[index];
                        let block_align = block_size;
                        let layout = Layout::from_size_align(block_size, block_align)
                            .unwrap();
                        allocator.fallback_alloc(layout)
                    }
                }
            }
            None => allocator.fallback_alloc(layout),
        }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let allocator = &mut *(self as *const Self as *mut Self);

        match list_index(&layout) {
            Some(index) => {
                let new_node = ListNode {
                    next: allocator.list_heads[index].take(),
                };
                let new_node_ptr = ptr as *mut ListNode;
                new_node_ptr.write(new_node);
                allocator.list_heads[index] = Some(&mut *new_node_ptr);
            }
            None => {
                let ptr = NonNull::new(ptr).unwrap();
                allocator.fallback_allocator.deallocate(ptr, layout);
            }
        }
    }
}

Arena Allocator

📝 Example: Arena/Region Allocator

use std::cell::Cell;
use std::marker::PhantomData;
use std::mem;
use std::ptr;

pub struct Arena<'a> {
    buffer: Vec,
    offset: Cell,
    _phantom: PhantomData<&'a ()>,
}

impl<'a> Arena<'a> {
    pub fn new(size: usize) -> Self {
        Arena {
            buffer: vec![0; size],
            offset: Cell::new(0),
            _phantom: PhantomData,
        }
    }

    pub fn alloc(&'a self, value: T) -> &'a mut T {
        let offset = self.offset.get();
        let align = mem::align_of::();
        let size = mem::size_of::();

        let aligned_offset = (offset + align - 1) & !(align - 1);
        let new_offset = aligned_offset + size;

        if new_offset > self.buffer.len() {
            panic!("Arena out of memory");
        }

        self.offset.set(new_offset);

        let ptr = unsafe {
            self.buffer.as_ptr().add(aligned_offset) as *mut T
        };

        unsafe {
            ptr.write(value);
            &mut *ptr
        }
    }

    pub fn alloc_slice(&'a self, slice: &[T]) -> &'a mut [T] {
        let offset = self.offset.get();
        let align = mem::align_of::();
        let size = mem::size_of::() * slice.len();

        let aligned_offset = (offset + align - 1) & !(align - 1);
        let new_offset = aligned_offset + size;

        if new_offset > self.buffer.len() {
            panic!("Arena out of memory");
        }

        self.offset.set(new_offset);

        let ptr = unsafe {
            self.buffer.as_ptr().add(aligned_offset) as *mut T
        };

        unsafe {
            ptr::copy_nonoverlapping(slice.as_ptr(), ptr, slice.len());
            std::slice::from_raw_parts_mut(ptr, slice.len())
        }
    }

    pub fn reset(&self) {
        self.offset.set(0);
    }
}

// Usage example
fn main() {
    let arena = Arena::new(1024);

    let x = arena.alloc(42i32);
    let y = arena.alloc(3.14f64);
    let arr = arena.alloc_slice(&[1, 2, 3, 4, 5]);

    println!("x = {}", x);
    println!("y = {}", y);
    println!("arr = {:?}", arr);

    // All allocations freed when arena is dropped
}

Part 2: SIMD Programming

Introduction to SIMD

ℹ️ What is SIMD? Single Instruction, Multiple Data allows processing multiple data points with a single CPU instruction. Modern CPUs support SSE, AVX, and AVX-512 instruction sets for vectorized operations.

Portable SIMD (std::simd)

📝 Example: Basic SIMD Operations

#![feature(portable_simd)]

use std::simd::*;

fn main() {
    // Create SIMD vectors
    let a = f32x8::from_array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]);
    let b = f32x8::from_array([8.0, 7.0, 6.0, 5.0, 4.0, 3.0, 2.0, 1.0]);

    // Vectorized addition
    let sum = a + b;
    println!("Sum: {:?}", sum.to_array());

    // Vectorized multiplication
    let product = a * b;
    println!("Product: {:?}", product.to_array());

    // Horizontal sum
    let total: f32 = sum.reduce_sum();
    println!("Total: {}", total);
}

// Vector dot product
fn dot_product_simd(a: &[f32], b: &[f32]) -> f32 {
    assert_eq!(a.len(), b.len());

    let mut sum = f32x8::splat(0.0);
    let chunks = a.len() / 8;

    for i in 0..chunks {
        let a_vec = f32x8::from_slice(&a[i * 8..]);
        let b_vec = f32x8::from_slice(&b[i * 8..]);
        sum += a_vec * b_vec;
    }

    let mut result = sum.reduce_sum();

    // Handle remainder
    for i in (chunks * 8)..a.len() {
        result += a[i] * b[i];
    }

    result
}

Architecture-Specific SIMD

📝 Example: AVX2 Intrinsics

#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::*;

#[target_feature(enable = "avx2")]
unsafe fn sum_avx2(data: &[f32]) -> f32 {
    let mut sum = _mm256_setzero_ps();

    let chunks = data.len() / 8;
    for i in 0..chunks {
        let values = _mm256_loadu_ps(data.as_ptr().add(i * 8));
        sum = _mm256_add_ps(sum, values);
    }

    // Horizontal sum
    let sum128 = _mm_add_ps(
        _mm256_extractf128_ps(sum, 0),
        _mm256_extractf128_ps(sum, 1),
    );

    let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128));
    let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 0x55));

    let mut result = _mm_cvtss_f32(sum32);

    // Handle remainder
    for i in (chunks * 8)..data.len() {
        result += data[i];
    }

    result
}

pub fn sum_fast(data: &[f32]) -> f32 {
    #[cfg(target_arch = "x86_64")]
    {
        if is_x86_feature_detected!("avx2") {
            unsafe { sum_avx2(data) }
        } else {
            data.iter().sum()
        }
    }

    #[cfg(not(target_arch = "x86_64"))]
    {
        data.iter().sum()
    }
}

Real-World SIMD: Image Processing

📝 Example: Grayscale Conversion

#![feature(portable_simd)]

use std::simd::*;

// Convert RGB to grayscale using SIMD
fn rgb_to_gray_simd(rgb: &[u8], gray: &mut [u8]) {
    assert_eq!(rgb.len(), gray.len() * 3);

    let r_weight = f32x8::splat(0.299);
    let g_weight = f32x8::splat(0.587);
    let b_weight = f32x8::splat(0.114);

    let chunks = gray.len() / 8;

    for i in 0..chunks {
        let idx = i * 8;
        let rgb_idx = idx * 3;

        // Load RGB values
        let mut r = [0.0f32; 8];
        let mut g = [0.0f32; 8];
        let mut b = [0.0f32; 8];

        for j in 0..8 {
            r[j] = rgb[rgb_idx + j * 3] as f32;
            g[j] = rgb[rgb_idx + j * 3 + 1] as f32;
            b[j] = rgb[rgb_idx + j * 3 + 2] as f32;
        }

        let r_vec = f32x8::from_array(r);
        let g_vec = f32x8::from_array(g);
        let b_vec = f32x8::from_array(b);

        // Weighted sum
        let result = r_vec * r_weight + g_vec * g_weight + b_vec * b_weight;

        // Convert to u8 and store
        for j in 0..8 {
            gray[idx + j] = result.to_array()[j] as u8;
        }
    }

    // Handle remainder
    for i in (chunks * 8)..gray.len() {
        let rgb_idx = i * 3;
        gray[i] = (
            rgb[rgb_idx] as f32 * 0.299 +
            rgb[rgb_idx + 1] as f32 * 0.587 +
            rgb[rgb_idx + 2] as f32 * 0.114
        ) as u8;
    }
}

// Benchmark comparison
#[cfg(test)]
mod bench {
    use super::*;
    use test::Bencher;

    #[bench]
    fn bench_simd(b: &mut Bencher) {
        let rgb = vec![128u8; 1920 * 1080 * 3];
        let mut gray = vec![0u8; 1920 * 1080];

        b.iter(|| {
            rgb_to_gray_simd(&rgb, &mut gray);
        });
    }
}

SIMD Matrix Operations

📝 Example: Matrix Multiplication

#![feature(portable_simd)]

use std::simd::*;

pub struct Matrix {
    data: Vec,
    rows: usize,
    cols: usize,
}

impl Matrix {
    pub fn new(rows: usize, cols: usize) -> Self {
        Matrix {
            data: vec![0.0; rows * cols],
            rows,
            cols,
        }
    }

    pub fn get(&self, row: usize, col: usize) -> f32 {
        self.data[row * self.cols + col]
    }

    pub fn set(&mut self, row: usize, col: usize, value: f32) {
        self.data[row * self.cols + col] = value;
    }

    // SIMD-accelerated matrix multiplication
    pub fn multiply_simd(&self, other: &Matrix) -> Matrix {
        assert_eq!(self.cols, other.rows);

        let mut result = Matrix::new(self.rows, other.cols);

        for i in 0..self.rows {
            for j in 0..other.cols {
                let mut sum = f32x8::splat(0.0);
                let chunks = self.cols / 8;

                for k in 0..chunks {
                    let a_vec = f32x8::from_slice(
                        &self.data[i * self.cols + k * 8..]
                    );

                    let mut b_values = [0.0f32; 8];
                    for l in 0..8 {
                        b_values[l] = other.get(k * 8 + l, j);
                    }
                    let b_vec = f32x8::from_array(b_values);

                    sum += a_vec * b_vec;
                }

                let mut total = sum.reduce_sum();

                // Handle remainder
                for k in (chunks * 8)..self.cols {
                    total += self.get(i, k) * other.get(k, j);
                }

                result.set(i, j, total);
            }
        }

        result
    }
}

Performance Optimization

🔑 Optimization Strategies

Technique	Benefit	Use Case
SIMD	2-8x speedup	Vector operations, image processing
Custom Allocator	Reduced fragmentation	Games, real-time systems
Arena Allocation	Fast bulk deallocation	Request handling, parsing
Pool Allocator	O(1) allocation	Fixed-size objects
Cache Alignment	Avoid false sharing	Multi-threaded code

Best Practices

⚠️ Performance Guidelines:

Profile First: Measure before optimizing
Test Thoroughly: SIMD code can have edge cases
Consider Portability: Not all CPUs support all SIMD
Handle Alignment: Misaligned access can be slow
Benchmark: Verify performance improvements

Summary

✅ You've learned:

Implementing custom memory allocators
Bump, pool, and arena allocation strategies
SIMD programming with portable_simd
Architecture-specific SIMD intrinsics
Real-world SIMD applications
Performance optimization techniques

Custom allocators and SIMD are powerful tools for achieving maximum performance in Rust. By understanding memory allocation patterns and leveraging CPU vectorization, you can build extremely efficient systems that rival or exceed C/C++ performance.