GCC Code Coverage Report

Directory: ./
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 57.6% 76 / 0 / 132
Functions: 76.2% 16 / 0 / 21
Branches: 29.4% 47 / 0 / 160
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 152 std::mutex& init_mutex() {
38
2/4
✓ Branch 0 taken 152 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 152 times.
 152 static std::mutex m;
39 152 return m;
40 }
41
42 458 std::atomic<bool>& init_done_flag() {
43 static std::atomic<bool> f{false};
44 458 return f;
45 }
46
47 // Global default device. Initialised lazily on first read.
48 162 std::atomic<Device>& global_default() {
49 static std::atomic<Device> d{Device::CPU};
50 162 return d;
51 }
52
53 329 std::atomic<bool>& global_default_set_flag() {
54 static std::atomic<bool> f{false};
55 329 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 167 Device pick_default_from_available() {
63
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 167 times.
 167 if (detail::is_registered(Device::CUDA)) return Device::CUDA;
64
2/2
✓ Branch 0 taken 152 times.
✓ Branch 1 taken 15 times.
 167 if (detail::is_registered(Device::Metal)) return Device::Metal;
65 15 return Device::CPU;
66 167 }
67
68 152 std::optional<Device> parse_env_device() {
69 152 const char* env = std::getenv("BROTENSOR_DEFAULT_DEVICE");
70
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 152 times.
 152 if (!env) return std::nullopt;
71 // case-insensitive compare against a few names
72 ✗ auto eq = [](const char* a, const char* b) {
73 ✗ while (*a && *b) {
74 ✗ char ca = *a, cb = *b;
75 ✗ if (ca >= 'A' && ca <= 'Z') ca = static_cast<char>(ca + 32);
76 ✗ if (cb >= 'A' && cb <= 'Z') cb = static_cast<char>(cb + 32);
77 ✗ if (ca != cb) return false;
78 ✗ ++a; ++b;
79 }
80 ✗ return *a == 0 && *b == 0;
81 ✗ };
82 ✗ if (eq(env, "cpu")) return Device::CPU;
83 ✗ if (eq(env, "cuda")) return Device::CUDA;
84 ✗ if (eq(env, "metal")) return Device::Metal;
85 ✗ return std::nullopt;
86 152 }
87
88 } // namespace
89
90 154 void init() {
91
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 152 times.
 154 if (init_done_flag().load(std::memory_order_acquire)) return;
92 152 std::lock_guard<std::mutex> lock(init_mutex());
93
2/4
✓ Branch 0 taken 152 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 152 times.
 152 if (init_done_flag().load(std::memory_order_relaxed)) return;
94
95 #if defined(BROTENSOR_HAS_CUDA)
96 try { brotensor_probe_and_register_cuda(); } catch (...) { /* no CUDA */ }
97 #endif
98 #if defined(BROTENSOR_HAS_METAL)
99
1/4
✓ Branch 0 taken 152 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 152 try { brotensor_probe_and_register_metal(); } catch (...) { /* no Metal */ }
100 #endif
101
102 // Determine default device once.
103
2/4
✓ Branch 0 taken 152 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 152 times.
✗ Branch 3 not taken.
 152 if (!global_default_set_flag().load(std::memory_order_relaxed)) {
104
2/4
✓ Branch 0 taken 152 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 152 times.
 152 if (auto envd = parse_env_device()) {
105 ✗ if (detail::is_registered(*envd)) {
106 ✗ global_default().store(*envd, std::memory_order_relaxed);
107 ✗ } else {
108 ✗ global_default().store(pick_default_from_available(),
109 std::memory_order_relaxed);
110 }
111 ✗ } else {
112
2/4
✓ Branch 0 taken 152 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 152 times.
✗ Branch 3 not taken.
 152 global_default().store(pick_default_from_available(),
113 std::memory_order_relaxed);
114 }
115
1/2
✓ Branch 0 taken 152 times.
✗ Branch 1 not taken.
 152 global_default_set_flag().store(true, std::memory_order_relaxed);
116 152 }
117
118
1/2
✓ Branch 0 taken 152 times.
✗ Branch 1 not taken.
 152 init_done_flag().store(true, std::memory_order_release);
119 154 }
120
121 ✗ void shutdown() {
122 ✗ detail::cpu::ThreadPool::instance().shutdown();
123 ✗ }
124
125 73 Device default_device() {
126
2/2
✓ Branch 0 taken 50 times.
✓ Branch 1 taken 23 times.
 73 if (tls_scope_override.has_value()) return *tls_scope_override;
127 // Lazy default: if init hasn't run, fall back to whatever's registered.
128
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 15 times.
 23 if (!global_default_set_flag().load(std::memory_order_acquire)) {
129 15 return pick_default_from_available();
130 }
131 8 return global_default().load(std::memory_order_acquire);
132 73 }
133
134 8 Dtype compute_dtype() {
135 8 return default_device() == Device::CPU ? Dtype::FP32 : Dtype::FP16;
136 }
137
138 2 void set_default_device(Device d) {
139
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 if (!detail::is_registered(d)) {
140 ✗ std::string m = "brotensor: set_default_device: backend ";
141 ✗ m += device_name(d);
142 ✗ m += " is not available";
143 ✗ throw std::runtime_error(m);
144 ✗ }
145 2 global_default().store(d, std::memory_order_release);
146 2 global_default_set_flag().store(true, std::memory_order_release);
147 2 }
148
149 2452 std::vector<Device> available_devices() {
150 2452 std::vector<Device> out;
151
3/6
✓ Branch 0 taken 2452 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2452 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2452 times.
✗ Branch 5 not taken.
 2452 if (detail::is_registered(Device::CPU)) out.push_back(Device::CPU);
152
2/6
✓ Branch 0 taken 2452 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2452 times.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 2452 if (detail::is_registered(Device::CUDA)) out.push_back(Device::CUDA);
153
3/6
✓ Branch 0 taken 2452 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2452 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2452 times.
✗ Branch 5 not taken.
 2452 if (detail::is_registered(Device::Metal)) out.push_back(Device::Metal);
154 2452 return out;
155 2452 }
156
157 426 bool is_available(Device d) {
158 426 return detail::is_registered(d);
159 }
160
161 // ─── DeviceScope ───────────────────────────────────────────────────────────
162
163 // Track whether tls_scope_override held a value before this scope was
164 // pushed. The Device prev_ field on the scope object stores the prior value
165 // when one existed; the parallel bool stack tracks "had a value".
166 namespace {
167 thread_local std::vector<bool> tls_scope_had_prev;
168 } // namespace
169
170 16 DeviceScope::DeviceScope(Device d) {
171
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 if (!detail::is_registered(d)) {
172 ✗ std::string m = "brotensor: DeviceScope: backend ";
173 ✗ m += device_name(d);
174 ✗ m += " is not available";
175 ✗ throw std::runtime_error(m);
176 ✗ }
177
2/2
✓ Branch 0 taken 1 time.
✓ Branch 1 taken 7 times.
 8 if (tls_scope_override.has_value()) {
178 1 prev_ = *tls_scope_override;
179 1 tls_scope_had_prev.push_back(true);
180 1 } else {
181 7 prev_ = Device::CPU; // unused when had_prev=false
182 7 tls_scope_had_prev.push_back(false);
183 }
184 8 tls_scope_override = d;
185 16 }
186
187 16 DeviceScope::~DeviceScope() {
188
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();
189
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();
190
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_;
191 7 else tls_scope_override.reset();
192 16 }
193
194 // ─── sync ──────────────────────────────────────────────────────────────────
195
196 11 void sync(Device d) {
197
1/2
✓ Branch 0 taken 11 times.
✗ Branch 1 not taken.
 11 if (!detail::is_registered(d)) {
198 ✗ std::string m = "brotensor: sync: backend ";
199 ✗ m += device_name(d);
200 ✗ m += " is not available";
201 ✗ throw std::runtime_error(m);
202 ✗ }
203 11 detail::alloc_for(d).sync();
204 11 }
205
206 2450 void sync_all() {
207
2/2
✓ Branch 0 taken 2450 times.
✓ Branch 1 taken 4900 times.
 7350 for (Device d : available_devices()) {
208
2/4
✓ Branch 0 taken 4900 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4900 times.
✗ Branch 3 not taken.
 4900 detail::alloc_for(d).sync();
209 }
210 2450 }
211
212 // ─── memory ────────────────────────────────────────────────────────────────
213
214 ✗ bool device_mem_info(Device d, std::size_t& free_bytes,
215 std::size_t& total_bytes) {
216 ✗ if (!detail::is_registered(d)) return false;
217 ✗ const auto fn = detail::alloc_for(d).mem_info;
218 ✗ if (fn == nullptr) return false;
219 ✗ return fn(&free_bytes, &total_bytes);
220 ✗ }
221
222 ✗ bool device_mem_trim(Device d, std::size_t keep_bytes) {
223 ✗ if (!detail::is_registered(d)) return false;
224 ✗ const auto fn = detail::alloc_for(d).mem_trim;
225 ✗ if (fn == nullptr) return false;
226 ✗ return fn(keep_bytes);
227 ✗ }
228
229 ✗ std::string device_product_name(Device d) {
230 ✗ if (!detail::is_registered(d)) return {};
231 ✗ const auto fn = detail::alloc_for(d).device_name;
232 ✗ if (fn == nullptr) return {};
233 ✗ char buf[256] = {0};
234 ✗ if (!fn(buf, sizeof(buf))) return {};
235 ✗ return std::string(buf);
236 ✗ }
237
238 } // namespace brotensor
239