spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#pragma once34#include "core/allocator.h"56#include <memory>7#include <mutex>8#include <string>9#include <unordered_map>1011namespace MTL {12class Device;13class CommandQueue;14class Library;15class ComputePipelineState;16class FunctionConstantValues;17}1819namespace forge {2021// Owns the Metal device, the single command queue, the kernel library and22// the pipeline cache. One instance per process (Device::get()).23//24// All pipelines are created lazily on first use and cached by25// "kernel_name" or "kernel_name/constants_key" so steady-state training26// never touches pipeline creation.27class Device {28public:29    static Device& get();3031    MTL::Device* mtl() const { return device_; }32    MTL::CommandQueue* queue() const { return queue_; }33    Allocator& allocator() { return *allocator_; }3435    // Plain kernel, no function constants.36    MTL::ComputePipelineState* pipeline(const std::string& kernel_name);3738    // Specialized kernel. constants_key must uniquely identify the constant39    // values (e.g. "am1_an0_ak1"); it is only used as a cache key.40    MTL::ComputePipelineState* pipeline(const std::string& kernel_name,41                                        const MTL::FunctionConstantValues* constants,42                                        const std::string& constants_key);4344    std::string name() const;45    size_t recommended_working_set() const;4647    Device(const Device&) = delete;48    Device& operator=(const Device&) = delete;4950private:51    Device();52    ~Device();5354    // Looks for forge.metallib next to the executable, in the current55    // directory, or at $FORGE_METALLIB.56    void load_library();5758    MTL::Device* device_ = nullptr;59    MTL::CommandQueue* queue_ = nullptr;60    MTL::Library* library_ = nullptr;61    std::unique_ptr<Allocator> allocator_;62    std::mutex pipeline_mutex_;63    std::unordered_map<std::string, MTL::ComputePipelineState*> pipelines_;64};6566} // namespace forge67