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 | // FP8 E4M3 (OCP "e4m3fn": 1 sign, 4 exponent bits bias 7, 3 mantissa | ||
| 34 | // bits, max 448, no infinities), 1 byte. Like INT8 a storage carrier for | ||
| 35 | // weight-only quantised ops (linear_forward_batched_wq_ex). | ||
| 36 | FP8_E4M3 = 6, | ||
| 37 | // GGUF legacy quants — 32-element blocks, opaque storage carriers only. | ||
| 38 | Q4_0 = 10, | ||
| 39 | Q4_1 = 11, | ||
| 40 | Q5_0 = 12, | ||
| 41 | Q5_1 = 13, | ||
| 42 | Q8_0 = 14, | ||
| 43 | Q8_1 = 15, | ||
| 44 | // GGUF K-quants — 256-element superblocks, opaque storage carriers only. | ||
| 45 | Q2_K = 20, | ||
| 46 | Q3_K = 21, | ||
| 47 | Q4_K = 22, | ||
| 48 | Q5_K = 23, | ||
| 49 | Q6_K = 24, | ||
| 50 | Q8_K = 25, | ||
| 51 | }; | ||
| 52 | |||
| 53 | // Bytes per scalar element. Returns 0 for quant dtypes (they aren't | ||
| 54 | // element-addressable — use dtype_storage_bytes() instead). | ||
| 55 | int dtype_size_bytes(Dtype); | ||
| 56 | |||
| 57 | // Elements per block. 1 for non-quant types; 32 for the legacy GGUF quants | ||
| 58 | // (Q4_0..Q8_1); 256 for the K-quants. | ||
| 59 | int dtype_block_size(Dtype); | ||
| 60 | |||
| 61 | // Bytes per block. Equals dtype_size_bytes(d) for non-quant dtypes; for quant | ||
| 62 | // dtypes it's the on-disk block size (e.g. Q4_K = 144). | ||
| 63 | int dtype_block_bytes(Dtype); | ||
| 64 | |||
| 65 | // Byte count for a tensor of `numel` elements stored as `d`. For non-quant | ||
| 66 | // types it's numel * dtype_size_bytes(d). For quant types `numel` must be a | ||
| 67 | // multiple of dtype_block_size(d) (throws std::runtime_error otherwise) and | ||
| 68 | // the result is (numel / block_size) * block_bytes. | ||
| 69 | std::size_t dtype_storage_bytes(Dtype d, std::int64_t numel); | ||
| 70 | |||
| 71 | // True iff `d` is a quant block carrier (Q*_*). | ||
| 72 | bool dtype_is_quant(Dtype); | ||
| 73 | |||
| 74 | // Human-readable name for a Dtype (e.g. "FP32", "Q4_K"). | ||
| 75 | const char* dtype_name(Dtype); | ||
| 76 | |||
| 77 | // ─── Device ──────────────────────────────────────────────────────────────── | ||
| 78 | // | ||
| 79 | // Runtime backend tag carried on every Tensor. CPU is always available; | ||
| 80 | // CUDA / Metal / Vulkan are registered at runtime by `brotensor::init()` | ||
| 81 | // if the corresponding backend was compiled into this binary. Supports | ||
| 82 | // multi-GPU device indexing (e.g. Device::cuda(0), Device::cuda(1)). | ||
| 83 | // (3 was HIP, removed; the values are kept stable.) | ||
| 84 | enum class DeviceType : int { CPU = 0, CUDA = 1, Metal = 2, VULKAN = 4 }; | ||
| 85 | |||
| 86 | struct Device { | ||
| 87 | 11 | DeviceType type = DeviceType::CPU; | |
| 88 | 11 | int index = 0; | |
| 89 | |||
| 90 | 33 | constexpr Device() = default; | |
| 91 | 16 | constexpr Device(DeviceType t, int idx = 0) : type(t), index(idx) {} | |
| 92 | |||
| 93 | 6 | static constexpr Device cpu() { return Device(DeviceType::CPU, 0); } | |
| 94 | 1 | static constexpr Device cuda(int idx = 0) { return Device(DeviceType::CUDA, idx); } | |
| 95 | ✗ | static constexpr Device metal(int idx = 0) { return Device(DeviceType::Metal, idx); } | |
| 96 | 1 | static constexpr Device vulkan(int idx = 0) { return Device(DeviceType::VULKAN, idx); } | |
| 97 | |||
| 98 | // Legacy enum compatibility constants so `Device::CPU`, `Device::CUDA`, `Device::Metal`, `Device::VULKAN` work verbatim. | ||
| 99 | static const Device CPU; | ||
| 100 | static const Device CUDA; | ||
| 101 | static const Device Metal; | ||
| 102 | static const Device VULKAN; | ||
| 103 | |||
| 104 | 32636642248 | constexpr bool operator==(const Device& o) const { | |
| 105 |
2/2✓ Branch 0 taken 12890 times.
✓ Branch 1 taken 32636629358 times.
|
32636642248 | return type == o.type && index == o.index; |
| 106 | } | ||
| 107 | 32636625407 | constexpr bool operator!=(const Device& o) const { | |
| 108 | 32636625407 | return !(*this == o); | |
| 109 | } | ||
| 110 | constexpr bool operator<(const Device& o) const { | ||
| 111 | if (type != o.type) return static_cast<int>(type) < static_cast<int>(o.type); | ||
| 112 | return index < o.index; | ||
| 113 | } | ||
| 114 | |||
| 115 | 29349 | constexpr bool is_cpu() const { return type == DeviceType::CPU; } | |
| 116 | 563671 | constexpr bool is_cuda() const { return type == DeviceType::CUDA; } | |
| 117 | 5070 | constexpr bool is_metal() const { return type == DeviceType::Metal; } | |
| 118 | 151 | constexpr bool is_vulkan() const { return type == DeviceType::VULKAN; } | |
| 119 | ✗ | constexpr bool is_gpu() const { return type != DeviceType::CPU; } | |
| 120 | }; | ||
| 121 | |||
| 122 | const char* device_name(Device); | ||
| 123 | std::string to_string(Device); | ||
| 124 | |||
| 125 | // ─── Tensor ──────────────────────────────────────────────────────────────── | ||
| 126 | // | ||
| 127 | // Unified tensor: a row-major (rows, cols) buffer tagged with both a Dtype | ||
| 128 | // and a Device. Storage is a single opaque `void*` allocated through the | ||
| 129 | // backend's alloc vtable (see detail/dispatch.h); the destructor frees via | ||
| 130 | // the same vtable. Rank is fixed at 2 (matrix) or 1 (vector — cols == 1). | ||
| 131 | // | ||
| 132 | // Copyable and movable: the copy ctor / copy assignment perform a | ||
| 133 | // device-aware deep copy (identical to clone()); move transfers ownership | ||
| 134 | // of the underlying buffer. Copying a GPU-resident tensor therefore | ||
| 135 | // allocates and copies on-device — pass by reference on hot paths and use | ||
| 136 | // clone() where the copy should be explicit. The CPU backend allocates | ||
| 137 | // plain host memory through the same vtable interface so the storage layout | ||
| 138 | // is uniform across devices; for CPU tensors the typed host accessors | ||
| 139 | // (host_f32, host_fp16, at, to_host_vector) give ergonomic access without a | ||
| 140 | // device sync. | ||
| 141 | struct Tensor { | ||
| 142 | 104662 | void* data = nullptr; | |
| 143 | 104662 | int rows = 0; // rank-1 tensors: rows = N, cols = 1 | |
| 144 | 104662 | int cols = 0; | |
| 145 | 104662 | Dtype dtype = Dtype::FP32; | |
| 146 | 104662 | Device device = Device::CPU; | |
| 147 | |||
| 148 | // Opaque outside brotensor::jit, which is the only code that ever sets or | ||
| 149 | // reads it. -1 on every ordinary tensor. | ||
| 150 | // | ||
| 151 | // >= 0 means this Tensor is a *symbolic* traced intermediate standing for | ||
| 152 | // node `jit_slot` of the trace currently recording: `data` is null, it | ||
| 153 | // owns nothing and frees nothing, and it carries only rows / cols / dtype | ||
| 154 | // / device. A fused kernel evaluates such a value in registers and never | ||
| 155 | // reads it back out of memory, so the tracer allocates nothing for it; | ||
| 156 | // end_trace() hands a real buffer to exactly those symbolic tensors the | ||
| 157 | // caller is still holding when the trace closes. Copying one throws — see | ||
| 158 | // the copy ctor. | ||
| 159 | 104662 | int jit_slot = -1; | |
| 160 | |||
| 161 | 313986 | Tensor() = default; | |
| 162 | ~Tensor(); | ||
| 163 | |||
| 164 | // Copyable + movable. The copy ctor / copy assignment perform a | ||
| 165 | // device-aware deep copy — identical to clone() — so a Tensor can be | ||
| 166 | // used with value semantics (caches, std::vector storage, by-value | ||
| 167 | // params). clone() remains for call sites that want the copy to be | ||
| 168 | // explicit. Copying a GPU-resident tensor allocates + copies on-device. | ||
| 169 | // | ||
| 170 | // Copying a symbolic traced intermediate (jit_slot >= 0) throws | ||
| 171 | // std::runtime_error: it stands for a node of a trace and owns no memory, | ||
| 172 | // so there is nothing to deep-copy, and two Tensors cannot both become | ||
| 173 | // the one buffer end_trace() will allocate for that node. | ||
| 174 | Tensor(const Tensor&); | ||
| 175 | Tensor& operator=(const Tensor&); | ||
| 176 | Tensor(Tensor&&) noexcept; | ||
| 177 | Tensor& operator=(Tensor&&) noexcept; | ||
| 178 | |||
| 179 | // ─── Factories ───────────────────────────────────────────────────────── | ||
| 180 | // | ||
| 181 | // zeros / empty allocate on the current default device (see runtime.h — | ||
| 182 | // controlled by set_default_device() / DeviceScope, or the | ||
| 183 | // BROTENSOR_DEFAULT_DEVICE env var). `zeros` memset-zeros the buffer | ||
| 184 | // via the backend's memset_zero hook; `empty` leaves contents undefined. | ||
| 185 | static Tensor zeros(int r, int c, Dtype dt = Dtype::FP32); | ||
| 186 | static Tensor empty(int r, int c, Dtype dt = Dtype::FP32); | ||
| 187 | |||
| 188 | // Explicit-device variants — bypass the thread-local default. Useful for | ||
| 189 | // tests, multi-device pipelines, and any code that wants to pin storage | ||
| 190 | // to a specific backend regardless of caller policy. | ||
| 191 | static Tensor zeros_on(Device, int r, int c, Dtype dt = Dtype::FP32); | ||
| 192 | static Tensor empty_on(Device, int r, int c, Dtype dt = Dtype::FP32); | ||
| 193 | |||
| 194 | // Host (CPU) FP32 factories. Always allocate zero-filled storage pinned | ||
| 195 | // to Device::CPU regardless of the current default device — a parameter- | ||
| 196 | // bearing layer builds its weights on the host, then migrates the whole | ||
| 197 | // layer with to(Device). `mat` is a (rows, cols) matrix; `vec` is a | ||
| 198 | // rank-1 (n, 1) column vector. | ||
| 199 | 7294 | static Tensor mat(int r, int c) { return zeros_on(Device::CPU, r, c); } | |
| 200 | 1498 | static Tensor vec(int n) { return zeros_on(Device::CPU, n, 1); } | |
| 201 | |||
| 202 | // Host bootstrap. Allocates on the current default device and uploads | ||
| 203 | // `r * c` floats (FP32) or uint16_t bit patterns (FP16) from `src`. | ||
| 204 | // For non-CPU defaults this performs a host→device copy via the | ||
| 205 | // backend's memcpy_h2d hook; for the CPU default it's a plain memcpy. | ||
| 206 | static Tensor from_host(const float* src, int r, int c); | ||
| 207 | static Tensor from_host_fp16(const uint16_t* src, int r, int c); | ||
| 208 | static Tensor from_host_bf16(const uint16_t* src, int r, int c); | ||
| 209 | // INT8 weights (W8A16): `r * c` int8_t values, e.g. the output of | ||
| 210 | // quantize_int8_per_row_host paired with FP32 per-row dequant scales. | ||
| 211 | static Tensor from_host_int8(const int8_t* src, int r, int c); | ||
| 212 | |||
| 213 | // Variant that pins to a specific device, bypassing the default. | ||
| 214 | static Tensor from_host_on(Device, const float* src, int r, int c); | ||
| 215 | static Tensor from_host_fp16_on(Device, const uint16_t* src, int r, int c); | ||
| 216 | static Tensor from_host_bf16_on(Device, const uint16_t* src, int r, int c); | ||
| 217 | static Tensor from_host_int8_on(Device, const int8_t* src, int r, int c); | ||
| 218 | |||
| 219 | // Non-owning view over an existing backend-resident pointer. The | ||
| 220 | // returned tensor's destructor will NOT free `data`. Caller is | ||
| 221 | // responsible for lifetime. Mirrors the legacy GpuTensor::view pattern. | ||
| 222 | static Tensor view(Device, void* data, int rows, int cols, Dtype = Dtype::FP32); | ||
| 223 | |||
| 224 | // Dtype-agnostic host bootstrap: allocates on `target` and copies | ||
| 225 | // `nbytes` raw bytes from `src` — a plain memcpy for Device::CPU, a | ||
| 226 | // single memcpy_h2d otherwise. Unlike from_host*_on, this works for any | ||
| 227 | // Dtype including the opaque GGUF block-quant carriers, since it copies | ||
| 228 | // bytes() rather than interpreting elements. `nbytes` must equal the | ||
| 229 | // resulting tensor's bytes() (i.e. dtype_storage_bytes(dt, r*c)). | ||
| 230 | static Tensor from_raw_bytes_on(Device target, const void* src, | ||
| 231 | int r, int c, Dtype dt, | ||
| 232 | std::size_t nbytes); | ||
| 233 | |||
| 234 | // ─── Migration ───────────────────────────────────────────────────────── | ||
| 235 | |||
| 236 | // Returns a fresh tensor on `target` with the same shape/dtype/contents | ||
| 237 | // as `*this`. No-op clone() if already on the target device. The source | ||
| 238 | // tensor is unchanged. Uses the backend pair's memcpy_h2d / memcpy_d2h / | ||
| 239 | // memcpy_d2d hooks as appropriate. | ||
| 240 | Tensor to(Device target) const; | ||
| 241 | |||
| 242 | // Device-preserving deep copy. | ||
| 243 | Tensor clone() const; | ||
| 244 | |||
| 245 | // ─── Mutators ────────────────────────────────────────────────────────── | ||
| 246 | |||
| 247 | // memset-zero the buffer over bytes(). Dispatches through the backend's | ||
| 248 | // memset_zero hook. | ||
| 249 | void zero(); | ||
| 250 | |||
| 251 | // Reshapes to (r, c, dt); leaves contents undefined (call zero() | ||
| 252 | // afterwards if needed). Device is preserved. Storage is kept whenever | ||
| 253 | // the requested shape fits the existing allocation (capacity = the | ||
| 254 | // high-water mark of this tensor's past sizes), so a scratch buffer | ||
| 255 | // cycling through shapes stabilises at its largest size instead of | ||
| 256 | // reallocating every call — which also keeps its device pointer stable, | ||
| 257 | // a requirement for CUDA-graph-captured op sequences. Reallocates only | ||
| 258 | // when growing past capacity. A no-op if the shape and dtype already | ||
| 259 | // match. Throws std::runtime_error on a negative dimension, or if called | ||
| 260 | // on a non-owning view (a tensor from view()) — reshaping a view would | ||
| 261 | // silently sever it, so allocate a fresh tensor instead. | ||
| 262 | void resize(int r, int c, Dtype dt = Dtype::FP32); | ||
| 263 | |||
| 264 | // ─── Accessors ───────────────────────────────────────────────────────── | ||
| 265 | |||
| 266 | 865044215 | int size() const { return rows * cols; } | |
| 267 | std::size_t bytes() const; | ||
| 268 | std::size_t row_bytes() const; | ||
| 269 | 3 | bool is_host() const { return device == Device::CPU; } | |
| 270 |
2/2✓ Branch 0 taken 5 times.
✓ Branch 1 taken 5182 times.
|
5187 | bool empty() const { return data == nullptr || size() == 0; } |
| 271 | |||
| 272 | // Host-side typed accessors. Throw std::runtime_error if device != CPU. | ||
| 273 | // `host_f32` additionally throws if dtype != FP32; `host_fp16` if | ||
| 274 | // dtype != FP16. `host_raw` is dtype-agnostic. | ||
| 275 | float* host_f32_mut(); | ||
| 276 | const float* host_f32() const; | ||
| 277 | uint16_t* host_fp16_mut(); | ||
| 278 | const uint16_t* host_fp16() const; | ||
| 279 | uint16_t* host_bf16_mut(); | ||
| 280 | const uint16_t* host_bf16() const; | ||
| 281 | void* host_raw_mut(); | ||
| 282 | const void* host_raw() const; | ||
| 283 | |||
| 284 | // Element access helpers (host-only, FP32-only — convenience for tests). | ||
| 285 | // Throw if device != CPU or dtype != FP32 or indices out of range. | ||
| 286 | float& at(int r, int c); | ||
| 287 | float at(int r, int c) const; | ||
| 288 | |||
| 289 | // Host (CPU) FP32 convenience accessors. Thin aliases over the typed | ||
| 290 | // host accessors above — they throw via the same checks if device != CPU | ||
| 291 | // or dtype != FP32. `ptr` is the raw row-major base pointer; operator() | ||
| 292 | // is bounds-checked (r, c) access; operator[] is flat element access. | ||
| 293 | 318339407 | float* ptr() { return host_f32_mut(); } | |
| 294 | 1216902 | const float* ptr() const { return host_f32(); } | |
| 295 | 53500 | float& operator()(int r, int c) { return at(r, c); } | |
| 296 | 42232 | float operator()(int r, int c) const { return at(r, c); } | |
| 297 | 3962885 | float& operator[](int i) { return host_f32_mut()[i]; } | |
| 298 | 2472537432 | float operator[](int i) const { return host_f32()[i]; } | |
| 299 | |||
| 300 | // ─── Host roundtrip helpers ──────────────────────────────────────────── | ||
| 301 | // | ||
| 302 | // `to_host_vector*` downloads (if on a GPU backend) and returns a | ||
| 303 | // std::vector containing the buffer's contents in the matching scalar | ||
| 304 | // type. The copy_to_host variants write into a caller-supplied buffer | ||
| 305 | // of at least size() elements. | ||
| 306 | std::vector<float> to_host_vector() const; // FP32 only | ||
| 307 | std::vector<uint16_t> to_host_vector_fp16() const; // FP16 only | ||
| 308 | std::vector<uint16_t> to_host_vector_bf16() const; // BF16 only | ||
| 309 | void copy_to_host(float* dst) const; // FP32 only | ||
| 310 | void copy_to_host_fp16(uint16_t* dst) const; // FP16 only | ||
| 311 | void copy_to_host_bf16(uint16_t* dst) const; // BF16 only | ||
| 312 | |||
| 313 | // Dtype-agnostic raw byte transfers against the tensor's EXISTING storage: | ||
| 314 | // no reallocation, so the device pointer stays put — how a caller feeds | ||
| 315 | // new inputs to (and reads outputs from) a CUDA-graph-captured op | ||
| 316 | // sequence. Copies the first `nbytes` bytes; nbytes must be <= bytes(). | ||
| 317 | void copy_from_host_raw(const void* src, std::size_t nbytes); | ||
| 318 | void copy_to_host_raw(void* dst, std::size_t nbytes) const; | ||
| 319 | |||
| 320 | private: | ||
| 321 | 104662 | bool owns_ = false; | |
| 322 | // Bytes actually allocated behind `data` when owns_ is true — resize() | ||
| 323 | // keeps the existing storage whenever the requested size fits, so the | ||
| 324 | // capacity is the high-water mark of past sizes. 0 for views, released, | ||
| 325 | // and default-constructed tensors. | ||
| 326 | 104662 | std::size_t cap_bytes_ = 0; | |
| 327 | void release_(); | ||
| 328 | }; | ||
| 329 | |||
| 330 | // ─── jit symbolic-tensor tracking hook ───────────────────────────────────── | ||
| 331 | // | ||
| 332 | // brotensor::jit has to be able to find the caller's Tensor object again when | ||
| 333 | // a trace closes, because a value that is still live then must end up owning | ||
| 334 | // a real buffer the caller can read. It keeps a slot -> Tensor* table, and | ||
| 335 | // Tensor's move ctor / move assignment / destructor keep that table honest | ||
| 336 | // through this one function pointer. | ||
| 337 | // | ||
| 338 | // It is null unless a trace has been started on this process, so an ordinary | ||
| 339 | // tensor's lifetime costs a single `jit_slot >= 0` test and nothing else, and | ||
| 340 | // tensor.cpp gains no dependency on the jit. `to == nullptr` means the | ||
| 341 | // symbolic tensor at `from` is gone; otherwise it moved from `from` to `to`. | ||
| 342 | namespace detail { | ||
| 343 | using JitSlotRetargetHook = void (*)(int slot, const Tensor* from, Tensor* to); | ||
| 344 | using JitClampInplaceHook = bool (*)(Tensor& y, float lo, float hi); | ||
| 345 | using JitCastOutHook = bool (*)(const Tensor& src, Tensor& dst, Dtype out_dtype); | ||
| 346 | extern JitSlotRetargetHook jit_slot_retarget; | ||
| 347 | extern JitClampInplaceHook jit_clamp_inplace_hook; | ||
| 348 | extern JitCastOutHook jit_cast_out_hook; | ||
| 349 | } // namespace detail | ||
| 350 | |||
| 351 | // ─── Allocation accounting ───────────────────────────────────────────────── | ||
| 352 | // | ||
| 353 | // Monotonic, process-wide counters over every Tensor storage allocation that | ||
| 354 | // goes through a backend allocator — every zeros/empty/from_host/clone/to and | ||
| 355 | // every growing resize(), on any device. Views and no-op resizes do not count. | ||
| 356 | // | ||
| 357 | // This exists so a test can assert what a piece of code did *not* allocate: | ||
| 358 | // the trace JIT's contract is that a fused expression whose result is store()d | ||
| 359 | // into a caller buffer allocates nothing at all for its intermediates, and the | ||
| 360 | // only way to hold it to that is to read the counter on both sides. | ||
| 361 | 4 | struct AllocStats { | |
| 362 | 4 | std::uint64_t count = 0; // allocations performed | |
| 363 | 4 | std::uint64_t bytes = 0; // bytes requested across them | |
| 364 | }; | ||
| 365 | AllocStats alloc_stats(); | ||
| 366 | |||
| 367 | // ─── FP16 / BF16 ↔ FP32 host-side conversion helpers ─────────────────────── | ||
| 368 | // | ||
| 369 | // Pure-CPU conversion. `fp16` is IEEE 754 binary16; `bf16` is bfloat16 — the | ||
| 370 | // high 16 bits of an FP32 with round-to-nearest-even. Useful for tests and | ||
| 371 | // small preprocessing where a GPU roundtrip would be wasteful. Not intended | ||
| 372 | // for hot loops. | ||
| 373 | uint16_t fp32_to_fp16_bits(float v); | ||
| 374 | float fp16_bits_to_fp32(uint16_t bits); | ||
| 375 | uint16_t fp32_to_bf16_bits(float v); | ||
| 376 | float bf16_bits_to_fp32(uint16_t bits); | ||
| 377 | // FP8 E4M3 (Dtype::FP8_E4M3): round to nearest even, saturating to +-448 | ||
| 378 | // (NaN stays NaN, 0x7f); decode is exact. | ||
| 379 | uint8_t fp32_to_fp8_e4m3_bits(float v); | ||
| 380 | float fp8_e4m3_bits_to_fp32(uint8_t bits); | ||
| 381 | |||
| 382 | } // namespace brotensor | ||
| 383 | |||
| 384 | namespace std { | ||
| 385 | template <> | ||
| 386 | struct hash<brotensor::Device> { | ||
| 387 | std::size_t operator()(const brotensor::Device& d) const noexcept { | ||
| 388 | return (static_cast<std::size_t>(d.type) << 16) | (static_cast<std::size_t>(d.index) & 0xFFFF); | ||
| 389 | } | ||
| 390 | }; | ||
| 391 | } | ||
| 392 |