MemoryArena.h

Go to the documentation of this file.

#pragma once
#include <cstddef>
#include <cstdlib>
#include <new>
#include <stdexcept>
#if defined(_WIN32)
#include <malloc.h>
#include <windows.h>
#else
#include <sys/mman.h>
#endif

 
namespace Deep
{
    namespace detail
    {
        // std::aligned_alloc isn't reliably available on Windows/MinGW: the
        // standard requires memory from it be freed with plain free(), but
        // Windows' native aligned allocator needs the opposite pairing
        // (_aligned_malloc / _aligned_free), so many MinGW toolchains simply
        // omit std::aligned_alloc rather than violate the standard's contract.

        inline void *AlignedAllocPortable(size_t alignment, size_t size)
        {
#if defined(_WIN32)
            return _aligned_malloc(size, alignment); // note: (size, alignment) -- reversed vs aligned_alloc
#else
            return std::aligned_alloc(alignment, size);
#endif
        }

        inline void AlignedFreePortable(void *ptr)
        {
#if defined(_WIN32)
            _aligned_free(ptr);
#else
            std::free(ptr);
#endif
        }

        inline void *HugePageAllocPortable(size_t size, bool &succeeded)
        {
            succeeded = false;
#if defined(_WIN32)
            void *ptr = VirtualAlloc(nullptr, size, MEM_RESERVE | MEM_COMMIT | MEM_LARGE_PAGES, PAGE_READWRITE);
            if (ptr)
                succeeded = true;
            return ptr;
#else
            void *ptr = mmap(nullptr, size, PROT_READ | PROT_WRITE,
                              MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
            if (ptr == MAP_FAILED)
                return nullptr;
 
            // Best-effort hint -- if this fails, we still have valid,
            // page-aligned (though not huge-page-backed) memory from the
            // mmap itself, so this isn't treated as a hard failure.
            madvise(ptr, size, MADV_HUGEPAGE);
            succeeded = true;
            return ptr;
#endif
        }

        inline void HugePageFreePortable(void *ptr, size_t size)
        {
#if defined(_WIN32)
            VirtualFree(ptr, 0, MEM_RELEASE);
#else
            munmap(ptr, size);
#endif
        }
    } // namespace detail

    class MemoryArena
    {
    private:
        float *base_ptr;
        size_t capacity_bytes;
        size_t offset_bytes;
        bool used_huge_pages;
 
    public:
        explicit MemoryArena(size_t total_floats, bool use_huge_pages = false)
            : used_huge_pages(false)
        {
            // Guarantee total capacity is a multiple of 64 bytes
            capacity_bytes = (total_floats * sizeof(float) + 63) & ~63;
 
            if (use_huge_pages)
            {
                bool succeeded = false;
                base_ptr = static_cast<float *>(detail::HugePageAllocPortable(capacity_bytes, succeeded));
                used_huge_pages = succeeded;
            }
 
            if (!used_huge_pages)
            {
                // Either huge pages weren't requested, or the request
                // wasn't honored (mmap/VirtualAlloc failed, or the
                // system lacks the required configuration/privilege) --
                // fall back to the standard aligned allocator rather
                // than treating this as a hard failure.
                base_ptr = static_cast<float *>(detail::AlignedAllocPortable(64, capacity_bytes));
            }
 
            if (!base_ptr)
            {
                throw std::bad_alloc();
            }
            offset_bytes = 0;
        }

        ~MemoryArena()
        {
            if (base_ptr)
            {
                if (used_huge_pages)
                    detail::HugePageFreePortable(base_ptr, capacity_bytes);
                else
                    detail::AlignedFreePortable(base_ptr);
            }
        }
 
        // Delete copy/move constructors to prevent double-free corruption

        MemoryArena(const MemoryArena &) = delete;
        MemoryArena &operator=(const MemoryArena &) = delete;

        float *AllocateFloats(size_t num_floats)
        {
            // Calculate how many bytes we need, padded to the nearest 64-byte boundary
            size_t allocation_size = (num_floats * sizeof(float) + 63) & ~63;
 
            if (offset_bytes + allocation_size > capacity_bytes)
            {
                throw std::runtime_error("Fatal: MemoryArena capacity exceeded during allocation.");
            }
 
            float *chunk = reinterpret_cast<float *>(
                reinterpret_cast<char *>(base_ptr) + offset_bytes);
 
            offset_bytes += allocation_size;
            return chunk;
        }

        size_t GetUsedBytes() const { return offset_bytes; }
        size_t GetCapacityBytes() const { return capacity_bytes; }
        bool UsedHugePages() const { return used_huge_pages; }
    };
} // namespace Deep