GCC Code Coverage Report

Directory: brotensor/
Coverage: low: ≥ 0% medium: ≥ 75.0% high: ≥ 90.0%
Coverage Exec / Excl / Total
Lines: 81.6% 146 / 0 / 179
Functions: 93.8% 15 / 0 / 16
Branches: 42.8% 125 / 0 / 292
src/api/native_tensor_safetensors.cpp
Line Branch Exec Source
1 // bro.tensor.openSafetensors / saveSafetensors and the SafetensorsFile
2 // handle, ported from the QuickJS binding (tensor_bindings_safetensors.cpp):
3 //
4 // openSafetensors(path) -> handle (mmap; header only, no payload read)
5 // handle.count / names() / header() / get(name, rows?, cols?, dtype?) / close()
6 // saveSafetensors(path, { name: GpuTensor, ... })
7 //
8 // The handle owns the mmap'd brotensor::safetensors::File; `file` is null
9 // after close(), and the class destructor unmaps on GC. names()/header() are
10 // assembled in js/tensor.js from the per-index natives here (nameAt /
11 // dtypeAt / shapeAt / nbytesAt): a native returns scalars, typed arrays or
12 // handles, never a JS object. Paths go through api.h's path resolver.
13
14 #include "native_tensor_decl.h"
15 #include "api_internal.h"
16
17 #include <brotensor/safetensors.h>
18
19 #include <cstring>
20 #include <memory>
21 #include <string>
22 #include <vector>
23
24 namespace stns = brotensor::safetensors;
25 namespace ev = bronze::embed;
26 using namespace brotensor::api;
27 using Value = bronze::Value;
28
29 namespace {
30
31 struct SafetensorsFileHandle {
32 std::unique_ptr<stns::File> file;
33 };
34
35 238 SafetensorsFileHandle* asFile(void* self) {
36 238 return static_cast<SafetensorsFileHandle*>(self);
37 }
38
39 // The view at index i of an open file, or nullptr (with the error set).
40 140 const stns::TensorView* viewAt(void* self, int32_t i, const char* label) {
41 140 auto* h = asFile(self);
42
2/4
✓ Branch 0 taken 140 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 140 times.
 140 if (!h || !h->file) {
43 ✗ setError(std::string(label) + "() on a closed safetensors file");
44 ✗ return nullptr;
45 }
46 140 const auto& tensors = h->file->tensors();
47
2/4
✓ Branch 0 taken 140 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 140 times.
 140 if (i < 0 || static_cast<size_t>(i) >= tensors.size()) {
48 ✗ setError(std::string(label) + ": index out of range");
49 ✗ return nullptr;
50 }
51 140 return &tensors[static_cast<size_t>(i)];
52 140 }
53
54 thread_local std::vector<double> tl_shape;
55
56 } // namespace
57
58 extern "C" {
59
60 ✗ void* bro_tensor_SafetensorsFile_ctor(void) {
61 ✗ return new SafetensorsFileHandle();
62 }
63
64 4 void bro_tensor_SafetensorsFile_dtor(void* self) {
65
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 4 times.
 4 delete asFile(self);
66 4 }
67
68 10 void* bro_tensor_openSafetensors(const char* path) {
69 10 auto* h = new SafetensorsFileHandle();
70 try {
71
6/10
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 10 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 10 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 8 times.
✓ Branch 7 taken 2 times.
✓ Branch 8 taken 8 times.
✗ Branch 9 not taken.
 12 h->file = std::make_unique<stns::File>(stns::File::open(resolvePath(path ? path : "")));
72
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 10 } catch (const std::exception& e) {
73
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 delete h;
74
3/6
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
 2 setError(std::string("openSafetensors: ") + e.what());
75 2 return nullptr;
76 2 }
77 8 return h;
78 12 }
79
80 42 bool bro_tensor_SafetensorsFile_isOpen(void* self) {
81 42 auto* h = asFile(self);
82
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 42 times.
 42 return h && h->file != nullptr;
83 }
84
85 26 int32_t bro_tensor_SafetensorsFile_count_get(void* self) {
86 26 auto* h = asFile(self);
87
3/4
✓ Branch 0 taken 26 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 24 times.
 26 return (h && h->file) ? static_cast<int32_t>(h->file->size()) : 0;
88 }
89
90 38 const char* bro_tensor_SafetensorsFile_nameAt(void* self, int32_t i) {
91 38 const auto* tv = viewAt(self, i, "names");
92
1/2
✓ Branch 0 taken 38 times.
✗ Branch 1 not taken.
 38 return tv ? tv->name.c_str() : "";
93 }
94
95 34 const char* bro_tensor_SafetensorsFile_dtypeAt(void* self, int32_t i) {
96 34 const auto* tv = viewAt(self, i, "header");
97
1/2
✓ Branch 0 taken 34 times.
✗ Branch 1 not taken.
 34 return tv ? stns::dtype_name(tv->dtype) : "";
98 }
99
100 34 void bro_tensor_SafetensorsFile_shapeAt(void* self, int32_t i, bronze_native_buffer* out) {
101
1/2
✓ Branch 0 taken 34 times.
✗ Branch 1 not taken.
 34 if (!out) return;
102 34 out->data = nullptr;
103 34 out->length = 0;
104 34 out->release = nullptr;
105 34 out->ctx = nullptr;
106 34 const auto* tv = viewAt(self, i, "header");
107
1/2
✓ Branch 0 taken 34 times.
✗ Branch 1 not taken.
 34 if (!tv) return;
108 34 tl_shape.clear();
109 34 tl_shape.reserve(tv->shape.size());
110
2/2
✓ Branch 0 taken 72 times.
✓ Branch 1 taken 34 times.
 106 for (int64_t d : tv->shape) tl_shape.push_back(static_cast<double>(d));
111 34 out->data = tl_shape.data();
112 34 out->length = static_cast<uint32_t>(tl_shape.size());
113 34 }
114
115 34 double bro_tensor_SafetensorsFile_nbytesAt(void* self, int32_t i) {
116 34 const auto* tv = viewAt(self, i, "header");
117
1/2
✓ Branch 0 taken 34 times.
✗ Branch 1 not taken.
 34 return tv ? static_cast<double>(tv->nbytes) : 0.0;
118 }
119
120 // get(name, rows, cols, mode): the old body's shape rule — an explicit
121 // (rows, cols) pair when both were numbers, else (shape[0], numel/shape[0]);
122 // mode "compute" -> upload_compute, "fp16" -> upload_fp16, anything else ->
123 // upload (the file's own dtype).
124 18 void* bro_tensor_SafetensorsFile_get(void* self, const char* name, int32_t rows, int32_t cols, const char* mode) {
125 18 auto* h = asFile(self);
126
2/4
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 18 times.
 18 if (!h || !h->file) {
127 ✗ setError("get() on a closed safetensors file");
128 ✗ return nullptr;
129 }
130
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 18 const std::string key(name ? name : "");
131
1/2
✓ Branch 0 taken 18 times.
✗ Branch 1 not taken.
 18 const stns::TensorView* tv = h->file->find(key);
132
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 16 times.
 18 if (!tv) {
133
3/6
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
 2 setError("get: no tensor named '" + key + "'");
134 2 return nullptr;
135 }
136 16 int r = rows;
137 16 int c = cols;
138
3/4
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 12 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 4 times.
 16 if (r <= 0 || c <= 0) {
139 try {
140
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 const int64_t numel = tv->numel();
141
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 12 times.
 12 r = tv->shape.empty() ? 1 : static_cast<int>(tv->shape[0]);
142
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 c = r > 0 ? static_cast<int>(numel / r) : 0;
143
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 12 } catch (const std::exception& e) {
144 ✗ setError(std::string("get: ") + e.what());
145 ✗ return nullptr;
146 ✗ }
147 12 }
148
1/2
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
 16 auto* gt = new GpuTensorHandle();
149
6/10
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 12 times.
✓ Branch 2 taken 4 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 12 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 4 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 16 times.
✗ Branch 9 not taken.
 16 gt->shape = (rows <= 0 || cols <= 0) ? tv->shape : std::vector<int64_t>{rows, cols};
150 try {
151
3/4
✓ Branch 0 taken 16 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 14 times.
 16 if (mode && std::strcmp(mode, "compute") == 0) {
152
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 2 stns::upload_compute(*tv, r, c, gt->tensor);
153
3/4
✓ Branch 0 taken 14 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 4 times.
✓ Branch 3 taken 10 times.
 16 } else if (mode && std::strcmp(mode, "fp16") == 0) {
154
1/2
✓ Branch 0 taken 4 times.
✗ Branch 1 not taken.
 4 stns::upload_fp16(*tv, r, c, gt->tensor);
155 4 } else {
156
1/2
✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
 10 stns::upload(*tv, r, c, gt->tensor);
157 }
158
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 16 } catch (const std::exception& e) {
159 ✗ delete gt;
160 ✗ setError(std::string("get: ") + e.what());
161 ✗ return nullptr;
162 ✗ }
163 16 return gt;
164 18 }
165
166 8 void bro_tensor_SafetensorsFile_close(void* self) {
167 8 auto* h = asFile(self);
168
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (h) h->file.reset();
169 8 }
170
171 6 bool bro_tensor_saveSafetensors(const char* path, uint64_t names_bits, uint64_t tensors_bits) {
172 6 ev::Persistent names(bronze::Value{names_bits});
173
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 ev::Persistent tensors(bronze::Value{tensors_bits});
174
6/12
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 6 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 6 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 6 times.
✗ Branch 11 not taken.
 6 if (!ev::isObject(names.get()) || !ev::isObject(tensors.get())) {
175 ✗ setError("saveSafetensors(path, {name: tensor, ...})");
176 ✗ return false;
177 }
178
4/8
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 6 times.
✗ Branch 7 not taken.
 6 const auto n = static_cast<uint32_t>(ev::toDouble(ev::getProperty(names.get(), "length")));
179
180 6 std::vector<stns::WriteEntry> entries;
181 // Backing host buffers, alive until write_file() returns; reserved so
182 // the inner vectors never reallocate out from under host_data.
183 6 std::vector<std::vector<float>> f32store;
184 6 std::vector<std::vector<uint16_t>> f16store;
185
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 entries.reserve(n);
186
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 f32store.reserve(n);
187
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 f16store.reserve(n);
188
189 try {
190
1/2
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
 6 brotensor::sync_all();
191
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 6 times.
 14 for (uint32_t i = 0; i < n; ++i) {
192
3/6
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 8 times.
✗ Branch 5 not taken.
 8 const std::string key = ev::toUtf8(ev::getElement(names.get(), i));
193
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 Value tVal = ev::getElement(tensors.get(), i);
194
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 void* d = ev::handleData(tVal);
195
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 auto* gh = d ? toHandle(d) : nullptr;
196
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 auto* gt = gh ? &gh->tensor : nullptr;
197
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (!gt) {
198 ✗ setError("saveSafetensors: value for '" + key + "' is not a tensor");
199 ✗ return false;
200 }
201 8 stns::WriteEntry e;
202
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 e.name = key;
203
204 // Preserve full n-dimensional tensor shape if present. The shape
205 // array is read across several allocating getProperty/getElement
206 // calls, so it lives in a Persistent: a raw Value is stale after
207 // the first of them.
208 8 bool shapeSet = false;
209
3/6
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 8 times.
✗ Branch 5 not taken.
 8 ev::Persistent shapeProp(ev::getProperty(tVal, "shape"));
210
3/6
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 8 times.
✗ Branch 5 not taken.
 8 if (ev::isObject(shapeProp.get())) {
211
3/6
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 8 times.
✗ Branch 5 not taken.
 8 Value lenVal = ev::getProperty(shapeProp.get(), "length");
212
2/4
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✗ Branch 3 not taken.
 8 if (ev::isNumber(lenVal)) {
213
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 uint32_t shapeLen = static_cast<uint32_t>(ev::toDouble(lenVal));
214
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (shapeLen > 0) {
215 8 std::vector<int64_t> customShape;
216
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 customShape.reserve(shapeLen);
217 8 int64_t numel = 1;
218
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 24 times.
 32 for (uint32_t s = 0; s < shapeLen; ++s) {
219
3/6
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 24 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 24 times.
✗ Branch 5 not taken.
 24 int64_t dim = static_cast<int64_t>(ev::toDouble(ev::getElement(shapeProp.get(), s)));
220
1/2
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
 24 customShape.push_back(dim);
221 // A negative dim can multiply back to rows*cols
222 // ([-2, -3] for 6); never let it reach the header.
223
1/2
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
 24 numel = dim < 0 ? -1 : numel * dim;
224
1/2
✓ Branch 0 taken 24 times.
✗ Branch 1 not taken.
 24 if (numel < 0) break;
225 24 }
226
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (numel == static_cast<int64_t>(gt->rows) * gt->cols) {
227 8 e.shape = std::move(customShape);
228 8 shapeSet = true;
229 8 }
230 8 }
231 8 }
232 8 }
233
234
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 8 if (!shapeSet && gh && !gh->shape.empty()) {
235 ✗ int64_t numel = 1;
236 ✗ for (int64_t dim : gh->shape) numel *= dim;
237 ✗ if (numel == static_cast<int64_t>(gt->rows) * gt->cols) {
238 ✗ e.shape = gh->shape;
239 ✗ shapeSet = true;
240 ✗ }
241 ✗ }
242
243
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 8 times.
 8 if (!shapeSet) {
244 ✗ e.shape = {gt->rows, gt->cols};
245 ✗ }
246
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 brotensor::Tensor host = gt->to(brotensor::Device::CPU);
247
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 2 times.
 8 if (host.dtype == brotensor::Dtype::FP32) {
248
2/4
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 6 times.
 6 f32store.push_back(host.to_host_vector());
249 6 e.dtype = stns::Dtype::F32;
250 6 e.host_data = f32store.back().data();
251 6 e.bytes = f32store.back().size() * sizeof(float);
252
1/2
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
 8 } else if (host.dtype == brotensor::Dtype::FP16) {
253
2/4
✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✗ Branch 3 not taken.
 2 f16store.push_back(host.to_host_vector_fp16());
254 2 e.dtype = stns::Dtype::F16;
255 2 e.host_data = f16store.back().data();
256 2 e.bytes = f16store.back().size() * sizeof(uint16_t);
257 2 } else {
258 ✗ setError("saveSafetensors: '" + key + "' has an unsupported dtype (FP32/FP16 only)");
259 ✗ return false;
260 }
261
1/2
✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
 8 entries.push_back(std::move(e));
262 8 }
263
4/8
✓ Branch 0 taken 6 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 6 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 6 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 6 times.
✗ Branch 7 not taken.
 6 stns::write_file(resolvePath(path ? path : ""), entries);
264
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 6 } catch (const std::exception& e) {
265 ✗ setError(std::string("saveSafetensors: ") + e.what());
266 ✗ return false;
267 ✗ }
268 6 return true;
269 6 }
270
271 } // extern "C"
272