GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 97.5% 118 / 0 / 121
Functions: 100.0% 24 / 0 / 24
Branches: 66.9% 79 / 0 / 118
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
11 #include <array>
12 #include <cstddef>
13 #include <cstring>
14 #include <initializer_list>
15 #include <stdexcept>
16 #include <string>
17
18 namespace brotensor::detail {
19
20 namespace {
21
22 constexpr int kNumDevices = 3; // CPU, CUDA, Metal
23
24 // Upper bound on operands a single dispatch call inspects. The widest op is
25 // resblock_backward (25 operands); 32 leaves headroom. dispatch_with_opts
26 // throws if a future op ever exceeds this rather than silently truncating.
27 constexpr std::size_t kMaxOperands = 32;
28
29 struct Slot {
30 OpsVTable ops{};
31 AllocVTable alloc{};
32 bool registered = false;
33 };
34
35 252279 std::array<Slot, kNumDevices>& slots() {
36 static std::array<Slot, kNumDevices> s{};
37 252279 return s;
38 }
39
40 14 const char* dev_name(Device d) {
41
2/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 14 switch (d) {
42 8 case Device::CPU: return "CPU";
43 6 case Device::CUDA: return "CUDA";
44 ✗ case Device::Metal: return "Metal";
45 }
46 ✗ return "?";
47 14 }
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 209189 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 // Resolve the op's device from the first committed operand; verify every other
74 // committed operand agrees (uncommitted operands are skipped). Returns the
75 // vtable for the resolved device.
76 64609 const OpsVTable& dispatch_v(std::initializer_list<const Tensor*> ts) {
77 64609 Device dev = Device::CPU;
78 64609 bool found = false;
79
2/2
✓ Branch 0 taken 18 times.
✓ Branch 1 taken 64613 times.
 64631 for (const Tensor* t : ts) {
80
3/4
✓ Branch 0 taken 64613 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 22 times.
✓ Branch 3 taken 64591 times.
 64613 if (t && committed(*t)) { dev = t->device; found = true; break; }
81 }
82
4/6
✓ Branch 0 taken 18 times.
✓ Branch 1 taken 64591 times.
✓ Branch 2 taken 18 times.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 18 times.
 64609 if (!found && ts.size() != 0 && *ts.begin() != nullptr) {
83 // All operands empty — degenerate but not an error (e.g. a fresh
84 // output passed before any input). Keep operand 0's tag.
85 18 dev = (*ts.begin())->device;
86 18 }
87 64609 int idx = 0;
88
2/2
✓ Branch 0 taken 139368 times.
✓ Branch 1 taken 64607 times.
 203975 for (const Tensor* t : ts) {
89
5/6
✓ Branch 0 taken 139368 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 85391 times.
✓ Branch 3 taken 53977 times.
✓ Branch 4 taken 2 times.
✓ Branch 5 taken 85389 times.
 139368 if (t && committed(*t) && t->device != dev) {
90 2 throw_device_mismatch(dev, idx, t->device);
91 }
92 139366 ++idx;
93 }
94 64607 return ops_for(dev);
95 }
96
97 // Same resolution rule as dispatch_v, over a plain operand array — used by the
98 // optional-operand dispatch path so it needs no heap allocation. Null entries
99 // (skipped optional operands) are tolerated.
100 754 const OpsVTable& resolve_over(const Tensor* const* all, std::size_t count) {
101 754 Device dev = Device::CPU;
102 754 bool found = false;
103
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 758 times.
 760 for (std::size_t i = 0; i < count; ++i) {
104
4/4
✓ Branch 0 taken 756 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 4 times.
✓ Branch 3 taken 752 times.
 758 if (all[i] && committed(*all[i])) {
105 752 dev = all[i]->device;
106 752 found = true;
107 752 break;
108 }
109 6 }
110
2/2
✓ Branch 0 taken 752 times.
✓ Branch 1 taken 2 times.
 754 if (!found) {
111 // No committed operand — adopt the first non-null operand's tag.
112
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 for (std::size_t i = 0; i < count; ++i) {
113
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (all[i]) { dev = all[i]->device; break; }
114 ✗ }
115 2 }
116
2/2
✓ Branch 0 taken 5162 times.
✓ Branch 1 taken 752 times.
 5914 for (std::size_t i = 0; i < count; ++i) {
117
6/6
✓ Branch 0 taken 4452 times.
✓ Branch 1 taken 710 times.
✓ Branch 2 taken 3763 times.
✓ Branch 3 taken 689 times.
✓ Branch 4 taken 2 times.
✓ Branch 5 taken 3761 times.
 5162 if (all[i] && committed(*all[i]) && all[i]->device != dev) {
118 2 throw_device_mismatch(dev, static_cast<int>(i), all[i]->device);
119 }
120 5160 }
121 752 return ops_for(dev);
122 }
123
124 1 [[noreturn]] void throw_too_many_operands(std::size_t n) {
125
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(
126
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) +
127
1/2
✓ Branch 0 taken 1 time.
✗ Branch 1 not taken.
 1 " exceeds the fixed dispatch buffer (" +
128
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) + ")");
129 1 }
130
131 } // namespace
132
133 315 void register_backend(Device d, const OpsVTable& ops, const AllocVTable& alloc) {
134 315 auto& s = slots()[static_cast<int>(d)];
135 315 std::memcpy(&s.ops, &ops, sizeof(OpsVTable));
136 315 std::memcpy(&s.alloc, &alloc, sizeof(AllocVTable));
137 315 s.registered = true;
138 315 }
139
140 8140 bool is_registered(Device d) {
141 8140 return slots()[static_cast<int>(d)].registered;
142 }
143
144 65450 const OpsVTable& ops_for(Device d) {
145 65450 auto& s = slots()[static_cast<int>(d)];
146
2/2
✓ Branch 0 taken 65449 times.
✓ Branch 1 taken 1 time.
 65450 if (!s.registered) throw_unregistered(d);
147 65449 return s.ops;
148 }
149
150 178374 const AllocVTable& alloc_for(Device d) {
151 178374 auto& s = slots()[static_cast<int>(d)];
152
2/2
✓ Branch 0 taken 178373 times.
✓ Branch 1 taken 1 time.
 178374 if (!s.registered) throw_unregistered(d);
153 178373 return s.alloc;
154 }
155
156 // ─── dispatch overloads ────────────────────────────────────────────────────
157 //
158 // Each forwards to dispatch_v, which resolves the device from the first
159 // committed operand and skips uncommitted (data == nullptr) wildcards.
160
161 1709 const OpsVTable& dispatch(const Tensor& a) {
162 1709 return dispatch_v({&a});
163 }
164
165 53978 const OpsVTable& dispatch(const Tensor& a, const Tensor& b) {
166 53978 return dispatch_v({&a, &b});
167 }
168
169 7151 const OpsVTable& dispatch(const Tensor& a, const Tensor& b, const Tensor& c) {
170 7151 return dispatch_v({&a, &b, &c});
171 }
172
173 1340 const OpsVTable& dispatch(const Tensor& a, const Tensor& b, const Tensor& c,
174 const Tensor& d) {
175 1340 return dispatch_v({&a, &b, &c, &d});
176 }
177
178 89 const OpsVTable& dispatch(const Tensor& a, const Tensor& b, const Tensor& c,
179 const Tensor& d, const Tensor& e) {
180 89 return dispatch_v({&a, &b, &c, &d, &e});
181 }
182
183 107 const OpsVTable& dispatch(const Tensor& a, const Tensor& b, const Tensor& c,
184 const Tensor& d, const Tensor& e, const Tensor& f) {
185 107 return dispatch_v({&a, &b, &c, &d, &e, &f});
186 }
187
188 77 const OpsVTable& dispatch(const Tensor& a, const Tensor& b, const Tensor& c,
189 const Tensor& d, const Tensor& e, const Tensor& f,
190 const Tensor& g) {
191 77 return dispatch_v({&a, &b, &c, &d, &e, &f, &g});
192 }
193
194 158 const OpsVTable& dispatch(const Tensor& a, const Tensor& b, const Tensor& c,
195 const Tensor& d, const Tensor& e, const Tensor& f,
196 const Tensor& g, const Tensor& h) {
197 158 return dispatch_v({&a, &b, &c, &d, &e, &f, &g, &h});
198 }
199
200 56 const OpsVTable& dispatch_with_opts(const Tensor& a,
201 std::initializer_list<const Tensor*> opts) {
202 56 const std::size_t count = 1 + opts.size();
203
2/2
✓ Branch 0 taken 55 times.
✓ Branch 1 taken 1 time.
 56 if (count > kMaxOperands) throw_too_many_operands(count);
204 const Tensor* all[kMaxOperands];
205 55 std::size_t n = 0;
206 55 all[n++] = &a;
207
2/2
✓ Branch 0 taken 117 times.
✓ Branch 1 taken 55 times.
 172 for (const Tensor* p : opts) all[n++] = p;
208 55 return resolve_over(all, n);
209 }
210
211 699 const OpsVTable& dispatch_with_opts(const Tensor& a, const Tensor& b,
212 std::initializer_list<const Tensor*> opts) {
213 699 const std::size_t count = 2 + opts.size();
214
1/2
✓ Branch 0 taken 699 times.
✗ Branch 1 not taken.
 699 if (count > kMaxOperands) throw_too_many_operands(count);
215 const Tensor* all[kMaxOperands];
216 699 std::size_t n = 0;
217 699 all[n++] = &a;
218 699 all[n++] = &b;
219
2/2
✓ Branch 0 taken 3592 times.
✓ Branch 1 taken 699 times.
 4291 for (const Tensor* p : opts) all[n++] = p;
220 699 return resolve_over(all, n);
221 }
222
223 // ─── output adoption ───────────────────────────────────────────────────────
224 //
225 // An uncommitted output tensor (data == nullptr) has no device affinity yet.
226 // The wrapper calls this after dispatch so the tensor is pinned to the op's
227 // device before the backend impl resizes/allocates it.
228 70679 void adopt_output(Tensor& t, Device d) {
229
2/2
✓ Branch 0 taken 12706 times.
✓ Branch 1 taken 57973 times.
 70679 if (t.data == nullptr) t.device = d;
230 70679 }
231
232 4 [[noreturn]] void throw_not_implemented(const char* op_name, Device d) {
233 4 std::string msg = "brotensor: ";
234
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 msg += op_name;
235
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 msg += ": not implemented on ";
236
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);
237
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);
238 4 }
239
240 } // namespace brotensor::detail
241