spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#include "core/tensor.h"34#include "core/device.h"56#include <Metal/Metal.hpp>78#include <cstdio>9#include <cstdlib>10#include <cstring>11#include <sstream>1213namespace forge {1415namespace {16[[noreturn]] void die(const std::string& msg) {17    std::fprintf(stderr, "forge: %s\n", msg.c_str());18    std::abort();19}20} // namespace2122// Owns one pooled MTLBuffer; returns it to the pool when the last Tensor23// view drops it. External storages (fmodel mmaps) retain a caller-owned24// buffer instead and release it directly — the allocator never sees it.25struct Tensor::Storage {26    MTL::Buffer* buffer = nullptr;27    bool external = false;2829    explicit Storage(size_t nbytes) {30        buffer = Device::get().allocator().acquire(nbytes);31    }32    Storage(MTL::Buffer* ext, bool) : buffer(ext), external(true) {33        buffer->retain();34    }35    ~Storage() {36        if (external) buffer->release();37        else Device::get().allocator().release(buffer);38    }39    Storage(const Storage&) = delete;40    Storage& operator=(const Storage&) = delete;41};4243std::vector<int64_t> Tensor::contiguous_strides(const std::vector<int64_t>& shape) {44    std::vector<int64_t> strides(shape.size());45    int64_t acc = 1;46    for (int64_t i = int64_t(shape.size()) - 1; i >= 0; --i) {47        strides[size_t(i)] = acc;48        acc *= shape[size_t(i)];49    }50    return strides;51}5253Tensor Tensor::empty(std::vector<int64_t> shape, DType dtype) {54    Tensor t;55    t.dtype_ = dtype;56    t.shape_ = std::move(shape);57    t.strides_ = contiguous_strides(t.shape_);58    int64_t n = 1;59    for (int64_t d : t.shape_) {60        if (d <= 0) die("Tensor::empty: non-positive dim");61        n *= d;62    }63    t.storage_ = std::make_shared<Storage>(size_t(n) * dtype_size(dtype));64    return t;65}6667Tensor Tensor::from_buffer(MTL::Buffer* buffer, size_t byte_offset,68                           std::vector<int64_t> shape, DType dtype) {69    Tensor t;70    t.dtype_ = dtype;71    t.shape_ = std::move(shape);72    t.strides_ = contiguous_strides(t.shape_);73    if (byte_offset % dtype_size(dtype) != 0)74        die("Tensor::from_buffer: offset not aligned to dtype size");75    t.offset_ = int64_t(byte_offset / dtype_size(dtype));76    t.storage_ = std::make_shared<Storage>(buffer, true);77    return t;78}7980Tensor Tensor::zeros(std::vector<int64_t> shape, DType dtype) {81    Tensor t = empty(std::move(shape), dtype);82    std::memset(t.raw(), 0, t.nbytes());83    return t;84}8586Tensor Tensor::full(std::vector<int64_t> shape, float value, DType dtype) {87    Tensor t = empty(std::move(shape), dtype);88    t.fill_(value);89    return t;90}9192int64_t Tensor::size(int64_t dim) const {93    if (dim < 0) dim += ndim();94    if (dim < 0 || dim >= ndim()) die("Tensor::size: dim out of range");95    return shape_[size_t(dim)];96}9798int64_t Tensor::numel() const {99    int64_t n = 1;100    for (int64_t d : shape_) n *= d;101    return n;102}103104bool Tensor::is_contiguous() const {105    return strides_ == contiguous_strides(shape_);106}107108void* Tensor::raw() const {109    if (!storage_) die("Tensor::raw: undefined tensor");110    return static_cast<char*>(storage_->buffer->contents()) + size_t(offset_) * itemsize();111}112113MTL::Buffer* Tensor::buffer() const {114    if (!storage_) die("Tensor::buffer: undefined tensor");115    return storage_->buffer;116}117118size_t Tensor::buffer_offset() const {119    return size_t(offset_) * itemsize();120}121122Tensor Tensor::view(std::vector<int64_t> new_shape) const {123    if (!is_contiguous()) die("Tensor::view: tensor not contiguous");124    int64_t n = 1;125    for (int64_t d : new_shape) n *= d;126    if (n != numel()) die("Tensor::view: numel mismatch");127    Tensor t = *this;128    t.shape_ = std::move(new_shape);129    t.strides_ = contiguous_strides(t.shape_);130    return t;131}132133Tensor Tensor::slice0(int64_t start, int64_t len) const {134    if (!is_contiguous()) die("Tensor::slice0: tensor not contiguous");135    if (ndim() < 1 || start < 0 || len <= 0 || start + len > shape_[0])136        die("Tensor::slice0: range out of bounds");137    Tensor t = *this;138    t.shape_[0] = len;139    t.offset_ = offset_ + start * strides_[0];140    return t;141}142143void Tensor::fill_(float value) {144    const int64_t n = numel();145    switch (dtype_) {146        case DType::F32: {147            float* p = data<float>();148            for (int64_t i = 0; i < n; ++i) p[i] = value;149            break;150        }151        case DType::F16: {152            f16_t* p = data<f16_t>();153            const f16_t v = f16_t(value);154            for (int64_t i = 0; i < n; ++i) p[i] = v;155            break;156        }157        case DType::BF16: {158            uint16_t* p = data<uint16_t>();159            const uint16_t v = float_to_bf16(value);160            for (int64_t i = 0; i < n; ++i) p[i] = v;161            break;162        }163        case DType::U16: {164            uint16_t* p = data<uint16_t>();165            const uint16_t v = uint16_t(value);166            for (int64_t i = 0; i < n; ++i) p[i] = v;167            break;168        }169        case DType::I32: {170            int32_t* p = data<int32_t>();171            const int32_t v = int32_t(value);172            for (int64_t i = 0; i < n; ++i) p[i] = v;173            break;174        }175    }176}177178float Tensor::item_at(int64_t i) const {179    switch (dtype_) {180        case DType::F32:  return data<float>()[i];181        case DType::F16:  return float(data<f16_t>()[i]);182        case DType::BF16: return bf16_to_float(data<uint16_t>()[i]);183        case DType::U16:  return float(data<uint16_t>()[i]);184        case DType::I32:  return float(data<int32_t>()[i]);185    }186    return 0.0f;187}188189void Tensor::set_item(int64_t i, float value) {190    switch (dtype_) {191        case DType::F32:  data<float>()[i] = value; break;192        case DType::F16:  data<f16_t>()[i] = f16_t(value); break;193        case DType::BF16: data<uint16_t>()[i] = float_to_bf16(value); break;194        case DType::U16:  data<uint16_t>()[i] = uint16_t(value); break;195        case DType::I32:  data<int32_t>()[i] = int32_t(value); break;196    }197}198199std::string Tensor::describe() const {200    std::ostringstream os;201    os << "Tensor(" << dtype_name(dtype_) << ", [";202    for (size_t i = 0; i < shape_.size(); ++i) {203        if (i) os << ", ";204        os << shape_[i];205    }206    os << "])";207    return os.str();208}209210} // namespace forge211