GCC Code Coverage Report

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