include/brotensor/tensor.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #pragma once | ||
| 2 | |||
| 3 | #include <cstddef> | ||
| 4 | #include <cstdint> | ||
| 5 | #include <string> | ||
| 6 | #include <vector> | ||
| 7 | |||
| 8 | namespace brotensor { | ||
| 9 | |||
| 10 | // ─── Dtype ───────────────────────────────────────────────────────────────── | ||
| 11 | // | ||
| 12 | // brotensor's tensor type carries a dtype tag so ops can pick the right | ||
| 13 | // kernel without a parallel tensor type per precision. Storage stays as a | ||
| 14 | // single raw `void*` (`data`); typed access is via the host_f32 / host_fp16 | ||
| 15 | // accessors. GPU backends reinterpret the same allocation for FP16 / INT8. | ||
| 16 | // | ||
| 17 | // Element sizes are fixed: FP32 = 4 bytes, FP16 = 2 bytes, BF16 = 2 bytes, | ||
| 18 | // INT8 = 1 byte, INT32 = 4 bytes. Allocation, clone, zero, and resize all use | ||
| 19 | // dtype-aware byte counts. BF16 (IEEE 754 bfloat16 — the high 16 bits of an | ||
| 20 | // FP32) is an arithmetic dtype carried only by the GPU backends; like FP16 it | ||
| 21 | // is stored as a uint16_t bit pattern on the host. Arithmetic ops dispatch on | ||
| 22 | // FP32/FP16/BF16. INT8 is currently only carried by weight-only quantised ops | ||
| 23 | // (W8A16 matmul/conv2d). INT32 is likewise a pure storage carrier — used for | ||
| 24 | // device-resident index/offset buffers (e.g. per-head offset tables for | ||
| 25 | // softmax_xent_fused_batched); no arithmetic op dispatches on it. | ||
| 26 | enum class Dtype : int { | ||
| 27 | FP32 = 0, | ||
| 28 | FP16 = 1, | ||
| 29 | INT8 = 2, | ||
| 30 | INT32 = 3, | ||
| 31 | BF16 = 4, | ||
| 32 | F64 = 5, | ||
| 33 | // GGUF legacy quants — 32-element blocks, opaque storage carriers only. | ||
| 34 | Q4_0 = 10, | ||
| 35 | Q4_1 = 11, | ||
| 36 | Q5_0 = 12, | ||
| 37 | Q5_1 = 13, | ||
| 38 | Q8_0 = 14, | ||
| 39 | Q8_1 = 15, | ||
| 40 | // GGUF K-quants — 256-element superblocks, opaque storage carriers only. | ||
| 41 | Q2_K = 20, | ||
| 42 | Q3_K = 21, | ||
| 43 | Q4_K = 22, | ||
| 44 | Q5_K = 23, | ||
| 45 | Q6_K = 24, | ||
| 46 | Q8_K = 25, | ||
| 47 | }; | ||
| 48 | |||
| 49 | // Bytes per scalar element. Returns 0 for quant dtypes (they aren't | ||
| 50 | // element-addressable — use dtype_storage_bytes() instead). | ||
| 51 | int dtype_size_bytes(Dtype); | ||
| 52 | |||
| 53 | // Elements per block. 1 for non-quant types; 32 for the legacy GGUF quants | ||
| 54 | // (Q4_0..Q8_1); 256 for the K-quants. | ||
| 55 | int dtype_block_size(Dtype); | ||
| 56 | |||
| 57 | // Bytes per block. Equals dtype_size_bytes(d) for non-quant dtypes; for quant | ||
| 58 | // dtypes it's the on-disk block size (e.g. Q4_K = 144). | ||
| 59 | int dtype_block_bytes(Dtype); | ||
| 60 | |||
| 61 | // Byte count for a tensor of `numel` elements stored as `d`. For non-quant | ||
| 62 | // types it's numel * dtype_size_bytes(d). For quant types `numel` must be a | ||
| 63 | // multiple of dtype_block_size(d) (throws std::runtime_error otherwise) and | ||
| 64 | // the result is (numel / block_size) * block_bytes. | ||
| 65 | std::size_t dtype_storage_bytes(Dtype d, std::int64_t numel); | ||
| 66 | |||
| 67 | // True iff `d` is a quant block carrier (Q*_*). | ||
| 68 | bool dtype_is_quant(Dtype); | ||
| 69 | |||
| 70 | // ─── Device ──────────────────────────────────────────────────────────────── | ||
| 71 | // | ||
| 72 | // Runtime backend tag carried on every Tensor. CPU is always available; | ||
| 73 | // CUDA / Metal are registered at runtime by `brotensor::init()` if the | ||
| 74 | // corresponding backend was compiled into this binary. Supports multi-GPU | ||
| 75 | // device indexing (e.g. Device::cuda(0), Device::cuda(1)). | ||
| 76 | enum class DeviceType : int { CPU = 0, CUDA = 1, Metal = 2 }; | ||
| 77 | |||
| 78 | struct Device { | ||
| 79 | 11 | DeviceType type = DeviceType::CPU; | |
| 80 | 11 | int index = 0; | |
| 81 | |||
| 82 | 33 | constexpr Device() = default; | |
| 83 | ✗ | constexpr Device(DeviceType t, int idx = 0) : type(t), index(idx) {} | |
| 84 | |||
| 85 | ✗ | static constexpr Device cpu() { return Device(DeviceType::CPU, 0); } | |
| 86 | ✗ | static constexpr Device cuda(int idx = 0) { return Device(DeviceType::CUDA, idx); } | |
| 87 | ✗ | static constexpr Device metal(int idx = 0) { return Device(DeviceType::Metal, idx); } | |
| 88 | |||
| 89 | // Legacy enum compatibility constants so `Device::CPU`, `Device::CUDA`, `Device::Metal` work verbatim. | ||
| 90 | static const Device CPU; | ||
| 91 | static const Device CUDA; | ||
| 92 | static const Device Metal; | ||
| 93 | |||
| 94 | 65931057 | constexpr bool operator==(const Device& o) const { | |
| 95 |
2/2✓ Branch 0 taken 7055 times.
✓ Branch 1 taken 65924002 times.
|
65931057 | return type == o.type && index == o.index; |
| 96 | } | ||
| 97 | 65921843 | constexpr bool operator!=(const Device& o) const { | |
| 98 | 65921843 | return !(*this == o); | |
| 99 | } | ||
| 100 | constexpr bool operator<(const Device& o) const { | ||
| 101 | if (type != o.type) return static_cast<int>(type) < static_cast<int>(o.type); | ||
| 102 | return index < o.index; | ||
| 103 | } | ||
| 104 | |||
| 105 | 18223 | constexpr bool is_cpu() const { return type == DeviceType::CPU; } | |
| 106 | 6630 | constexpr bool is_cuda() const { return type == DeviceType::CUDA; } | |
| 107 | 3294 | constexpr bool is_metal() const { return type == DeviceType::Metal; } | |
| 108 | constexpr bool is_gpu() const { return type != DeviceType::CPU; } | ||
| 109 | }; | ||
| 110 | |||
| 111 | const char* device_name(Device); | ||
| 112 | std::string to_string(Device); | ||
| 113 | |||
| 114 | // ─── Tensor ──────────────────────────────────────────────────────────────── | ||
| 115 | // | ||
| 116 | // Unified tensor: a row-major (rows, cols) buffer tagged with both a Dtype | ||
| 117 | // and a Device. Storage is a single opaque `void*` allocated through the | ||
| 118 | // backend's alloc vtable (see detail/dispatch.h); the destructor frees via | ||
| 119 | // the same vtable. Rank is fixed at 2 (matrix) or 1 (vector — cols == 1). | ||
| 120 | // | ||
| 121 | // Copyable and movable: the copy ctor / copy assignment perform a | ||
| 122 | // device-aware deep copy (identical to clone()); move transfers ownership | ||
| 123 | // of the underlying buffer. Copying a GPU-resident tensor therefore | ||
| 124 | // allocates and copies on-device — pass by reference on hot paths and use | ||
| 125 | // clone() where the copy should be explicit. The CPU backend allocates | ||
| 126 | // plain host memory through the same vtable interface so the storage layout | ||
| 127 | // is uniform across devices; for CPU tensors the typed host accessors | ||
| 128 | // (host_f32, host_fp16, at, to_host_vector) give ergonomic access without a | ||
| 129 | // device sync. | ||
| 130 | struct Tensor { | ||
| 131 | 81875 | void* data = nullptr; | |
| 132 | 81875 | int rows = 0; // rank-1 tensors: rows = N, cols = 1 | |
| 133 | 81875 | int cols = 0; | |
| 134 | 81875 | Dtype dtype = Dtype::FP32; | |
| 135 | 81875 | Device device = Device::CPU; | |
| 136 | |||
| 137 | 245625 | Tensor() = default; | |
| 138 | ~Tensor(); | ||
| 139 | |||
| 140 | // Copyable + movable. The copy ctor / copy assignment perform a | ||
| 141 | // device-aware deep copy — identical to clone() — so a Tensor can be | ||
| 142 | // used with value semantics (caches, std::vector storage, by-value | ||
| 143 | // params). clone() remains for call sites that want the copy to be | ||
| 144 | // explicit. Copying a GPU-resident tensor allocates + copies on-device. | ||
| 145 | Tensor(const Tensor&); | ||
| 146 | Tensor& operator=(const Tensor&); | ||
| 147 | Tensor(Tensor&&) noexcept; | ||
| 148 | Tensor& operator=(Tensor&&) noexcept; | ||
| 149 | |||
| 150 | // ─── Factories ───────────────────────────────────────────────────────── | ||
| 151 | // | ||
| 152 | // zeros / empty allocate on the current default device (see runtime.h — | ||
| 153 | // controlled by set_default_device() / DeviceScope, or the | ||
| 154 | // BROTENSOR_DEFAULT_DEVICE env var). `zeros` memset-zeros the buffer | ||
| 155 | // via the backend's memset_zero hook; `empty` leaves contents undefined. | ||
| 156 | static Tensor zeros(int r, int c, Dtype dt = Dtype::FP32); | ||
| 157 | static Tensor empty(int r, int c, Dtype dt = Dtype::FP32); | ||
| 158 | |||
| 159 | // Explicit-device variants — bypass the thread-local default. Useful for | ||
| 160 | // tests, multi-device pipelines, and any code that wants to pin storage | ||
| 161 | // to a specific backend regardless of caller policy. | ||
| 162 | static Tensor zeros_on(Device, int r, int c, Dtype dt = Dtype::FP32); | ||
| 163 | static Tensor empty_on(Device, int r, int c, Dtype dt = Dtype::FP32); | ||
| 164 | |||
| 165 | // Host (CPU) FP32 factories. Always allocate zero-filled storage pinned | ||
| 166 | // to Device::CPU regardless of the current default device — a parameter- | ||
| 167 | // bearing layer builds its weights on the host, then migrates the whole | ||
| 168 | // layer with to(Device). `mat` is a (rows, cols) matrix; `vec` is a | ||
| 169 | // rank-1 (n, 1) column vector. | ||
| 170 | 5874 | static Tensor mat(int r, int c) { return zeros_on(Device::CPU, r, c); } | |
| 171 | 1422 | static Tensor vec(int n) { return zeros_on(Device::CPU, n, 1); } | |
| 172 | |||
| 173 | // Host bootstrap. Allocates on the current default device and uploads | ||
| 174 | // `r * c` floats (FP32) or uint16_t bit patterns (FP16) from `src`. | ||
| 175 | // For non-CPU defaults this performs a host→device copy via the | ||
| 176 | // backend's memcpy_h2d hook; for the CPU default it's a plain memcpy. | ||
| 177 | static Tensor from_host(const float* src, int r, int c); | ||
| 178 | static Tensor from_host_fp16(const uint16_t* src, int r, int c); | ||
| 179 | static Tensor from_host_bf16(const uint16_t* src, int r, int c); | ||
| 180 | // INT8 weights (W8A16): `r * c` int8_t values, e.g. the output of | ||
| 181 | // quantize_int8_per_row_host paired with FP32 per-row dequant scales. | ||
| 182 | static Tensor from_host_int8(const int8_t* src, int r, int c); | ||
| 183 | |||
| 184 | // Variant that pins to a specific device, bypassing the default. | ||
| 185 | static Tensor from_host_on(Device, const float* src, int r, int c); | ||
| 186 | static Tensor from_host_fp16_on(Device, const uint16_t* src, int r, int c); | ||
| 187 | static Tensor from_host_bf16_on(Device, const uint16_t* src, int r, int c); | ||
| 188 | static Tensor from_host_int8_on(Device, const int8_t* src, int r, int c); | ||
| 189 | |||
| 190 | // Non-owning view over an existing backend-resident pointer. The | ||
| 191 | // returned tensor's destructor will NOT free `data`. Caller is | ||
| 192 | // responsible for lifetime. Mirrors the legacy GpuTensor::view pattern. | ||
| 193 | static Tensor view(Device, void* data, int rows, int cols, Dtype = Dtype::FP32); | ||
| 194 | |||
| 195 | // Dtype-agnostic host bootstrap: allocates on `target` and copies | ||
| 196 | // `nbytes` raw bytes from `src` — a plain memcpy for Device::CPU, a | ||
| 197 | // single memcpy_h2d otherwise. Unlike from_host*_on, this works for any | ||
| 198 | // Dtype including the opaque GGUF block-quant carriers, since it copies | ||
| 199 | // bytes() rather than interpreting elements. `nbytes` must equal the | ||
| 200 | // resulting tensor's bytes() (i.e. dtype_storage_bytes(dt, r*c)). | ||
| 201 | static Tensor from_raw_bytes_on(Device target, const void* src, | ||
| 202 | int r, int c, Dtype dt, | ||
| 203 | std::size_t nbytes); | ||
| 204 | |||
| 205 | // ─── Migration ───────────────────────────────────────────────────────── | ||
| 206 | |||
| 207 | // Returns a fresh tensor on `target` with the same shape/dtype/contents | ||
| 208 | // as `*this`. No-op clone() if already on the target device. The source | ||
| 209 | // tensor is unchanged. Uses the backend pair's memcpy_h2d / memcpy_d2h / | ||
| 210 | // memcpy_d2d hooks as appropriate. | ||
| 211 | Tensor to(Device target) const; | ||
| 212 | |||
| 213 | // Device-preserving deep copy. | ||
| 214 | Tensor clone() const; | ||
| 215 | |||
| 216 | // ─── Mutators ────────────────────────────────────────────────────────── | ||
| 217 | |||
| 218 | // memset-zero the buffer over bytes(). Dispatches through the backend's | ||
| 219 | // memset_zero hook. | ||
| 220 | void zero(); | ||
| 221 | |||
| 222 | // Reshapes to (r, c, dt); leaves contents undefined (call zero() | ||
| 223 | // afterwards if needed). Device is preserved. Storage is kept whenever | ||
| 224 | // the requested shape fits the existing allocation (capacity = the | ||
| 225 | // high-water mark of this tensor's past sizes), so a scratch buffer | ||
| 226 | // cycling through shapes stabilises at its largest size instead of | ||
| 227 | // reallocating every call — which also keeps its device pointer stable, | ||
| 228 | // a requirement for CUDA-graph-captured op sequences. Reallocates only | ||
| 229 | // when growing past capacity. A no-op if the shape and dtype already | ||
| 230 | // match. Throws std::runtime_error on a negative dimension, or if called | ||
| 231 | // on a non-owning view (a tensor from view()) — reshaping a view would | ||
| 232 | // silently sever it, so allocate a fresh tensor instead. | ||
| 233 | void resize(int r, int c, Dtype dt = Dtype::FP32); | ||
| 234 | |||
| 235 | // ─── Accessors ───────────────────────────────────────────────────────── | ||
| 236 | |||
| 237 | 50365431 | int size() const { return rows * cols; } | |
| 238 | std::size_t bytes() const; | ||
| 239 | 1 | bool is_host() const { return device == Device::CPU; } | |
| 240 |
2/2✓ Branch 0 taken 4 times.
✓ Branch 1 taken 4423 times.
|
4427 | bool empty() const { return data == nullptr || size() == 0; } |
| 241 | |||
| 242 | // Host-side typed accessors. Throw std::runtime_error if device != CPU. | ||
| 243 | // `host_f32` additionally throws if dtype != FP32; `host_fp16` if | ||
| 244 | // dtype != FP16. `host_raw` is dtype-agnostic. | ||
| 245 | float* host_f32_mut(); | ||
| 246 | const float* host_f32() const; | ||
| 247 | uint16_t* host_fp16_mut(); | ||
| 248 | const uint16_t* host_fp16() const; | ||
| 249 | uint16_t* host_bf16_mut(); | ||
| 250 | const uint16_t* host_bf16() const; | ||
| 251 | void* host_raw_mut(); | ||
| 252 | const void* host_raw() const; | ||
| 253 | |||
| 254 | // Element access helpers (host-only, FP32-only — convenience for tests). | ||
| 255 | // Throw if device != CPU or dtype != FP32 or indices out of range. | ||
| 256 | float& at(int r, int c); | ||
| 257 | float at(int r, int c) const; | ||
| 258 | |||
| 259 | // Host (CPU) FP32 convenience accessors. Thin aliases over the typed | ||
| 260 | // host accessors above — they throw via the same checks if device != CPU | ||
| 261 | // or dtype != FP32. `ptr` is the raw row-major base pointer; operator() | ||
| 262 | // is bounds-checked (r, c) access; operator[] is flat element access. | ||
| 263 | 46841205 | float* ptr() { return host_f32_mut(); } | |
| 264 | const float* ptr() const { return host_f32(); } | ||
| 265 | 7985 | float& operator()(int r, int c) { return at(r, c); } | |
| 266 | 1024 | float operator()(int r, int c) const { return at(r, c); } | |
| 267 | 3955087 | float& operator[](int i) { return host_f32_mut()[i]; } | |
| 268 | 12961814 | float operator[](int i) const { return host_f32()[i]; } | |
| 269 | |||
| 270 | // ─── Host roundtrip helpers ──────────────────────────────────────────── | ||
| 271 | // | ||
| 272 | // `to_host_vector*` downloads (if on a GPU backend) and returns a | ||
| 273 | // std::vector containing the buffer's contents in the matching scalar | ||
| 274 | // type. The copy_to_host variants write into a caller-supplied buffer | ||
| 275 | // of at least size() elements. | ||
| 276 | std::vector<float> to_host_vector() const; // FP32 only | ||
| 277 | std::vector<uint16_t> to_host_vector_fp16() const; // FP16 only | ||
| 278 | std::vector<uint16_t> to_host_vector_bf16() const; // BF16 only | ||
| 279 | void copy_to_host(float* dst) const; // FP32 only | ||
| 280 | void copy_to_host_fp16(uint16_t* dst) const; // FP16 only | ||
| 281 | void copy_to_host_bf16(uint16_t* dst) const; // BF16 only | ||
| 282 | |||
| 283 | private: | ||
| 284 | 81875 | bool owns_ = false; | |
| 285 | // Bytes actually allocated behind `data` when owns_ is true — resize() | ||
| 286 | // keeps the existing storage whenever the requested size fits, so the | ||
| 287 | // capacity is the high-water mark of past sizes. 0 for views, released, | ||
| 288 | // and default-constructed tensors. | ||
| 289 | 81875 | std::size_t cap_bytes_ = 0; | |
| 290 | void release_(); | ||
| 291 | }; | ||
| 292 | |||
| 293 | // ─── FP16 / BF16 ↔ FP32 host-side conversion helpers ─────────────────────── | ||
| 294 | // | ||
| 295 | // Pure-CPU conversion. `fp16` is IEEE 754 binary16; `bf16` is bfloat16 — the | ||
| 296 | // high 16 bits of an FP32 with round-to-nearest-even. Useful for tests and | ||
| 297 | // small preprocessing where a GPU roundtrip would be wasteful. Not intended | ||
| 298 | // for hot loops. | ||
| 299 | uint16_t fp32_to_fp16_bits(float v); | ||
| 300 | float fp16_bits_to_fp32(uint16_t bits); | ||
| 301 | uint16_t fp32_to_bf16_bits(float v); | ||
| 302 | float bf16_bits_to_fp32(uint16_t bits); | ||
| 303 | |||
| 304 | } // namespace brotensor | ||
| 305 | |||
| 306 | namespace std { | ||
| 307 | template <> | ||
| 308 | struct hash<brotensor::Device> { | ||
| 309 | std::size_t operator()(const brotensor::Device& d) const noexcept { | ||
| 310 | return (static_cast<std::size_t>(d.type) << 16) | (static_cast<std::size_t>(d.index) & 0xFFFF); | ||
| 311 | } | ||
| 312 | }; | ||
| 313 | } | ||
| 314 |