Hands-on Practical: Implementing a Simple Allocator

Efficient memory management is fundamental to performant ML runtime systems. Tensors, especially intermediate activation tensors, can consume significant memory. Standard general-purpose allocators like malloc might introduce overhead (metadata storage, fragmentation, slower allocation/deallocation) that becomes noticeable under the intense allocation patterns of ML inference.

One common strategy employed in specialized runtimes is the use of an arena allocator, also known as a bump allocator. The core idea is simple: pre-allocate a large, contiguous block of memory (the arena) and then satisfy allocation requests by simply "bumping" a pointer forward within that block. Deallocation is typically achieved by resetting the pointer back to the beginning of the arena, effectively freeing all memory allocated within it simultaneously. This approach is particularly well-suited for scenarios like single-inference execution where many temporary tensors are created and destroyed together.

Let's implement a basic arena allocator in C++ to illustrate the concept.

Arena Allocator Concept

Imagine the arena as a large reserved memory region. We maintain a pointer, initially pointing to the start of this region. When an allocation request arrives, we check if there's enough space remaining, adjust the pointer forward by the requested size (plus any necessary alignment padding), and return the original pointer position.

Progression of an arena allocator. Allocations increment the offset pointer within the pre-allocated buffer. Resetting moves the offset back to zero.

Basic C++ Implementation

"We'll create a simple ArenaAllocator class. For simplicity, this example omits thread safety and uses basic C++ features. In a practical scenario, you'd likely use platform-specific APIs for aligned memory allocation and potentially incorporate locking mechanisms for multi-threaded use."

#include <cstddef> // for std::size_t, std::byte
#include <cstdint> // for std::uintptr_t
#include <memory>  // for std::unique_ptr
#include <stdexcept> // for std::runtime_error, std::bad_alloc
#include <vector>
#include <numeric> // for std::lcm

// Helper function to align pointers
inline void* align_pointer(void* ptr, std::size_t alignment) {
    if (alignment == 0) {
        return ptr; // No alignment needed
    }
    std::uintptr_t int_ptr = reinterpret_cast<std::uintptr_t>(ptr);
    std::size_t remainder = int_ptr % alignment;
    if (remainder != 0) {
        int_ptr += (alignment - remainder);
    }
    return reinterpret_cast<void*>(int_ptr);
}

class ArenaAllocator {
public:
    // Constructor: Allocates the arena buffer
    explicit ArenaAllocator(std::size_t size_bytes)
        : buffer_size_(size_bytes), current_offset_(0) {
        // Allocate raw memory. In production, consider aligned allocation.
        raw_buffer_ = std::make_unique<std::byte[]>(buffer_size_);
        if (!raw_buffer_) {
            throw std::bad_alloc();
        }
        start_ptr_ = raw_buffer_.get();
        current_ptr_ = start_ptr_;
        // std::cout << "Arena allocated: " << buffer_size_ << " bytes at " << start_ptr_ << std::endl;
    }

    // Deleted copy constructor and assignment operator
    ArenaAllocator(const ArenaAllocator&) = delete;
    ArenaAllocator& operator=(const ArenaAllocator&) = delete;

    // Allocate memory from the arena
    void* allocate(std::size_t size_bytes, std::size_t alignment = alignof(std::max_align_t)) {
        if (size_bytes == 0) {
            return nullptr; // Standard behavior for zero-size allocation
        }

        // Ensure alignment is a power of 2
        if (alignment == 0 || (alignment & (alignment - 1)) != 0) {
             throw std::invalid_argument("Alignment must be a power of 2");
        }

        // Calculate the aligned pointer
        void* aligned_ptr = align_pointer(current_ptr_, alignment);
        std::size_t alignment_padding = reinterpret_cast<std::byte*>(aligned_ptr) - reinterpret_cast<std::byte*>(current_ptr_);

        // Calculate total size needed (including alignment padding)
        std::size_t total_size_needed = alignment_padding + size_bytes;

        // Check if there's enough space
        if (current_offset_ + total_size_needed > buffer_size_) {
             // std::cerr << "Arena allocation failed: Requested " << size_bytes
             //           << " (total " << total_size_needed << "), available "
             //           << (buffer_size_ - current_offset_) << std::endl;
            throw std::bad_alloc(); // Not enough space
        }

        // "Allocate" the memory by bumping the pointers/offset
        void* result_ptr = aligned_ptr;
        current_ptr_ = reinterpret_cast<std::byte*>(aligned_ptr) + size_bytes;
        current_offset_ += total_size_needed;

        // std::cout << "Allocated " << size_bytes << " bytes (aligned at " << alignment << "). "
        //           << "Offset: " << current_offset_ << "/" << buffer_size_ << std::endl;

        return result_ptr;
    }

    // Reset the arena, effectively deallocating all memory within it
    void reset() {
        current_offset_ = 0;
        current_ptr_ = start_ptr_;
        // std::cout << "Arena reset." << std::endl;
    }

    // Get remaining free space (approximate, doesn't account for future alignment needs)
    std::size_t get_free_space() const {
        return buffer_size_ - current_offset_;
    }

    // Get total size of the arena
    std::size_t get_total_size() const {
        return buffer_size_;
    }

private:
    std::unique_ptr<std::byte[]> raw_buffer_; // Owns the memory
    void* start_ptr_;       // Start of the usable buffer
    void* current_ptr_;     // Current allocation position
    std::size_t buffer_size_;   // Total size of the buffer
    std::size_t current_offset_; // Current offset from the start (tracks total used including padding)
};

Using the Allocator

Here's how you might use this allocator in a simplified inference context:

#include <iostream>
#include <vector>

// Dummy Tensor class (replace with actual tensor implementation)
struct Tensor {
    void* data;
    std::size_t size;
    // ... other metadata
};

int main() {
    const std::size_t ARENA_SIZE = 1024 * 1024; // 1 MiB Arena
    ArenaAllocator allocator(ARENA_SIZE);

    try {
        std::cout << "Initial free space: " << allocator.get_free_space() << std::endl;

        // Simulate allocating tensors for different layers
        Tensor activation1;
        activation1.size = 256 * 1024; // 256 KiB
        activation1.data = allocator.allocate(activation1.size, 64); // Request 64-byte alignment
        std::cout << "Allocated Tensor 1 (" << activation1.size << " bytes). Free space: " << allocator.get_free_space() << std::endl;

        Tensor weights1;
        weights1.size = 128 * 1024; // 128 KiB
        weights1.data = allocator.allocate(weights1.size, 32); // Request 32-byte alignment
        std::cout << "Allocated Tensor 2 (" << weights1.size << " bytes). Free space: " << allocator.get_free_space() << std::endl;

        Tensor activation2;
        activation2.size = 512 * 1024; // 512 KiB
        activation2.data = allocator.allocate(activation2.size); // Default alignment
        std::cout << "Allocated Tensor 3 (" << activation2.size << " bytes). Free space: " << allocator.get_free_space() << std::endl;

        // Simulate end of inference run
        std::cout << "Inference complete. Resetting arena." << std::endl;
        allocator.reset();
        std::cout << "Free space after reset: " << allocator.get_free_space() << std::endl;

        // Can reuse the arena for the next inference
        Tensor next_run_tensor;
        next_run_tensor.size = 100 * 1024;
        next_run_tensor.data = allocator.allocate(next_run_tensor.size);
        std::cout << "Allocated tensor for next run. Free space: " << allocator.get_free_space() << std::endl;

    } catch (const std::bad_alloc& e) {
        std::cerr << "Allocation failed: Out of memory in arena." << std::endl;
    } catch (const std::exception& e) {
        std::cerr << "An error occurred: " << e.what() << std::endl;
    }

    return 0;
}

Discussion and Trade-offs

Advantages:

1. Speed: Allocation is extremely fast, often just pointer arithmetic and a bounds check. Deallocation (reset) is instantaneous ($O(1)$), regardless of the number of allocations.
2. Reduced Fragmentation: Since memory is allocated contiguously, external fragmentation is eliminated within the arena. Internal fragmentation can still occur due to alignment padding.
3. Locality: Allocating related tensors (like activations in sequential layers) consecutively can potentially improve cache locality, although this depends heavily on access patterns.

Disadvantages:

1. Fixed Size: The arena has a fixed size. If an inference run requires more memory than the arena provides, allocation fails. Determining the required size upfront (memory planning) is important.
2. No Individual Deallocation: Memory can only be freed all at once by resetting the arena. This makes it unsuitable for scenarios with long-lived objects or objects with very different lifespans that need individual management.
3. Potential Waste: If the peak memory usage during one phase is much higher than another, but the arena needs to accommodate the peak, memory might be held unnecessarily long. The entire arena buffer is held even if only a small portion is actively used after the high-water mark.

Extensions:

• Thread Safety: Add mutexes or use thread-local arenas.
• Multiple Arenas: Use different arenas for different types of allocations (e.g., persistent weights vs. transient activations).
• Fallback Allocator: If the arena runs out of space, potentially fall back to a standard allocator (though this complicates deallocation).
• Growth: Implement logic to grow the arena (reallocate and copy) if needed, although this negates some of the performance benefits.

This hands-on example provides a foundational understanding of arena allocation. While simple, this pattern is surprisingly effective and forms the basis for memory management in many high-performance ML runtime systems where allocation/deallocation patterns are predictable within a specific scope, like a single model execution. Building upon this, runtimes often incorporate more sophisticated static analysis (memory planning) to determine optimal arena sizes and allocation strategies.