spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#pragma once34#include "core/dtype.h"56#include <cstdint>7#include <memory>8#include <string>9#include <vector>1011namespace MTL {12class Buffer;13}1415namespace forge {1617// Refcounted view over a pooled MTLBuffer (MTLStorageModeShared), so the18// same memory is addressable from CPU code and GPU kernels with zero copies.19// Row-major, element strides. Views (reshape/view/slice) share storage.20class Tensor {21public:22    Tensor() = default;2324    static Tensor empty(std::vector<int64_t> shape, DType dtype = DType::F32);25    static Tensor zeros(std::vector<int64_t> shape, DType dtype = DType::F32);26    static Tensor full(std::vector<int64_t> shape, float value, DType dtype = DType::F32);2728    // Wrap an externally-owned MTLBuffer region (e.g. an mmapped .fmodel29    // shard bridged with newBuffer(bytesNoCopy)). The buffer is retained for30    // the storage's lifetime and released — NOT returned to the allocator31    // pool — when the last view drops. byte_offset must be a multiple of the32    // dtype size (fmodel aligns to the 16 KB page, far stricter).33    static Tensor from_buffer(MTL::Buffer* buffer, size_t byte_offset,34                              std::vector<int64_t> shape, DType dtype = DType::F32);3536    bool defined() const { return storage_ != nullptr; }37    DType dtype() const { return dtype_; }38    int64_t ndim() const { return int64_t(shape_.size()); }39    const std::vector<int64_t>& shape() const { return shape_; }40    const std::vector<int64_t>& strides() const { return strides_; }41    int64_t size(int64_t dim) const;42    int64_t numel() const;43    size_t itemsize() const { return dtype_size(dtype_); }44    size_t nbytes() const { return size_t(numel()) * itemsize(); }45    bool is_contiguous() const;4647    // CPU access. Valid at all times (unified memory) — but only touch it48    // when no in-flight command buffer may write the same storage.49    void* raw() const;50    template <typename T>51    T* data() const { return static_cast<T*>(raw()); }5253    // GPU access.54    MTL::Buffer* buffer() const;55    size_t buffer_offset() const; // bytes from the start of the MTLBuffer5657    // Views (no copy). Both require contiguous layout.58    Tensor view(std::vector<int64_t> new_shape) const;59    Tensor reshape(std::vector<int64_t> new_shape) const { return view(std::move(new_shape)); }60    // Slice along dim 0: rows [start, start+len). Contiguous only.61    Tensor slice0(int64_t start, int64_t len) const;6263    void fill_(float value);64    void zero_() { fill_(0.0f); }6566    // Read/write a single element as float, whatever the dtype (test/debug).67    float item_at(int64_t linear_index) const;68    void set_item(int64_t linear_index, float value);6970    std::string describe() const;7172private:73    struct Storage;7475    std::shared_ptr<Storage> storage_;76    std::vector<int64_t> shape_;77    std::vector<int64_t> strides_; // in elements78    int64_t offset_ = 0;           // in elements79    DType dtype_ = DType::F32;8081    static std::vector<int64_t> contiguous_strides(const std::vector<int64_t>& shape);82};8384} // namespace forge85