DeviceMemoryArena.h

Go to the documentation of this file.

#pragma once
#if defined(DEEPITY_USE_CUDA)
#include <deepity/backend/IComputeBackend.h>
#include <new>
#include <stdexcept>
#include <iostream>

 
namespace Deep
{
    class DeviceMemoryArena
    {
    private:
        IComputeBackend *backend;
        float *base_ptr;
        size_t capacity_bytes;
        size_t offset_bytes;
 
    public:
        DeviceMemoryArena(IComputeBackend *backend, size_t total_floats)
            : backend(backend)
        {
            capacity_bytes = (total_floats * sizeof(float) + 63) & ~(size_t)63;
 
            base_ptr = backend->Allocate(capacity_bytes / sizeof(float));
            if (!base_ptr)
            {
                throw std::bad_alloc();
            }
            offset_bytes = 0;
        }

        ~DeviceMemoryArena()
        {
            if (base_ptr)
            {
                backend->Free(base_ptr);
            }
        }
 
        // Delete copy/move constructors to prevent double-free corruption
        // -- the original version of this class had no such guard, a real
        // gap relative to MemoryArena's own established protection.

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

        float *AllocateFloats(size_t num_floats)
        {
            size_t allocation_size = (num_floats * sizeof(float) + 63) & ~(size_t)63;
 
            if (offset_bytes + allocation_size > capacity_bytes)
            {
                throw std::runtime_error("Fatal: DeviceMemoryArena 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; }
    };
}
#endif