GCC Code Coverage Report

Directory: brotensor/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 52.1% 88 / 0 / 169
Functions: 77.3% 17 / 0 / 22
Branches: 22.3% 53 / 0 / 238
src/init.cpp
Line Branch Exec Source
1 // brotensor runtime: init(), default-device policy, DeviceScope, sync.
2 //
3 // The CPU backend self-registers from a static-init object in
4 // src/cpu/register.cpp. init() probes CUDA / Metal / Vulkan if the
5 // corresponding backend was compiled in. When a backend isn't built,
6 // BROTENSOR_HAS_CUDA / BROTENSOR_HAS_METAL are not defined so the probe
7 // branches compile out.
8
9 #include <brotensor/runtime.h>
10 #include <brotensor/tensor.h>
11 #include <brotensor/detail/dispatch.h>
12 #include <brotensor/detail/cpu/thread_pool.h>
13
14 #include <atomic>
15 #include <initializer_list>
16 #include <cstdlib>
17 #include <cstring>
18 #include <mutex>
19 #include <optional>
20 #include <stdexcept>
21 #include <string>
22 #include <vector>
23
24 #if defined(BROTENSOR_HAS_VULKAN)
25 // Defined in src/vulkan/register.cpp.
26 extern "C" void brotensor_probe_and_register_vulkan();
27 #endif
28
29 #if defined(BROTENSOR_HAS_CUDA)
30 // Defined in src/cuda/init.cu.
31 extern "C" void brotensor_probe_and_register_cuda();
32 #endif
33
34 #if defined(BROTENSOR_HAS_METAL)
35 // Defined in src/metal/init.mm.
36 extern "C" void brotensor_probe_and_register_metal();
37 #endif
38
39 namespace brotensor {
40
41 namespace {
42
43 167 std::mutex& init_mutex() {
44
2/4
✓ Branch 0 taken 167 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 167 times.
 167 static std::mutex m;
45 167 return m;
46 }
47
48 503 std::atomic<bool>& init_done_flag() {
49 static std::atomic<bool> f{false};
50 503 return f;
51 }
52
53 // Global default device. Initialised lazily on first read.
54 2830 std::atomic<Device>& global_default() {
55
3/4
✓ Branch 0 taken 167 times.
✓ Branch 1 taken 2663 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 167 times.
 2830 static std::atomic<Device> d{Device::CPU};
56 2830 return d;
57 }
58
59 3012 std::atomic<bool>& global_default_set_flag() {
60 static std::atomic<bool> f{false};
61 3012 return f;
62 }
63
64 // Thread-local DeviceScope override stack — we only need the *current* value
65 // since DeviceScope ctor saves the previous on the local stack frame.
66 thread_local std::optional<Device> tls_scope_override;
67
68 // CUDA first, then Metal, then Vulkan (the AMD GPU path, docs/vulkan.md).
69 182 Device pick_default_from_available() {
70
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 182 times.
 182 if (detail::is_registered(Device::CUDA)) return Device::CUDA;
71
2/2
✓ Branch 0 taken 167 times.
✓ Branch 1 taken 15 times.
 182 if (detail::is_registered(Device::Metal)) return Device::Metal;
72
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 15 times.
 15 if (detail::is_registered(Device::VULKAN)) return Device::VULKAN;
73 15 return Device::CPU;
74 182 }
75
76 167 std::optional<Device> parse_env_device() {
77 167 const char* env = std::getenv("BROTENSOR_DEFAULT_DEVICE");
78
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 167 times.
 167 if (!env) return std::nullopt;
79 ✗ std::string s(env);
80 ✗ for (char& c : s) {
81 ✗ if (c >= 'A' && c <= 'Z') c = static_cast<char>(c + 32);
82 }
83 ✗ if (s == "cpu") return Device::cpu();
84 ✗ if (s == "cuda" || s == "cuda:0") return Device::cuda(0);
85 ✗ if (s.rfind("cuda:", 0) == 0) {
86 ✗ int idx = std::atoi(s.c_str() + 5);
87 ✗ return Device::cuda(idx);
88 }
89 ✗ for (const char* pfx : {"vulkan", "vk"}) {
90 ✗ const std::string p(pfx);
91 ✗ if (s == p) return Device::vulkan(0);
92 ✗ if (s.rfind(p + ":", 0) == 0) return Device::vulkan(std::atoi(s.c_str() + p.size() + 1));
93 ✗ }
94 ✗ if (s == "metal" || s == "metal:0") return Device::metal(0);
95 ✗ if (s.rfind("metal:", 0) == 0) {
96 ✗ int idx = std::atoi(s.c_str() + 6);
97 ✗ return Device::metal(idx);
98 }
99 ✗ return std::nullopt;
100 167 }
101
102 } // namespace
103
104 169 void init() {
105
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 167 times.
 169 if (init_done_flag().load(std::memory_order_acquire)) return;
106 167 std::lock_guard<std::mutex> lock(init_mutex());
107
2/4
✓ Branch 0 taken 167 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 167 times.
 167 if (init_done_flag().load(std::memory_order_relaxed)) return;
108
109 #if defined(BROTENSOR_HAS_CUDA)
110 try { brotensor_probe_and_register_cuda(); } catch (...) { /* no CUDA */ }
111 #endif
112 #if defined(BROTENSOR_HAS_METAL)
113
1/4
✓ Branch 0 taken 167 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 167 try { brotensor_probe_and_register_metal(); } catch (...) { /* no Metal */ }
114 #endif
115 #if defined(BROTENSOR_HAS_VULKAN)
116 try { brotensor_probe_and_register_vulkan(); } catch (...) { /* no Vulkan */ }
117 #endif
118
119 // Determine default device once.
120
2/4
✓ Branch 0 taken 167 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 167 times.
✗ Branch 3 not taken.
 167 if (!global_default_set_flag().load(std::memory_order_relaxed)) {
121
2/4
✓ Branch 0 taken 167 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 167 times.
 167 if (auto envd = parse_env_device()) {
122 ✗ const Device want = detail::resolve_device_alias(*envd);
123 ✗ if (detail::is_registered(want)) {
124 ✗ global_default().store(want, std::memory_order_relaxed);
125 ✗ } else {
126 ✗ global_default().store(pick_default_from_available(),
127 std::memory_order_relaxed);
128 }
129 ✗ } else {
130
2/4
✓ Branch 0 taken 167 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 167 times.
✗ Branch 3 not taken.
 167 global_default().store(pick_default_from_available(),
131 std::memory_order_relaxed);
132 }
133
1/2
✓ Branch 0 taken 167 times.
✗ Branch 1 not taken.
 167 global_default_set_flag().store(true, std::memory_order_relaxed);
134 167 }
135
136
1/2
✓ Branch 0 taken 167 times.
✗ Branch 1 not taken.
 167 init_done_flag().store(true, std::memory_order_release);
137 169 }
138
139 ✗ void shutdown() {
140 ✗ if (auto* pool = detail::cpu::ThreadPool::existing()) pool->shutdown();
141 ✗ }
142
143 2725 Device default_device() {
144
2/2
✓ Branch 0 taken 50 times.
✓ Branch 1 taken 2675 times.
 2725 if (tls_scope_override.has_value()) return *tls_scope_override;
145 // Lazy default: if init hasn't run, fall back to whatever's registered.
146
2/2
✓ Branch 0 taken 2660 times.
✓ Branch 1 taken 15 times.
 2675 if (!global_default_set_flag().load(std::memory_order_acquire)) {
147 15 return pick_default_from_available();
148 }
149 2660 return global_default().load(std::memory_order_acquire);
150 2725 }
151
152 10 Dtype compute_dtype() {
153 10 return compute_dtype(default_device());
154 }
155
156 10 Dtype compute_dtype(Device d) {
157 10 return d.is_cpu() ? Dtype::FP32 : Dtype::FP16;
158 }
159
160 3 void set_default_device(Device d) {
161 3 d = detail::resolve_device_alias(d);
162
1/2
✓ Branch 0 taken 3 times.
✗ Branch 1 not taken.
 3 if (!detail::is_registered(d)) {
163 ✗ std::string m = "brotensor: set_default_device: backend ";
164 ✗ m += device_name(d);
165 ✗ m += " is not available";
166 ✗ throw std::runtime_error(m);
167 ✗ }
168 3 global_default().store(d, std::memory_order_release);
169 3 global_default_set_flag().store(true, std::memory_order_release);
170 3 }
171
172 4608 std::vector<Device> available_devices() {
173 4608 std::vector<Device> out;
174
2/4
✓ Branch 0 taken 4608 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4608 times.
✗ Branch 3 not taken.
 4608 if (detail::is_registered(Device::CPU)) {
175
1/2
✓ Branch 0 taken 4608 times.
✗ Branch 1 not taken.
 4608 out.push_back(Device::CPU);
176 4608 }
177
2/4
✓ Branch 0 taken 4608 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 4608 times.
 4608 if (detail::is_registered(Device::CUDA)) {
178 ✗ int count = cuda_device_count();
179 ✗ if (count <= 0) count = 1;
180 ✗ for (int i = 0; i < count; ++i) {
181 ✗ out.push_back(Device::cuda(i));
182 ✗ }
183 ✗ }
184
2/4
✓ Branch 0 taken 4608 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4608 times.
✗ Branch 3 not taken.
 4608 if (detail::is_registered(Device::Metal)) {
185
1/2
✓ Branch 0 taken 4608 times.
✗ Branch 1 not taken.
 4608 out.push_back(Device::Metal);
186 4608 }
187
2/4
✓ Branch 0 taken 4608 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 4608 times.
 4608 for (int i = 0; i < vulkan_device_count(); ++i) {
188 ✗ if (detail::is_registered(Device::vulkan(i))) out.push_back(Device::vulkan(i));
189 ✗ }
190 4608 return out;
191 4608 }
192
193 // Truthful about the hardware: not alias-resolved (detail/dispatch.h).
194 284 bool is_available(Device d) {
195 284 return detail::is_registered(d);
196 }
197
198 // ─── DeviceScope ───────────────────────────────────────────────────────────
199
200 // Track whether tls_scope_override held a value before this scope was
201 // pushed. The Device prev_ field on the scope object stores the prior value
202 // when one existed; the parallel bool stack tracks "had a value".
203 namespace {
204 thread_local std::vector<bool> tls_scope_had_prev;
205 } // namespace
206
207 16 DeviceScope::DeviceScope(Device d) {
208 8 d = detail::resolve_device_alias(d);
209
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!detail::is_registered(d)) {
210 ✗ std::string m = "brotensor: DeviceScope: backend ";
211 ✗ m += device_name(d);
212 ✗ m += " is not available";
213 ✗ throw std::runtime_error(m);
214 ✗ }
215
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 7 times.
 8 if (tls_scope_override.has_value()) {
216 1 prev_ = *tls_scope_override;
217 1 tls_scope_had_prev.push_back(true);
218 1 } else {
219 7 prev_ = Device::CPU; // unused when had_prev=false
220 7 tls_scope_had_prev.push_back(false);
221 }
222 8 tls_scope_override = d;
223 16 }
224
225 16 DeviceScope::~DeviceScope() {
226
2/4
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 bool had = tls_scope_had_prev.empty() ? false : tls_scope_had_prev.back();
227
2/4
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 if (!tls_scope_had_prev.empty()) tls_scope_had_prev.pop_back();
228
3/4
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 7 times.
✓ Branch 2 taken 1 time.
✗ Branch 3 not taken.
 8 if (had) tls_scope_override = prev_;
229 7 else tls_scope_override.reset();
230 16 }
231
232 // ─── sync ──────────────────────────────────────────────────────────────────
233
234 84 void sync(Device d) {
235 84 d = detail::resolve_device_alias(d);
236
1/2
✓ Branch 0 taken 84 times.
✗ Branch 1 not taken.
 84 if (!detail::is_registered(d)) {
237 ✗ std::string m = "brotensor: sync: backend ";
238 ✗ m += device_name(d);
239 ✗ m += " is not available";
240 ✗ throw std::runtime_error(m);
241 ✗ }
242 84 detail::alloc_for(d).sync(d.index);
243 84 }
244
245 4606 void sync_all() {
246
2/2
✓ Branch 0 taken 4606 times.
✓ Branch 1 taken 9212 times.
 13818 for (Device d : available_devices()) {
247
2/4
✓ Branch 0 taken 9212 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 9212 times.
✗ Branch 3 not taken.
 9212 detail::alloc_for(d).sync(d.index);
248 }
249 4606 }
250
251 ✗ bool idle(Device d) {
252 ✗ d = detail::resolve_device_alias(d);
253 ✗ if (!detail::is_registered(d)) return true;
254 ✗ const auto& t = detail::alloc_for(d);
255 ✗ if (t.idle) return t.idle(d.index);
256 ✗ t.sync(d.index);
257 ✗ return true;
258 ✗ }
259
260 // ─── memory ────────────────────────────────────────────────────────────────
261
262 ✗ bool device_mem_info(Device d, std::size_t& free_bytes,
263 std::size_t& total_bytes) {
264 ✗ d = detail::resolve_device_alias(d);
265 ✗ if (!detail::is_registered(d)) return false;
266 ✗ const auto fn = detail::alloc_for(d).mem_info;
267 ✗ if (fn == nullptr) return false;
268 ✗ return fn(&free_bytes, &total_bytes, d.index);
269 ✗ }
270
271 ✗ bool device_mem_trim(Device d, std::size_t keep_bytes) {
272 ✗ d = detail::resolve_device_alias(d);
273 ✗ if (!detail::is_registered(d)) return false;
274 ✗ const auto fn = detail::alloc_for(d).mem_trim;
275 ✗ if (fn == nullptr) return false;
276 ✗ return fn(keep_bytes, d.index);
277 ✗ }
278
279 ✗ std::string device_product_name(Device d) {
280 ✗ d = detail::resolve_device_alias(d);
281 ✗ if (!detail::is_registered(d)) return {};
282 ✗ const auto fn = detail::alloc_for(d).device_name;
283 ✗ if (fn == nullptr) return {};
284 ✗ char buf[256] = {0};
285 ✗ if (!fn(buf, sizeof(buf), d.index)) return {};
286 ✗ return std::string(buf);
287 ✗ }
288
289 } // namespace brotensor
290