GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 92.7% 127 / 0 / 137
Functions: 92.9% 26 / 0 / 28
Branches: 66.4% 85 / 0 / 128
src/dispatch.cpp
Line Branch Exec Source
1 // Backend registration + per-operand dispatch: the dispatcher core.
2 //
3 // Storage is two parallel std::array<>s indexed by static_cast<int>(Device).
4 // Registration happens during static init (CPU) and at brotensor::init()
5 // time (CUDA/Metal probe success). Both are single-threaded entry points
6 // before any op call, so no mutex is required.
7
8 #include <brotensor/detail/dispatch.h>
9 #include <brotensor/tensor.h>
10 #include <brotensor/runtime.h>
11
12 #include <array>
13 #include <cstddef>
14 #include <cstring>
15 #include <initializer_list>
16 #include <stdexcept>
17 #include <string>
18
19 #if defined(BROTENSOR_HAS_CUDA)
20 #include <cuda_runtime.h>
21 #endif
22
23 namespace brotensor::detail {
24
25 namespace {
26
27 constexpr int kNumDevices = 3; // CPU, CUDA, Metal
28
29 // Upper bound on operands a single dispatch call inspects. The widest op is
30 // resblock_backward (25 operands); 32 leaves headroom. dispatch_with_opts
31 // throws if a future op ever exceeds this rather than silently truncating.
32 constexpr std::size_t kMaxOperands = 32;
33
34 struct Slot {
35 OpsVTable ops{};
36 AllocVTable alloc{};
37 bool registered = false;
38 };
39
40 261144 std::array<Slot, kNumDevices>& slots() {
41 static std::array<Slot, kNumDevices> s{};
42 261144 return s;
43 }
44
45 14 const char* dev_name(Device d) {
46 14 return device_name(d);
47 }
48
49 2 [[noreturn]] void throw_unregistered(Device d) {
50 2 std::string msg = "brotensor: backend ";
51
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 msg += dev_name(d);
52
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 msg += " not registered";
53
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
 2 throw std::runtime_error(msg);
54 2 }
55
56 // A tensor is "committed" to a device once it owns real storage. A freshly
57 // default-constructed (or resized-to-empty) tensor has data == nullptr and
58 // carries no device affinity yet — it is a wildcard that adopts the op's
59 // device. This is what lets callers pass an unsized output tensor to a GPU op
60 // without first pinning it to the right backend by hand.
61 211440 inline bool committed(const Tensor& t) { return t.data != nullptr; }
62
63 4 [[noreturn]] void throw_device_mismatch(Device resolved, int idx, Device got) {
64 4 std::string msg = "brotensor: dispatch: operand ";
65
2/4
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 msg += std::to_string(idx);
66
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 msg += " is on ";
67
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
 4 msg += dev_name(got);
68
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 msg += " but the op resolved to ";
69
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
 4 msg += dev_name(resolved);
70
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
 4 throw std::runtime_error(msg);
71 4 }
72
73 65831 inline void activate_device_context(Device dev) {
74 #if defined(BROTENSOR_HAS_CUDA)
75 if (dev.is_cuda()) {
76 cudaSetDevice(dev.index);
77 }
78 #else
79 (void)dev;
80 #endif
81 65831 }
82
83 // Resolve the op's device from the first committed operand; verify every other
84 // committed operand agrees (uncommitted operands are skipped). Returns the
85 // vtable for the resolved device.
86 64973 const OpsVTable& dispatch_v(std::initializer_list<const Tensor*> ts) {
87 64973 Device dev = Device::CPU;
88 64973 bool found = false;
89
2/2
✓ Branch 0 taken 18 times.
✓ Branch 1 taken 64977 times.
 64995 for (const Tensor* t : ts) {
90
3/4
✓ Branch 0 taken 64977 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 22 times.
✓ Branch 3 taken 64955 times.
 64977 if (t && committed(*t)) { dev = t->device; found = true; break; }
91 }
92
4/6
✓ Branch 0 taken 18 times.
✓ Branch 1 taken 64955 times.
✓ Branch 2 taken 18 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 18 times.
 64973 if (!found && ts.size() != 0 && *ts.begin() != nullptr) {
93 // All operands empty — degenerate but not an error (e.g. a fresh
94 // output passed before any input). Keep operand 0's tag.
95 18 dev = (*ts.begin())->device;
96 18 }
97 64973 int idx = 0;
98
2/2
✓ Branch 0 taken 140772 times.
✓ Branch 1 taken 64971 times.
 205743 for (const Tensor* t : ts) {
99
5/6
✓ Branch 0 taken 140772 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 86753 times.
✓ Branch 3 taken 54019 times.
✓ Branch 4 taken 2 times.
✓ Branch 5 taken 86751 times.
 140772 if (t && committed(*t) && t->device != dev) {
100 2 throw_device_mismatch(dev, idx, t->device);
101 }
102 140770 ++idx;
103 }
104 64971 activate_device_context(dev);
105 64971 return ops_for(dev);
106 }
107
108 // Same resolution rule as dispatch_v, over a plain operand array — used by the
109 // optional-operand dispatch path so it needs no heap allocation. Null entries
110 // (skipped optional operands) are tolerated.
111 862 const OpsVTable& resolve_over(const Tensor* const* all, std::size_t count) {
112 862 Device dev = Device::CPU;
113 862 bool found = false;
114
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 866 times.
 868 for (std::size_t i = 0; i < count; ++i) {
115
4/4
✓ Branch 0 taken 864 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 4 times.
✓ Branch 3 taken 860 times.
 866 if (all[i] && committed(*all[i])) {
116 860 dev = all[i]->device;
117 860 found = true;
118 860 break;
119 }
120 6 }
121
2/2
✓ Branch 0 taken 860 times.
✓ Branch 1 taken 2 times.
 862 if (!found) {
122 // No committed operand — adopt the first non-null operand's tag.
123
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 for (std::size_t i = 0; i < count; ++i) {
124
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (all[i]) { dev = all[i]->device; break; }
125 ✗ }
126 2 }
127
2/2
✓ Branch 0 taken 5588 times.
✓ Branch 1 taken 860 times.
 6448 for (std::size_t i = 0; i < count; ++i) {
128
6/6
✓ Branch 0 taken 4827 times.
✓ Branch 1 taken 761 times.
✓ Branch 2 taken 4030 times.
✓ Branch 3 taken 797 times.
✓ Branch 4 taken 2 times.
✓ Branch 5 taken 4028 times.
 5588 if (all[i] && committed(*all[i]) && all[i]->device != dev) {
129 2 throw_device_mismatch(dev, static_cast<int>(i), all[i]->device);
130 }
131 5586 }
132 860 activate_device_context(dev);
133 860 return ops_for(dev);
134 }
135
136 1 [[noreturn]] void throw_too_many_operands(std::size_t n) {
137
3/6
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 1 time.
✓ Branch 4 taken 1 time.
✗ Branch 5 not taken.
 2 throw std::runtime_error(
138
3/6
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
✓ Branch 4 taken 1 time.
✗ Branch 5 not taken.
 1 "brotensor: dispatch: operand count " + std::to_string(n) +
139
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 " exceeds the fixed dispatch buffer (" +
140
2/4
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
 2 std::to_string(kMaxOperands) + ")");
141 1 }
142
143 } // namespace
144
145 static int g_cuda_device_count = 0;
146
147 ✗ void set_cuda_device_count(int count) {
148 ✗ g_cuda_device_count = count;
149 ✗ }
150
151 319 void register_backend(DeviceType dt, const OpsVTable& ops, const AllocVTable& alloc) {
152 319 auto& s = slots()[static_cast<int>(dt)];
153 319 std::memcpy(&s.ops, &ops, sizeof(OpsVTable));
154 319 std::memcpy(&s.alloc, &alloc, sizeof(AllocVTable));
155 319 s.registered = true;
156 319 }
157
158 319 void register_backend(Device d, const OpsVTable& ops, const AllocVTable& alloc) {
159 319 register_backend(d.type, ops, alloc);
160 319 }
161
162 9540 bool is_registered(Device d) {
163
2/2
✓ Branch 0 taken 2928 times.
✓ Branch 1 taken 6612 times.
 9540 if (d.is_cpu()) return slots()[static_cast<int>(DeviceType::CPU)].registered;
164
2/2
✓ Branch 0 taken 3319 times.
✓ Branch 1 taken 3293 times.
 6612 if (d.is_cuda()) {
165
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 3319 times.
 3319 if (!slots()[static_cast<int>(DeviceType::CUDA)].registered) return false;
166 ✗ int count = ::brotensor::cuda_device_count();
167 ✗ if (count <= 0) count = 1;
168 ✗ return d.index >= 0 && d.index < count;
169 }
170
1/2
✓ Branch 0 taken 3293 times.
✗ Branch 1 not taken.
 3293 if (d.is_metal()) {
171
2/2
✓ Branch 0 taken 15 times.
✓ Branch 1 taken 3278 times.
 3293 return slots()[static_cast<int>(DeviceType::Metal)].registered && d.index == 0;
172 }
173 ✗ return false;
174 9540 }
175
176 65922 const OpsVTable& ops_for(Device d) {
177 65922 auto& s = slots()[static_cast<int>(d.type)];
178
2/2
✓ Branch 0 taken 65921 times.
✓ Branch 1 taken 1 time.
 65922 if (!s.registered) throw_unregistered(d);
179 65921 return s.ops;
180 }
181
182 185363 const AllocVTable& alloc_for(Device d) {
183 185363 auto& s = slots()[static_cast<int>(d.type)];
184
2/2
✓ Branch 0 taken 185362 times.
✓ Branch 1 taken 1 time.
 185363 if (!s.registered) throw_unregistered(d);
185 185362 return s.alloc;
186 }
187
188 // ─── dispatch overloads ────────────────────────────────────────────────────
189 //
190 // Each forwards to dispatch_v, which resolves the device from the first
191 // committed operand and skips uncommitted (data == nullptr) wildcards.
192
193 1720 const OpsVTable& dispatch(const Tensor& a) {
194 1720 return dispatch_v({&a});
195 }
196
197 53982 const OpsVTable& dispatch(const Tensor& a, const Tensor& b) {
198 53982 return dispatch_v({&a, &b});
199 }
200
201 7162 const OpsVTable& dispatch(const Tensor& a, const Tensor& b, const Tensor& c) {
202 7162 return dispatch_v({&a, &b, &c});
203 }
204
205 1678 const OpsVTable& dispatch(const Tensor& a, const Tensor& b, const Tensor& c,
206 const Tensor& d) {
207 1678 return dispatch_v({&a, &b, &c, &d});
208 }
209
210 89 const OpsVTable& dispatch(const Tensor& a, const Tensor& b, const Tensor& c,
211 const Tensor& d, const Tensor& e) {
212 89 return dispatch_v({&a, &b, &c, &d, &e});
213 }
214
215 107 const OpsVTable& dispatch(const Tensor& a, const Tensor& b, const Tensor& c,
216 const Tensor& d, const Tensor& e, const Tensor& f) {
217 107 return dispatch_v({&a, &b, &c, &d, &e, &f});
218 }
219
220 77 const OpsVTable& dispatch(const Tensor& a, const Tensor& b, const Tensor& c,
221 const Tensor& d, const Tensor& e, const Tensor& f,
222 const Tensor& g) {
223 77 return dispatch_v({&a, &b, &c, &d, &e, &f, &g});
224 }
225
226 158 const OpsVTable& dispatch(const Tensor& a, const Tensor& b, const Tensor& c,
227 const Tensor& d, const Tensor& e, const Tensor& f,
228 const Tensor& g, const Tensor& h) {
229 158 return dispatch_v({&a, &b, &c, &d, &e, &f, &g, &h});
230 }
231
232 158 const OpsVTable& dispatch_with_opts(const Tensor& a,
233 std::initializer_list<const Tensor*> opts) {
234 158 const std::size_t count = 1 + opts.size();
235
2/2
✓ Branch 0 taken 157 times.
✓ Branch 1 taken 1 time.
 158 if (count > kMaxOperands) throw_too_many_operands(count);
236 const Tensor* all[kMaxOperands];
237 157 std::size_t n = 0;
238 157 all[n++] = &a;
239
2/2
✓ Branch 0 taken 423 times.
✓ Branch 1 taken 157 times.
 580 for (const Tensor* p : opts) all[n++] = p;
240 157 return resolve_over(all, n);
241 }
242
243 705 const OpsVTable& dispatch_with_opts(const Tensor& a, const Tensor& b,
244 std::initializer_list<const Tensor*> opts) {
245 705 const std::size_t count = 2 + opts.size();
246
1/2
✓ Branch 0 taken 705 times.
✗ Branch 1 not taken.
 705 if (count > kMaxOperands) throw_too_many_operands(count);
247 const Tensor* all[kMaxOperands];
248 705 std::size_t n = 0;
249 705 all[n++] = &a;
250 705 all[n++] = &b;
251
2/2
✓ Branch 0 taken 3598 times.
✓ Branch 1 taken 705 times.
 4303 for (const Tensor* p : opts) all[n++] = p;
252 705 return resolve_over(all, n);
253 }
254
255 // ─── output adoption ───────────────────────────────────────────────────────
256 //
257 // An uncommitted output tensor (data == nullptr) has no device affinity yet.
258 // The wrapper calls this after dispatch so the tensor is pinned to the op's
259 // device before the backend impl resizes/allocates it.
260 71151 void adopt_output(Tensor& t, Device d) {
261
2/2
✓ Branch 0 taken 13028 times.
✓ Branch 1 taken 58123 times.
 71151 if (t.data == nullptr) t.device = d;
262 71151 }
263
264 4 [[noreturn]] void throw_not_implemented(const char* op_name, Device d) {
265 4 std::string msg = "brotensor: ";
266
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 msg += op_name;
267
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 msg += ": not implemented on ";
268
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
 4 msg += dev_name(d);
269
2/4
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
 4 throw std::runtime_error(msg);
270 4 }
271
272 } // namespace brotensor::detail
273
274 namespace brotensor {
275 ✗ int cuda_device_count() {
276 ✗ return detail::g_cuda_device_count;
277 }
278 }
279