spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#pragma once34#include <cstddef>5#include <mutex>6#include <unordered_map>7#include <vector>89namespace MTL {10class Buffer;11class Device;12}1314namespace forge {1516// Size-bucketed MTLBuffer pool. Every tensor allocation goes through here so17// the training loop never touches newBuffer (a kernel-level VM allocation).18// Buffers use MTLStorageModeShared: the CPU pointer is valid at all times.19//20// Sizes are rounded up to power-of-two buckets (min 256 bytes, which also21// satisfies every setBuffer offset-alignment rule). release() returns the22// buffer to its bucket; it is the caller's job (Storage refcounting, and23// later the command-buffer retire list) to only release buffers no in-flight24// command buffer references.25class Allocator {26public:27    explicit Allocator(MTL::Device* device);28    ~Allocator();2930    Allocator(const Allocator&) = delete;31    Allocator& operator=(const Allocator&) = delete;3233    // Returns a retained buffer of capacity >= nbytes.34    MTL::Buffer* acquire(size_t nbytes);35    void release(MTL::Buffer* buffer);3637    // While deferring (a command buffer is being encoded / in flight),38    // released buffers park on a retire list instead of the free pool:39    // recycling them would let CPU-side writes (zeros, fills) race pending40    // GPU reads. The Stream flips this on at first encode and flushes at41    // sync.42    void set_defer(bool defer);43    void flush_retired();4445    // Drop all pooled (free) buffers back to the OS.46    void trim();4748    size_t bytes_live() const { return bytes_live_; }49    size_t bytes_pooled() const { return bytes_pooled_; }5051private:52    static size_t bucket_size(size_t nbytes);5354    MTL::Device* device_; // borrowed, owned by forge::Device55    std::mutex mutex_;56    std::unordered_map<size_t, std::vector<MTL::Buffer*>> pool_;57    std::vector<MTL::Buffer*> retired_;58    bool defer_ = false;59    size_t bytes_live_ = 0;60    size_t bytes_pooled_ = 0;61};6263} // namespace forge64