spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#include "core/device.h"34#include <Foundation/Foundation.hpp>5#include <Metal/Metal.hpp>67#include <cstdio>8#include <cstdlib>9#include <filesystem>1011#include <mach-o/dyld.h>1213namespace forge {1415namespace {1617std::string executable_dir() {18    char buf[4096];19    uint32_t size = sizeof(buf);20    if (_NSGetExecutablePath(buf, &size) != 0) return {};21    std::error_code ec;22    auto canonical = std::filesystem::canonical(buf, ec);23    if (ec) return {};24    return canonical.parent_path().string();25}2627[[noreturn]] void die(const char* msg, NS::Error* err) {28    std::fprintf(stderr, "forge: %s", msg);29    if (err && err->localizedDescription()) {30        std::fprintf(stderr, ": %s", err->localizedDescription()->utf8String());31    }32    std::fprintf(stderr, "\n");33    std::abort();34}3536} // namespace3738Device& Device::get() {39    static Device instance;40    return instance;41}4243Device::Device() {44    device_ = MTL::CreateSystemDefaultDevice();45    if (!device_) die("no Metal device found", nullptr);46    queue_ = device_->newCommandQueue();47    if (!queue_) die("failed to create command queue", nullptr);48    allocator_ = std::make_unique<Allocator>(device_);49    load_library();50}5152Device::~Device() {53    for (auto& [key, pso] : pipelines_) pso->release();54    if (library_) library_->release();55    allocator_.reset(); // must drop pooled buffers before the device56    if (queue_) queue_->release();57    if (device_) device_->release();58}5960void Device::load_library() {61    std::vector<std::string> candidates;62    if (const char* env = std::getenv("FORGE_METALLIB")) candidates.push_back(env);63    if (auto dir = executable_dir(); !dir.empty()) candidates.push_back(dir + "/forge.metallib");64    candidates.push_back("forge.metallib");6566    for (const auto& path : candidates) {67        if (!std::filesystem::exists(path)) continue;68        NS::Error* err = nullptr;69        NS::String* nspath = NS::String::string(path.c_str(), NS::UTF8StringEncoding);70        library_ = device_->newLibrary(nspath, &err);71        if (library_) return;72        die("failed to load forge.metallib", err);73    }74    die("forge.metallib not found (looked next to the executable, in cwd, and at $FORGE_METALLIB)",75        nullptr);76}7778MTL::ComputePipelineState* Device::pipeline(const std::string& kernel_name) {79    return pipeline(kernel_name, nullptr, "");80}8182MTL::ComputePipelineState* Device::pipeline(const std::string& kernel_name,83                                            const MTL::FunctionConstantValues* constants,84                                            const std::string& constants_key) {85    const std::string key =86        constants_key.empty() ? kernel_name : kernel_name + "/" + constants_key;8788    std::lock_guard<std::mutex> lock(pipeline_mutex_);89    if (auto it = pipelines_.find(key); it != pipelines_.end()) return it->second;9091    NS::Error* err = nullptr;92    NS::String* nsname = NS::String::string(kernel_name.c_str(), NS::UTF8StringEncoding);93    MTL::Function* fn = constants94        ? library_->newFunction(nsname, constants, &err)95        : library_->newFunction(nsname);96    if (!fn) die(("kernel not found: " + kernel_name).c_str(), err);9798    MTL::ComputePipelineState* pso = device_->newComputePipelineState(fn, &err);99    fn->release();100    if (!pso) die(("pipeline creation failed: " + key).c_str(), err);101102    pipelines_.emplace(key, pso);103    return pso;104}105106std::string Device::name() const {107    return device_->name()->utf8String();108}109110size_t Device::recommended_working_set() const {111    return device_->recommendedMaxWorkingSetSize();112}113114} // namespace forge115