include/brotensor/detail/dispatch.h
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | #pragma once | ||
| 2 | |||
| 3 | #include "../tensor.h" | ||
| 4 | #include "op_table.h" | ||
| 5 | |||
| 6 | #include <cstddef> | ||
| 7 | #include <cstdint> | ||
| 8 | #include <vector> | ||
| 9 | |||
| 10 | namespace brotensor::detail { | ||
| 11 | |||
| 12 | // ─── OpsVTable ───────────────────────────────────────────────────────────── | ||
| 13 | // | ||
| 14 | // One function-pointer slot per public op. The slot signature matches the | ||
| 15 | // public op signature in <brotensor/ops.h> exactly. Generated from the | ||
| 16 | // single op list in detail/op_table.h so the public surface, the wrappers | ||
| 17 | // in src/ops.cpp, and each backend's registration table stay in sync. | ||
| 18 | // | ||
| 19 | // A null slot means "this backend does not implement this op". The | ||
| 20 | // dispatcher (`brotensor::detail::dispatch(...)` callers, or the wrapper | ||
| 21 | // itself) checks for null and throws a "not implemented on this backend" | ||
| 22 | // std::runtime_error. | ||
| 23 | |||
| 24 | #define BROTENSOR_VTABLE_FIELD(name, ret, params) ret (*name) params; | ||
| 25 | |||
| 26 | struct OpsVTable { | ||
| 27 | BROTENSOR_FOR_EACH_OP(BROTENSOR_VTABLE_FIELD) | ||
| 28 | }; | ||
| 29 | |||
| 30 | #undef BROTENSOR_VTABLE_FIELD | ||
| 31 | |||
| 32 | // ─── AllocVTable ─────────────────────────────────────────────────────────── | ||
| 33 | // | ||
| 34 | // Per-backend memory operations. Used by Tensor's destructor / factories / | ||
| 35 | // to() / clone() / resize() / zero() to manage storage without backend | ||
| 36 | // awareness leaking into the public Tensor API. | ||
| 37 | // | ||
| 38 | // Conventions: | ||
| 39 | // * `alloc(0)` returns nullptr; `free(nullptr)` is a no-op. | ||
| 40 | // * `memset_zero` is byte-wise zero (matches the GpuTensor::zero contract | ||
| 41 | // — fine because FP32, FP16, and INT8 all encode +0 as all-bits-zero). | ||
| 42 | // * `sync` waits for pending work on this backend to drain. CPU's slot is | ||
| 43 | // a no-op. | ||
| 44 | struct AllocVTable { | ||
| 45 | void* (*alloc)(std::size_t bytes, int device_index); | ||
| 46 | void (*free)(void* ptr, int device_index); | ||
| 47 | void (*memcpy_h2d)(void* dst, const void* src, std::size_t n, int device_index); | ||
| 48 | void (*memcpy_d2h)(void* dst, const void* src, std::size_t n, int device_index); | ||
| 49 | void (*memcpy_d2d)(void* dst, const void* src, std::size_t n, int device_index); | ||
| 50 | void (*memcpy_peer)(void* dst, int dst_device_index, const void* src, int src_device_index, std::size_t n); | ||
| 51 | void (*memset_zero)(void* dst, std::size_t n, int device_index); | ||
| 52 | void (*sync)(int device_index); | ||
| 53 | // Optional (backends that can't report leave it null via aggregate | ||
| 54 | // value-init): device-wide free/total memory in bytes, e.g. | ||
| 55 | // cudaMemGetInfo. Serves brotensor::device_mem_info(). | ||
| 56 | bool (*mem_info)(std::size_t* free_bytes, std::size_t* total_bytes, int device_index); | ||
| 57 | // Optional: return the allocator's cached-but-unused blocks to the | ||
| 58 | // driver, keeping at most `keep_bytes` cached (e.g. cudaMemPoolTrimTo | ||
| 59 | // after a device sync). Serves brotensor::device_mem_trim(). | ||
| 60 | bool (*mem_trim)(std::size_t keep_bytes, int device_index); | ||
| 61 | // Optional: the device's human-readable name (e.g. cudaDeviceProp.name, | ||
| 62 | // "NVIDIA GeForce RTX 4090"), written NUL-terminated into `out` (truncated | ||
| 63 | // to `cap`). Serves brotensor::device_name(). Null when the backend can't | ||
| 64 | // report (leaves the field null via aggregate value-init). | ||
| 65 | bool (*device_name)(char* out, std::size_t cap, int device_index); | ||
| 66 | // Optional: true when the work queued on the device's current stream has | ||
| 67 | // finished, without waiting (e.g. cudaStreamQuery). Serves | ||
| 68 | // brotensor::idle(); null makes idle() sync and answer true. | ||
| 69 | bool (*idle)(int device_index); | ||
| 70 | }; | ||
| 71 | |||
| 72 | // ─── Backend registration ────────────────────────────────────────────────── | ||
| 73 | // | ||
| 74 | // Each backend registers itself once. CPU does so from a static-init object | ||
| 75 | // (so CPU tensors work without a prior `init()` call); CUDA / Metal do so | ||
| 76 | // from `brotensor::init()` after a successful driver probe. Registering the | ||
| 77 | // same Device twice replaces the previous tables — callers should not rely | ||
| 78 | // on that (it exists for testability). | ||
| 79 | void register_backend(DeviceType dt, const OpsVTable& ops, const AllocVTable& alloc); | ||
| 80 | void register_backend(Device d, const OpsVTable& ops, const AllocVTable& alloc); | ||
| 81 | |||
| 82 | // ─── The Device::CUDA -> Metal / Vulkan alias ────────────────────────────── | ||
| 83 | // | ||
| 84 | // Much code names the GPU `Device::CUDA`. So that such code runs unchanged on | ||
| 85 | // an Apple, AMD (or any non-NVIDIA) machine, when the CUDA backend is not | ||
| 86 | // registered, a Device::cuda(i) *naming a place to put or run work* means the | ||
| 87 | // GPU that is: Device::Metal when the Metal backend is registered, else | ||
| 88 | // Device::vulkan(i) when the Vulkan backend is. With both Metal and Vulkan | ||
| 89 | // registered, a Vulkan default device (BROTENSOR_DEFAULT_DEVICE=vulkan, | ||
| 90 | // set_default_device, DeviceScope) picks Vulkan; otherwise Metal, the order | ||
| 91 | // the default device itself is picked in. The rule, in full: | ||
| 92 | // | ||
| 93 | // * Everything that acts on a device resolves the alias first: the tensor | ||
| 94 | // factories and view() (the tensor comes back tagged Vulkan), the ops / alloc | ||
| 95 | // tables, dispatch's device activation, sync, set_default_device, | ||
| 96 | // DeviceScope, device_mem_info / device_mem_trim / device_product_name. | ||
| 97 | // * Everything that reports what hardware exists does not: is_available / | ||
| 98 | // is_registered(Device::CUDA) stay false, available_devices() lists | ||
| 99 | // Metal / Vulkan devices only, cuda_device_count() is 0. Code that probes for CUDA | ||
| 100 | // before using it therefore sees the truth and picks its non-CUDA path; | ||
| 101 | // code that just says Device::CUDA gets the GPU that is there. | ||
| 102 | // | ||
| 103 | // With a CUDA backend registered (or neither Metal nor Vulkan) the alias is | ||
| 104 | // the identity. Resolution happens here and nowhere else. | ||
| 105 | Device resolve_device_alias(Device d); | ||
| 106 | |||
| 107 | // Dispatcher lookups (alias-resolved). Throw std::runtime_error if `d` is not | ||
| 108 | // currently registered. | ||
| 109 | const OpsVTable& ops_for(Device d); | ||
| 110 | const AllocVTable& alloc_for(Device d); | ||
| 111 | |||
| 112 | // True iff `d` has been registered (CPU is always true after static init). | ||
| 113 | // Not alias-resolved: see above. | ||
| 114 | bool is_registered(Device d); | ||
| 115 | |||
| 116 | // ─── Operand-consistency helpers ─────────────────────────────────────────── | ||
| 117 | // | ||
| 118 | // Every public op wrapper in `src/ops.cpp` starts with one of these calls. | ||
| 119 | // They resolve the op's device from the first *committed* operand (one whose | ||
| 120 | // `data != nullptr`) and verify every other committed operand agrees, | ||
| 121 | // returning that backend's vtable. On mismatch they throw std::runtime_error | ||
| 122 | // with a human-readable "operand <i> is on <dev> but the op resolved to | ||
| 123 | // <dev>" message. | ||
| 124 | // | ||
| 125 | // An *uncommitted* operand — a freshly default-constructed or resized-to-empty | ||
| 126 | // tensor with `data == nullptr` — has no device affinity yet and is treated as | ||
| 127 | // a wildcard: it is skipped by the consistency check. This is what lets a | ||
| 128 | // caller pass an unsized output tensor to any op without first pinning it to | ||
| 129 | // the right backend. The wrapper then calls `adopt_output` (below) so the | ||
| 130 | // tensor is tagged with the resolved device before the backend impl allocates | ||
| 131 | // it. | ||
| 132 | // | ||
| 133 | // Overloads cover up to 8 operands — enough for the heavier ops (attention, | ||
| 134 | // resblock, conv2d backward). Optional `const Tensor*` operands should be | ||
| 135 | // passed via the `_opt` variants, which skip null pointers. | ||
| 136 | const OpsVTable& dispatch(const Tensor& a); | ||
| 137 | const OpsVTable& dispatch(const Tensor& a, const Tensor& b); | ||
| 138 | const OpsVTable& dispatch(const Tensor& a, const Tensor& b, const Tensor& c); | ||
| 139 | const OpsVTable& dispatch(const Tensor& a, const Tensor& b, const Tensor& c, | ||
| 140 | const Tensor& d); | ||
| 141 | const OpsVTable& dispatch(const Tensor& a, const Tensor& b, const Tensor& c, | ||
| 142 | const Tensor& d, const Tensor& e); | ||
| 143 | const OpsVTable& dispatch(const Tensor& a, const Tensor& b, const Tensor& c, | ||
| 144 | const Tensor& d, const Tensor& e, const Tensor& f); | ||
| 145 | const OpsVTable& dispatch(const Tensor& a, const Tensor& b, const Tensor& c, | ||
| 146 | const Tensor& d, const Tensor& e, const Tensor& f, | ||
| 147 | const Tensor& g); | ||
| 148 | const OpsVTable& dispatch(const Tensor& a, const Tensor& b, const Tensor& c, | ||
| 149 | const Tensor& d, const Tensor& e, const Tensor& f, | ||
| 150 | const Tensor& g, const Tensor& h); | ||
| 151 | |||
| 152 | // Variant for ops whose surface includes both required and optional | ||
| 153 | // operands. Pass any number of required operands followed by a (possibly | ||
| 154 | // empty) list of optional pointers; null pointers are skipped. The dispatch | ||
| 155 | // device is the first required operand's device. | ||
| 156 | const OpsVTable& dispatch_with_opts(const Tensor& a, | ||
| 157 | std::initializer_list<const Tensor*> opts); | ||
| 158 | const OpsVTable& dispatch_with_opts(const Tensor& a, const Tensor& b, | ||
| 159 | std::initializer_list<const Tensor*> opts); | ||
| 160 | |||
| 161 | // Pin an uncommitted output tensor to the op's device. If `t` already owns | ||
| 162 | // storage (data != nullptr) this is a no-op — its device is left untouched so | ||
| 163 | // a genuine mismatch would have been caught by dispatch(). The wrapper calls | ||
| 164 | // this for every non-const Tensor& output, after dispatch and before invoking | ||
| 165 | // the backend impl, so the impl resizes/allocates on the right backend. | ||
| 166 | void adopt_output(Tensor& t, Device d); | ||
| 167 | |||
| 168 | // Set registered CUDA device count. | ||
| 169 | void set_cuda_device_count(int count); | ||
| 170 | |||
| 171 | // Set registered Vulkan device count. | ||
| 172 | void set_vulkan_device_count(int count); | ||
| 173 | |||
| 174 | // Throw a "<op>: not implemented on <device>" std::runtime_error. Called | ||
| 175 | // by the public wrapper when the chosen vtable slot is null. | ||
| 176 | [[noreturn]] void throw_not_implemented(const char* op_name, Device d); | ||
| 177 | |||
| 178 | // ─── Spatial & block-quant dimension validation guards ───────────────────── | ||
| 179 | // | ||
| 180 | // Inline hot-path checks (a few integer comparisons + 64-bit multiply) that | ||
| 181 | // delegate formatting + throwing to cold [[noreturn]] helpers in dispatch.cpp. | ||
| 182 | // The spatial ops index their input as contiguous N*C*H*W floats, and | ||
| 183 | // callers hand the same storage over as (N, C*H*W), (1, N*C*H*W) or a matrix | ||
| 184 | // they are transposing (CLAP's mel image, Whisper's embedding), so a spatial | ||
| 185 | // guard checks the element count, not the row split. | ||
| 186 | [[noreturn]] void throw_spatial_1d_error(const char* op, const char* tensor_name, | ||
| 187 | int N, int C, int L, int rows, int cols); | ||
| 188 | [[noreturn]] void throw_spatial_2d_error(const char* op, const char* tensor_name, | ||
| 189 | int N, int C, int H, int W, int rows, int cols); | ||
| 190 | [[noreturn]] void throw_spatial_3d_error(const char* op, const char* tensor_name, | ||
| 191 | int N, int C, int T, int H, int W, int rows, int cols); | ||
| 192 | [[noreturn]] void throw_quant_cols_error(const char* op, const char* tensor_name, | ||
| 193 | Dtype dt, int rows, int cols); | ||
| 194 | |||
| 195 | 941 | inline void check_spatial_1d(const char* op, const Tensor& t, const char* name, | |
| 196 | int N, int C, int L) { | ||
| 197 |
2/2✓ Branch 0 taken 939 times.
✓ Branch 1 taken 2 times.
|
941 | if (N <= 0 || C <= 0 || L <= 0 || |
| 198 | 939 | static_cast<std::int64_t>(t.rows) * t.cols != static_cast<std::int64_t>(N) * C * L) { | |
| 199 | 2 | throw_spatial_1d_error(op, name, N, C, L, t.rows, t.cols); | |
| 200 | } | ||
| 201 | 939 | } | |
| 202 | |||
| 203 | 1434 | inline void check_spatial_2d(const char* op, const Tensor& t, const char* name, | |
| 204 | int N, int C, int H, int W) { | ||
| 205 |
2/2✓ Branch 0 taken 1425 times.
✓ Branch 1 taken 9 times.
|
1434 | if (N <= 0 || C <= 0 || H <= 0 || W <= 0 || |
| 206 | 1425 | static_cast<std::int64_t>(t.rows) * t.cols != static_cast<std::int64_t>(N) * C * H * W) { | |
| 207 | 9 | throw_spatial_2d_error(op, name, N, C, H, W, t.rows, t.cols); | |
| 208 | } | ||
| 209 | 1425 | } | |
| 210 | |||
| 211 | 57 | inline void check_spatial_3d(const char* op, const Tensor& t, const char* name, | |
| 212 | int N, int C, int T, int H, int W) { | ||
| 213 |
2/2✓ Branch 0 taken 54 times.
✓ Branch 1 taken 3 times.
|
57 | if (N <= 0 || C <= 0 || T <= 0 || H <= 0 || W <= 0 || |
| 214 | 54 | static_cast<std::int64_t>(t.rows) * t.cols != static_cast<std::int64_t>(N) * C * T * H * W) { | |
| 215 | 3 | throw_spatial_3d_error(op, name, N, C, T, H, W, t.rows, t.cols); | |
| 216 | } | ||
| 217 | 54 | } | |
| 218 | |||
| 219 | 224 | inline void check_quant_cols(const char* op, const Tensor& t, const char* name) { | |
| 220 |
2/2✓ Branch 0 taken 94 times.
✓ Branch 1 taken 130 times.
|
224 | if (dtype_is_quant(t.dtype)) { |
| 221 | 130 | const int bs = dtype_block_size(t.dtype); | |
| 222 |
2/2✓ Branch 0 taken 127 times.
✓ Branch 1 taken 3 times.
|
130 | if (t.rows <= 0 || t.cols <= 0 || (t.cols % bs) != 0) { |
| 223 | 3 | throw_quant_cols_error(op, name, t.dtype, t.rows, t.cols); | |
| 224 | } | ||
| 225 | 127 | } | |
| 226 | 221 | } | |
| 227 | |||
| 228 | #define BROTENSOR_CHECK_SPATIAL_1D(op, t, N, C, L) \ | ||
| 229 | ::brotensor::detail::check_spatial_1d((op), (t), #t, (N), (C), (L)) | ||
| 230 | #define BROTENSOR_CHECK_SPATIAL_2D(op, t, N, C, H, W) \ | ||
| 231 | ::brotensor::detail::check_spatial_2d((op), (t), #t, (N), (C), (H), (W)) | ||
| 232 | #define BROTENSOR_CHECK_SPATIAL_3D(op, t, N, C, T, H, W) \ | ||
| 233 | ::brotensor::detail::check_spatial_3d((op), (t), #t, (N), (C), (T), (H), (W)) | ||
| 234 | #define BROTENSOR_CHECK_QUANT_COLS(op, t) \ | ||
| 235 | ::brotensor::detail::check_quant_cols((op), (t), #t) | ||
| 236 | |||
| 237 | } // namespace brotensor::detail | ||
| 238 |