spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#include "core/allocator.h"34#include <Metal/Metal.hpp>56#include <cstdio>7#include <cstdlib>89namespace forge {1011Allocator::Allocator(MTL::Device* device) : device_(device) {}1213Allocator::~Allocator() {14    trim();15}1617size_t Allocator::bucket_size(size_t nbytes) {18    size_t b = 256;19    while (b < nbytes) b <<= 1;20    return b;21}2223MTL::Buffer* Allocator::acquire(size_t nbytes) {24    const size_t bucket = bucket_size(nbytes);25    {26        std::lock_guard<std::mutex> lock(mutex_);27        auto it = pool_.find(bucket);28        if (it != pool_.end() && !it->second.empty()) {29            MTL::Buffer* buf = it->second.back();30            it->second.pop_back();31            bytes_pooled_ -= bucket;32            bytes_live_ += bucket;33            return buf;34        }35    }36    MTL::Buffer* buf = device_->newBuffer(bucket, MTL::ResourceStorageModeShared);37    if (!buf) {38        std::fprintf(stderr, "forge: Metal buffer allocation failed (%zu bytes)\n", bucket);39        std::abort();40    }41    std::lock_guard<std::mutex> lock(mutex_);42    bytes_live_ += bucket;43    return buf;44}4546void Allocator::release(MTL::Buffer* buffer) {47    if (!buffer) return;48    const size_t bucket = buffer->length();49    std::lock_guard<std::mutex> lock(mutex_);50    if (defer_) {51        retired_.push_back(buffer);52    } else {53        pool_[bucket].push_back(buffer);54        bytes_pooled_ += bucket;55    }56    bytes_live_ -= bucket;57}5859void Allocator::set_defer(bool defer) {60    std::lock_guard<std::mutex> lock(mutex_);61    defer_ = defer;62}6364void Allocator::flush_retired() {65    std::lock_guard<std::mutex> lock(mutex_);66    for (MTL::Buffer* buf : retired_) {67        pool_[buf->length()].push_back(buf);68        bytes_pooled_ += buf->length();69    }70    retired_.clear();71}7273void Allocator::trim() {74    std::lock_guard<std::mutex> lock(mutex_);75    for (auto& [bucket, buffers] : pool_) {76        for (MTL::Buffer* buf : buffers) buf->release();77    }78    pool_.clear();79    bytes_pooled_ = 0;80}8182} // namespace forge83