spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#include "core/fmodel.h"34#include "core/device.h"56#include <Metal/Metal.hpp>78#include <nlohmann/json.hpp>910#include <fcntl.h>11#include <sys/mman.h>12#include <sys/stat.h>13#include <unistd.h>1415#include <cinttypes>16#include <cstdio>17#include <cstdlib>18#include <cstring>19#include <filesystem>20#include <fstream>21#include <unordered_map>22#include <unordered_set>2324namespace forge::fmodel {2526namespace {2728namespace fs = std::filesystem;29using nlohmann::json;3031// Apple Silicon VM page — the alignment contract of the whole format.32constexpr size_t kPage = 16384;33constexpr uint32_t kShardMagic = 0x53475246; // "FRGS"34constexpr uint32_t kVersion = 1;3536[[noreturn]] void die(const std::string& msg) {37    std::fprintf(stderr, "forge/fmodel: %s\n", msg.c_str());38    std::abort();39}4041size_t page_align(size_t n) { return (n + kPage - 1) & ~(kPage - 1); }4243uint64_t fnv1a64(const void* data, size_t n) {44    const uint8_t* p = static_cast<const uint8_t*>(data);45    uint64_t h = 1469598103934665603ull;46    for (size_t i = 0; i < n; ++i) {47        h ^= p[i];48        h *= 1099511628211ull;49    }50    return h;51}5253std::string hex64(uint64_t v) {54    char buf[17];55    std::snprintf(buf, sizeof(buf), "%016" PRIx64, v);56    return buf;57}5859DType dtype_from_name(const std::string& s) {60    if (s == "f32") return DType::F32;61    if (s == "f16") return DType::F16;62    if (s == "bf16") return DType::BF16;63    if (s == "u16") return DType::U16;64    if (s == "i32") return DType::I32;65    die("unknown dtype in manifest: " + s);66}6768// Convert one f32 tensor into `dtype` bytes (identity for f32).69std::vector<uint8_t> convert_out(const Tensor& t, DType dtype) {70    const int64_t n = t.numel();71    const float* src = t.data<float>();72    std::vector<uint8_t> out(size_t(n) * dtype_size(dtype));73    if (dtype == DType::F32) {74        std::memcpy(out.data(), src, out.size());75    } else if (dtype == DType::F16) {76        f16_t* d = reinterpret_cast<f16_t*>(out.data());77        for (int64_t i = 0; i < n; ++i) d[i] = f16_t(src[i]);78    } else if (dtype == DType::BF16) {79        uint16_t* d = reinterpret_cast<uint16_t*>(out.data());80        for (int64_t i = 0; i < n; ++i) d[i] = float_to_bf16(src[i]);81    } else {82        die("save: unsupported storage dtype");83    }84    return out;85}8687std::vector<std::pair<std::string, Var>> dedupe(88    const std::vector<std::pair<std::string, Var>>& named) {89    std::vector<std::pair<std::string, Var>> out;90    std::unordered_set<const void*> seen;91    for (const auto& [name, p] : named) {92        if (!p.defined()) continue;93        if (seen.insert(p.id()).second) out.emplace_back(name, p);94    }95    return out;96}9798json load_json(const fs::path& p) {99    std::ifstream in(p);100    if (!in) die("cannot open " + p.string());101    return json::parse(in);102}103104} // namespace105106std::string save(const std::string& repo_dir, const std::string& config_json,107                 const std::vector<std::pair<std::string, Var>>& named_params,108                 const SaveOptions& opts) {109    const fs::path repo(repo_dir);110    fs::create_directories(repo / "manifests");111    fs::create_directories(repo / "objects");112113    // Parent manifest (if any): hash -> existing location, for delta reuse.114    json parent;115    std::string parent_name;116    const fs::path head = repo / "manifest-latest.json";117    if (fs::exists(head)) {118        parent = load_json(head);119        parent_name = parent.value("self", "");120    }121    std::unordered_map<std::string, json> known; // content hash -> tensor entry122    if (parent.contains("tensors"))123        for (const auto& [name, e] : parent.at("tensors").items())124            known.emplace(e.at("hash").get<std::string>(), e);125126    // Hash every (converted) tensor; split into reused vs new.127    struct Pending {128        std::string name;129        std::vector<uint8_t> bytes;130        std::string hash;131        std::vector<int64_t> shape;132    };133    json tensors = json::object();134    std::vector<Pending> fresh;135    size_t reused = 0, reused_bytes = 0;136    for (const auto& [name, p] : dedupe(named_params)) {137        Pending pd;138        pd.name = name;139        pd.shape = p.value().shape();140        pd.bytes = convert_out(p.value(), opts.dtype);141        pd.hash = hex64(fnv1a64(pd.bytes.data(), pd.bytes.size()));142        auto it = known.find(pd.hash);143        if (it != known.end()) {144            tensors[name] = it->second; // unchanged since parent: reference it145            ++reused;146            reused_bytes += pd.bytes.size();147        } else {148            fresh.push_back(std::move(pd));149        }150    }151152    // Pack new tensors into page-aligned shards capped at shard_mb.153    const size_t cap = size_t(opts.shard_mb) * 1024 * 1024;154    size_t si = 0;155    size_t written_bytes = 0;156    while (si < fresh.size()) {157        std::vector<uint8_t> blob(kPage, 0); // shard header page158        {159            uint32_t* h = reinterpret_cast<uint32_t*>(blob.data());160            h[0] = kShardMagic;161            h[1] = kVersion;162        }163        json members = json::array();164        uint64_t shard_hash = 1469598103934665603ull;165        size_t end = si;166        while (end < fresh.size()) {167            const size_t start = page_align(blob.size());168            const size_t next = page_align(start + fresh[end].bytes.size());169            if (end > si && next > cap + kPage) break; // shard full170            blob.resize(start, 0);171            blob.insert(blob.end(), fresh[end].bytes.begin(), fresh[end].bytes.end());172            members.push_back(json{{"i", end}, {"offset", start}});173            shard_hash ^= fnv1a64(fresh[end].hash.data(), fresh[end].hash.size());174            shard_hash *= 1099511628211ull;175            ++end;176        }177        blob.resize(page_align(blob.size()), 0);178179        const std::string shard_rel = "objects/" + hex64(shard_hash) + ".fshard";180        const fs::path shard_path = repo / shard_rel;181        if (!fs::exists(shard_path)) { // content-addressed: identical = skip182            const fs::path tmp = shard_path.string() + ".tmp";183            std::ofstream out(tmp, std::ios::binary);184            out.write(reinterpret_cast<const char*>(blob.data()),185                      std::streamsize(blob.size()));186            out.close();187            fs::rename(tmp, shard_path);188        }189        written_bytes += blob.size();190191        for (const auto& m : members) {192            const Pending& pd = fresh[m.at("i").get<size_t>()];193            tensors[pd.name] = json{194                {"dtype", dtype_name(opts.dtype)}, {"shape", pd.shape},195                {"hash", pd.hash},                 {"shard", shard_rel},196                {"offset", m.at("offset")},        {"nbytes", pd.bytes.size()}};197        }198        si = end;199    }200201    // Commit: numbered manifest + refreshed HEAD.202    int64_t seq = 0;203    for (const auto& e : fs::directory_iterator(repo / "manifests")) {204        (void)e;205        ++seq;206    }207    char mname[64];208    std::snprintf(mname, sizeof(mname), "manifest-%06lld.json",209                  static_cast<long long>(seq));210    json manifest{{"fmodel", kVersion},211                  {"self", mname},212                  {"parent", parent_name},213                  {"step", opts.step},214                  {"tag", opts.tag},215                  {"dtype", dtype_name(opts.dtype)},216                  {"config", json::parse(config_json.empty() ? "{}" : config_json)},217                  {"tensors", tensors}};218    const fs::path mpath = repo / "manifests" / mname;219    {220        std::ofstream out(mpath);221        out << manifest.dump(1);222    }223    fs::copy_file(mpath, head, fs::copy_options::overwrite_existing);224225    std::printf("fmodel: %s — %zu tensors reused (%.1f MB), %zu written (%.1f MB)\n",226                (repo / mname).c_str(), reused, double(reused_bytes) / 1e6,227                fresh.size(), double(written_bytes) / 1e6);228    return mpath.string();229}230231// ---- Snapshot ---------------------------------------------------------------232233struct Snapshot::Impl {234    struct Mapped {235        void* base = nullptr;236        size_t len = 0;237        MTL::Buffer* buffer = nullptr;238    };239    struct Entry {240        DType dtype;241        std::vector<int64_t> shape;242        std::string shard;243        size_t offset, nbytes;244    };245246    fs::path repo;247    std::string config;248    int64_t step = -1;249    std::vector<std::string> order;250    std::unordered_map<std::string, Entry> index;251    std::unordered_map<std::string, Mapped> maps; // shard rel-path -> mapping252253    ~Impl() {254        for (auto& [rel, m] : maps) {255            if (m.buffer) m.buffer->release();256            if (m.base) munmap(m.base, m.len);257        }258    }259260    const Mapped& mapped(const std::string& rel) {261        auto it = maps.find(rel);262        if (it != maps.end()) return it->second;263264        const fs::path p = repo / rel;265        const int fd = ::open(p.c_str(), O_RDONLY);266        if (fd < 0) die("cannot open shard " + p.string());267        struct stat st;268        if (fstat(fd, &st) != 0) die("fstat failed: " + p.string());269        const size_t len = size_t(st.st_size);270        if (len % kPage != 0) die("shard not page-padded: " + p.string());271272        // MAP_PRIVATE + RW: pages stay clean (backed by the file cache)273        // unless someone writes — and the read-only contract means nobody274        // does. newBuffer(bytesNoCopy) needs page alignment + page-multiple275        // length, which mmap and the writer guarantee.276        void* base = mmap(nullptr, len, PROT_READ | PROT_WRITE, MAP_PRIVATE, fd, 0);277        ::close(fd);278        if (base == MAP_FAILED) die("mmap failed: " + p.string());279        if (reinterpret_cast<const uint32_t*>(base)[0] != kShardMagic)280            die("bad shard magic: " + p.string());281282        MTL::Buffer* buf = Device::get().mtl()->newBuffer(283            base, len, MTL::ResourceStorageModeShared, nullptr);284        if (!buf) die("newBuffer(bytesNoCopy) failed: " + p.string());285        return maps.emplace(rel, Mapped{base, len, buf}).first->second;286    }287};288289Snapshot Snapshot::open(const std::string& path) {290    fs::path manifest(path);291    fs::path repo;292    if (fs::is_directory(manifest)) {293        repo = manifest;294        manifest = repo / "manifest-latest.json";295    } else {296        repo = manifest.parent_path();297        if (repo.filename() == "manifests") repo = repo.parent_path();298    }299    if (!fs::exists(manifest)) die("no manifest at " + manifest.string());300301    auto impl = std::make_shared<Impl>();302    impl->repo = repo;303    json j = load_json(manifest);304    impl->config = j.value("config", json::object()).dump();305    impl->step = j.value("step", int64_t(-1));306    for (const auto& [name, e] : j.at("tensors").items()) {307        Impl::Entry ent;308        ent.dtype = dtype_from_name(e.at("dtype").get<std::string>());309        ent.shape = e.at("shape").get<std::vector<int64_t>>();310        ent.shard = e.at("shard").get<std::string>();311        ent.offset = e.at("offset").get<size_t>();312        ent.nbytes = e.at("nbytes").get<size_t>();313        impl->order.push_back(name);314        impl->index.emplace(name, std::move(ent));315    }316317    Snapshot s;318    s.impl_ = std::move(impl);319    return s;320}321322const std::string& Snapshot::config_json() const { return impl_->config; }323int64_t Snapshot::step() const { return impl_->step; }324std::vector<std::string> Snapshot::names() const { return impl_->order; }325bool Snapshot::has(const std::string& name) const { return impl_->index.count(name); }326327const std::vector<int64_t>& Snapshot::shape(const std::string& name) const {328    auto it = impl_->index.find(name);329    if (it == impl_->index.end()) die("no tensor " + name);330    return it->second.shape;331}332333Tensor Snapshot::tensor_f32(const std::string& name) const {334    auto it = impl_->index.find(name);335    if (it == impl_->index.end()) die("no tensor " + name);336    const Impl::Entry& e = it->second;337    const Impl::Mapped& m = impl_->mapped(e.shard);338339    if (e.dtype == DType::F32)340        return Tensor::from_buffer(m.buffer, e.offset, e.shape); // zero-copy341342    // f16/bf16 storage: convert into a fresh f32 tensor.343    Tensor out = Tensor::empty(e.shape);344    const int64_t n = out.numel();345    float* dst = out.data<float>();346    const uint8_t* src = static_cast<const uint8_t*>(m.base) + e.offset;347    if (e.dtype == DType::F16) {348        const f16_t* s = reinterpret_cast<const f16_t*>(src);349        for (int64_t i = 0; i < n; ++i) dst[i] = float(s[i]);350    } else if (e.dtype == DType::BF16) {351        const uint16_t* s = reinterpret_cast<const uint16_t*>(src);352        for (int64_t i = 0; i < n; ++i) dst[i] = bf16_to_float(s[i]);353    } else {354        die("tensor_f32: unsupported stored dtype for " + name);355    }356    return out;357}358359} // namespace forge::fmodel360