Branch data Line data Source code
1 : : /*
2 : : * Copyright (C) 2026 by Thun Lu. All rights reserved.
3 : : * Author: Thun Lu <thun.lu@zohomail.cn>
4 : : * Repo: https://github.com/thun-res/vlink
5 : : * _ __ __ _ __
6 : : * | | / / / / (_) ____ / /__
7 : : * | | / / / / / / / __ \ / //_/
8 : : * | |/ / / /___ / / / / / / / ,<
9 : : * |___/ /_____/ /_/ /_/ /_/ /_/|_|
10 : : *
11 : : * Licensed under the Apache License, Version 2.0 (the "License");
12 : : * you may not use this file except in compliance with the License.
13 : : * You may obtain a copy of the License at
14 : : *
15 : : * http://www.apache.org/licenses/LICENSE-2.0
16 : : *
17 : : * Unless required by applicable law or agreed to in writing, software
18 : : * distributed under the License is distributed on an "AS IS" BASIS,
19 : : * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
20 : : * See the License for the specific language governing permissions and
21 : : * limitations under the License.
22 : : */
23 : :
24 : : #include "./zerocopy/tensor.h"
25 : :
26 : : #include <cstdint>
27 : :
28 : : namespace vlink {
29 : :
30 : : namespace zerocopy {
31 : :
32 : : // Tensor
33 : 67 : Tensor::Tensor() noexcept {
34 : : #if defined(__arm__) || defined(__x86__) || defined(__i386__)
35 : : #ifndef __ANDROID__
36 : : #warning "[Tensor] No support for 32-bit architecture."
37 : : #endif
38 : : #else
39 : : static_assert(sizeof(Tensor) == 248, "Sizeof must be 248 bytes.");
40 : : #endif
41 : 67 : }
42 : :
43 : 69 : Tensor::~Tensor() noexcept {
44 [ + + + - : 69 : if (is_owner_ && data_ && size_ != 0) {
+ - ]
45 : 33 : Bytes::bytes_free(data_, size_);
46 : : }
47 : 69 : }
48 : :
49 : 1 : Tensor::Tensor(const Tensor& target) noexcept { deep_copy(target); }
50 : :
51 : 1 : Tensor::Tensor(Tensor&& target) noexcept { move_copy(target); }
52 : :
53 : 1 : Tensor& Tensor::operator=(const Tensor& target) noexcept {
54 [ - + ]: 1 : if VUNLIKELY (this == &target) {
55 : : return *this; // LCOV_EXCL_LINE GCOVR_EXCL_LINE
56 : : }
57 : :
58 : 1 : deep_copy(target);
59 : :
60 : 1 : return *this;
61 : : }
62 : :
63 : 1 : Tensor& Tensor::operator=(Tensor&& target) noexcept {
64 [ - + ]: 1 : if VUNLIKELY (this == &target) {
65 : : return *this; // LCOV_EXCL_LINE GCOVR_EXCL_LINE
66 : : }
67 : :
68 : 1 : move_copy(target);
69 : :
70 : 1 : return *this;
71 : : }
72 : :
73 : 9 : bool Tensor::operator<<(const Bytes& bytes) noexcept {
74 : : static constexpr size_t kMagicNumberBeginSize = sizeof(kMagicNumberBegin);
75 : : static constexpr size_t kVersionSize = sizeof(kWireVersion);
76 : : // static constexpr size_t kMagicNumberEndSize = sizeof(kMagicNumberEnd);
77 : :
78 [ + + ]: 9 : if VUNLIKELY (bytes.empty()) {
79 : 1 : return false;
80 : : }
81 : :
82 [ + + ]: 8 : if VUNLIKELY (!check_valid(bytes)) {
83 : 2 : return false;
84 : : }
85 : :
86 : 6 : uint32_t wire_version = 0;
87 : 6 : std::memcpy(&wire_version, bytes.data() + kMagicNumberBeginSize, kVersionSize);
88 : :
89 [ - + ]: 6 : if VUNLIKELY (version_major(wire_version) != version_major(kWireVersion)) {
90 : : return false; // LCOV_EXCL_LINE GCOVR_EXCL_LINE
91 : : }
92 : :
93 [ + + + - : 6 : if (is_owner_ && data_ && size_ != 0) {
+ - ]
94 : 1 : Bytes::bytes_free(data_, size_);
95 : : }
96 : :
97 : : #if defined(__GNUC__) && !defined(__clang__)
98 : : #pragma GCC diagnostic push
99 : : #pragma GCC diagnostic ignored "-Wclass-memaccess"
100 : : #if __GNUC__ >= 11
101 : : #pragma GCC diagnostic ignored "-Wstringop-overread"
102 : : #endif
103 : : #endif
104 : :
105 : 6 : auto* target_ptr = reinterpret_cast<uint8_t*>(this);
106 : :
107 : 6 : std::memcpy(target_ptr, bytes.data() + kMagicNumberBeginSize + kVersionSize, sizeof(Tensor));
108 : :
109 : : #if defined(__GNUC__) && !defined(__clang__)
110 : : #pragma GCC diagnostic pop
111 : : #endif
112 : :
113 : 6 : data_ = const_cast<uint8_t*>(bytes.data() + kMagicNumberBeginSize + kVersionSize + sizeof(Tensor));
114 : 6 : is_owner_ = false;
115 : :
116 [ + + ]: 6 : if VUNLIKELY (bytes.size() != get_serialized_size()) {
117 : 1 : clear();
118 : 1 : return false;
119 : : }
120 : :
121 [ + + ]: 5 : if VUNLIKELY (rank_ > kMaxRank) {
122 : 1 : rank_ = kMaxRank;
123 : : }
124 : :
125 : 5 : element_size_ = element_size_of(dtype_);
126 : :
127 : 5 : return true;
128 : : }
129 : :
130 : 10 : bool Tensor::operator>>(Bytes& bytes) const noexcept {
131 : : static constexpr size_t kMagicNumberBeginSize = sizeof(kMagicNumberBegin);
132 : : static constexpr size_t kVersionSize = sizeof(kWireVersion);
133 : : static constexpr size_t kMagicNumberEndSize = sizeof(kMagicNumberEnd);
134 : :
135 [ + + - + : 10 : if (bytes.empty() || bytes.size() != get_serialized_size()) {
+ + ]
136 : 9 : bytes = Bytes::create(get_serialized_size());
137 : : }
138 : :
139 : 10 : std::memcpy(bytes.data(), &kMagicNumberBegin, kMagicNumberBeginSize);
140 : :
141 : 10 : std::memcpy(bytes.data() + kMagicNumberBeginSize, &kWireVersion, kVersionSize);
142 : :
143 : : // NOLINTNEXTLINE(bugprone-undefined-memory-manipulation)
144 : 10 : std::memcpy(bytes.data() + kMagicNumberBeginSize + kVersionSize, this, sizeof(Tensor));
145 : :
146 [ + + + - : 10 : if VLIKELY (data_ != nullptr && size_ != 0) {
+ + ]
147 : 9 : std::memcpy(bytes.data() + kMagicNumberBeginSize + kVersionSize + sizeof(Tensor), data_, size_);
148 : : }
149 : :
150 : 10 : std::memcpy(bytes.data() + kMagicNumberBeginSize + kVersionSize + sizeof(Tensor) + size_, &kMagicNumberEnd,
151 : : kMagicNumberEndSize);
152 : :
153 : 10 : return true;
154 : : }
155 : :
156 : 15 : bool Tensor::check_valid(const Bytes& bytes) noexcept {
157 : : static constexpr size_t kMagicNumberBeginSize = sizeof(kMagicNumberBegin);
158 : : static constexpr size_t kVersionSize = sizeof(kWireVersion);
159 : : static constexpr size_t kMagicNumberEndSize = sizeof(kMagicNumberEnd);
160 : :
161 [ + + ]: 15 : if VUNLIKELY (bytes.size() < kMagicNumberBeginSize + kVersionSize + sizeof(Tensor) + kMagicNumberEndSize) {
162 : 2 : return false;
163 : : }
164 : :
165 : 13 : uint32_t check_magic = 0;
166 : :
167 : 13 : std::memcpy(&check_magic, bytes.begin(), kMagicNumberBeginSize);
168 : :
169 [ + + ]: 13 : if VUNLIKELY (check_magic != kMagicNumberBegin) {
170 : 1 : return false;
171 : : }
172 : :
173 : 12 : uint32_t wire_version = 0;
174 : 12 : std::memcpy(&wire_version, bytes.data() + kMagicNumberBeginSize, kVersionSize);
175 : :
176 [ + + ]: 12 : if VUNLIKELY (version_major(wire_version) != version_major(kWireVersion)) {
177 : 2 : return false;
178 : : }
179 : :
180 : 10 : std::memcpy(&check_magic, bytes.end() - kMagicNumberEndSize, kMagicNumberEndSize);
181 : :
182 [ + + ]: 10 : if VUNLIKELY (check_magic != kMagicNumberEnd) {
183 : 1 : return false;
184 : : }
185 : :
186 : 9 : return true;
187 : : }
188 : :
189 : 20 : size_t Tensor::get_serialized_size() const noexcept {
190 : : static constexpr size_t kMagicNumberBeginSize = sizeof(kMagicNumberBegin);
191 : : static constexpr size_t kVersionSize = sizeof(kWireVersion);
192 : : static constexpr size_t kMagicNumberEndSize = sizeof(kMagicNumberEnd);
193 : :
194 : 20 : return kMagicNumberBeginSize + kVersionSize + sizeof(Tensor) + size_ + kMagicNumberEndSize;
195 : : }
196 : :
197 [ + + + - ]: 19 : bool Tensor::is_valid() const noexcept { return data_ != nullptr && size_ != 0; }
198 : :
199 : 9 : bool Tensor::shallow_copy(const Tensor& target) noexcept {
200 [ + + ]: 9 : if VUNLIKELY (this == &target) {
201 : 2 : return false;
202 : : }
203 : :
204 [ - + - - : 7 : if (is_owner_ && data_ && size_ != 0) {
- - ]
205 : 0 : Bytes::bytes_free(data_, size_);
206 : : }
207 : :
208 : 7 : header = target.header;
209 : :
210 : 7 : update_time_ns_ = target.update_time_ns_;
211 : 7 : num_elements_ = target.num_elements_;
212 : 7 : std::memcpy(name_, target.name_, sizeof(name_));
213 : 7 : std::memcpy(model_id_, target.model_id_, sizeof(model_id_));
214 : 7 : std::memcpy(layout_, target.layout_, sizeof(layout_));
215 : 7 : std::memcpy(shape_, target.shape_, sizeof(shape_));
216 : 7 : std::memcpy(strides_, target.strides_, sizeof(strides_));
217 : 7 : channel_ = target.channel_;
218 : 7 : freq_ = target.freq_;
219 : 7 : batch_size_ = target.batch_size_;
220 : 7 : quant_scale_ = target.quant_scale_;
221 : 7 : quant_zero_point_ = target.quant_zero_point_;
222 : 7 : dtype_ = target.dtype_;
223 : 7 : rank_ = target.rank_;
224 : 7 : device_ = target.device_;
225 : 7 : element_size_ = target.element_size_;
226 : 7 : reserved_buf_ = target.reserved_buf_;
227 : 7 : reserved_buf2_ = target.reserved_buf2_;
228 : 7 : reserved_buf3_ = target.reserved_buf3_;
229 : 7 : is_owner_ = false;
230 : 7 : data_ = target.data_;
231 : 7 : size_ = target.size_;
232 : :
233 : 7 : return true;
234 : : }
235 : :
236 : 5 : bool Tensor::deep_copy(const Tensor& target) noexcept {
237 [ + + + - : 5 : if VLIKELY (data_ && is_owner_ && target.data_ && size_ != 0 && size_ == target.size_) {
+ + + - +
+ + - + +
+ - + + ]
238 [ + + ]: 2 : if VUNLIKELY (this == &target) {
239 : 1 : return false;
240 : : }
241 : :
242 : 1 : header = target.header;
243 : :
244 : 1 : update_time_ns_ = target.update_time_ns_;
245 : 1 : num_elements_ = target.num_elements_;
246 : 1 : std::memcpy(name_, target.name_, sizeof(name_));
247 : 1 : std::memcpy(model_id_, target.model_id_, sizeof(model_id_));
248 : 1 : std::memcpy(layout_, target.layout_, sizeof(layout_));
249 : 1 : std::memcpy(shape_, target.shape_, sizeof(shape_));
250 : 1 : std::memcpy(strides_, target.strides_, sizeof(strides_));
251 : 1 : channel_ = target.channel_;
252 : 1 : freq_ = target.freq_;
253 : 1 : batch_size_ = target.batch_size_;
254 : 1 : quant_scale_ = target.quant_scale_;
255 : 1 : quant_zero_point_ = target.quant_zero_point_;
256 : 1 : dtype_ = target.dtype_;
257 : 1 : rank_ = target.rank_;
258 : 1 : device_ = target.device_;
259 : 1 : element_size_ = target.element_size_;
260 : 1 : reserved_buf_ = target.reserved_buf_;
261 : 1 : reserved_buf2_ = target.reserved_buf2_;
262 : 1 : reserved_buf3_ = target.reserved_buf3_;
263 : :
264 : 1 : std::memcpy(data_, target.data_, size_);
265 : :
266 : 1 : return true;
267 : : }
268 : :
269 [ - + ]: 3 : if VUNLIKELY (!shallow_copy(target)) {
270 : : return false; // LCOV_EXCL_LINE GCOVR_EXCL_LINE
271 : : }
272 : :
273 [ + - + - ]: 3 : if (data_ && size_ != 0) {
274 : 3 : data_ = Bytes::bytes_malloc(size_);
275 : :
276 : 3 : std::memcpy(data_, target.data_, size_);
277 : :
278 : 3 : is_owner_ = true;
279 : : }
280 : :
281 : 3 : return true;
282 : : }
283 : :
284 : 4 : bool Tensor::move_copy(Tensor& target) noexcept {
285 [ + + ]: 4 : if VUNLIKELY (!shallow_copy(target)) {
286 : 1 : return false;
287 : : }
288 : :
289 : 3 : is_owner_ = target.is_owner_;
290 : :
291 : 3 : target.update_time_ns_ = 0;
292 : 3 : target.num_elements_ = 0;
293 : 3 : std::memset(target.name_, 0, sizeof(target.name_));
294 : 3 : std::memset(target.model_id_, 0, sizeof(target.model_id_));
295 : 3 : std::memset(target.layout_, 0, sizeof(target.layout_));
296 : 3 : std::memset(target.shape_, 0, sizeof(target.shape_));
297 : 3 : std::memset(target.strides_, 0, sizeof(target.strides_));
298 : 3 : target.channel_ = 0;
299 : 3 : target.freq_ = 0;
300 : 3 : target.batch_size_ = 0;
301 : 3 : target.quant_scale_ = 0;
302 : 3 : target.quant_zero_point_ = 0;
303 : 3 : target.dtype_ = kDataUnknown;
304 : 3 : target.rank_ = 0;
305 : 3 : target.device_ = kDeviceCpu;
306 : 3 : target.element_size_ = 0;
307 : 3 : target.reserved_buf_ = 0;
308 : 3 : target.reserved_buf2_ = 0;
309 : 3 : target.reserved_buf3_ = 0;
310 : 3 : target.is_owner_ = false;
311 : 3 : target.data_ = nullptr;
312 : 3 : target.size_ = 0;
313 : :
314 : : #if defined(__GNUC__) && !defined(__clang__)
315 : : #pragma GCC diagnostic push
316 : : #pragma GCC diagnostic ignored "-Wclass-memaccess"
317 : : #if __GNUC__ >= 11
318 : : #pragma GCC diagnostic ignored "-Wstringop-overread"
319 : : #endif
320 : : #endif
321 : :
322 : 3 : std::memset(&target.header, 0, sizeof(header));
323 : :
324 : : #if defined(__GNUC__) && !defined(__clang__)
325 : : #pragma GCC diagnostic pop
326 : : #endif
327 : :
328 : 3 : return true;
329 : : }
330 : :
331 : 38 : bool Tensor::create(size_t _size) noexcept {
332 [ + + ]: 38 : if VUNLIKELY (_size == 0) {
333 : 1 : return false;
334 : : }
335 : :
336 [ + + + - : 37 : if (is_owner_ && data_ && size_ != 0) {
+ - ]
337 : 2 : Bytes::bytes_free(data_, size_);
338 : : }
339 : :
340 : 37 : size_ = _size;
341 : :
342 : 37 : data_ = Bytes::bytes_malloc(size_);
343 : :
344 : 37 : is_owner_ = true;
345 : :
346 : 37 : return true;
347 : : }
348 : :
349 : 5 : void Tensor::clear() noexcept {
350 [ + + + - : 5 : if (is_owner_ && data_ && size_ != 0) {
+ - ]
351 : 4 : Bytes::bytes_free(data_, size_);
352 : : }
353 : :
354 : 5 : update_time_ns_ = 0;
355 : 5 : num_elements_ = 0;
356 : 5 : std::memset(name_, 0, sizeof(name_));
357 : 5 : std::memset(model_id_, 0, sizeof(model_id_));
358 : 5 : std::memset(layout_, 0, sizeof(layout_));
359 : 5 : std::memset(shape_, 0, sizeof(shape_));
360 : 5 : std::memset(strides_, 0, sizeof(strides_));
361 : 5 : channel_ = 0;
362 : 5 : freq_ = 0;
363 : 5 : batch_size_ = 0;
364 : 5 : quant_scale_ = 0;
365 : 5 : quant_zero_point_ = 0;
366 : 5 : dtype_ = kDataUnknown;
367 : 5 : rank_ = 0;
368 : 5 : device_ = kDeviceCpu;
369 : 5 : element_size_ = 0;
370 : 5 : is_owner_ = false;
371 : 5 : data_ = nullptr;
372 : 5 : size_ = 0;
373 : :
374 : : #if defined(__GNUC__) && !defined(__clang__)
375 : : #pragma GCC diagnostic push
376 : : #pragma GCC diagnostic ignored "-Wclass-memaccess"
377 : : #if __GNUC__ >= 11
378 : : #pragma GCC diagnostic ignored "-Wstringop-overread"
379 : : #endif
380 : : #endif
381 : :
382 : 5 : std::memset(&header, 0, sizeof(header));
383 : :
384 : : #if defined(__GNUC__) && !defined(__clang__)
385 : : #pragma GCC diagnostic pop
386 : : #endif
387 : 5 : }
388 : :
389 : 5 : bool Tensor::shallow_copy(uint8_t* data, size_t size) noexcept {
390 [ + + + + : 5 : if VUNLIKELY (!data || size == 0) {
+ + ]
391 : 2 : return false;
392 : : }
393 : :
394 [ + + ]: 3 : if VUNLIKELY (data_ == data) {
395 : 1 : return false;
396 : : }
397 : :
398 [ - + - - : 2 : if (is_owner_ && data_ && size_ != 0) {
- - ]
399 : 0 : Bytes::bytes_free(data_, size_);
400 : : }
401 : :
402 : 2 : is_owner_ = false;
403 : :
404 : 2 : data_ = data;
405 : 2 : size_ = size;
406 : :
407 : 2 : return true;
408 : : }
409 : :
410 : 6 : bool Tensor::deep_copy(uint8_t* data, size_t size) noexcept {
411 [ + + + + : 6 : if VUNLIKELY (!data || size == 0) {
+ + ]
412 : 2 : return false;
413 : : }
414 : :
415 [ + + ]: 4 : if (is_owner_) {
416 [ - + ]: 2 : if VUNLIKELY (!data_) {
417 : : return false; // LCOV_EXCL_LINE GCOVR_EXCL_LINE
418 : : }
419 : :
420 [ + + ]: 2 : if VUNLIKELY (data_ == data) {
421 : 1 : return false;
422 : : }
423 : :
424 [ + - ]: 1 : if VUNLIKELY (size_ != size) {
425 : 1 : create(size);
426 : : }
427 : : } else {
428 : 2 : create(size);
429 : : }
430 : :
431 : 3 : std::memcpy(data_, data, size);
432 : :
433 : 3 : return true;
434 : : }
435 : :
436 : 1 : bool Tensor::fill_data(uint8_t* data, size_t size) noexcept { return deep_copy(data, size); }
437 : :
438 : 3 : uint64_t Tensor::update_time_ns() const noexcept { return update_time_ns_; }
439 : :
440 : 6 : uint64_t Tensor::num_elements() const noexcept { return num_elements_; }
441 : :
442 : 9 : std::string_view Tensor::name() const noexcept { return {name_, ::strnlen(name_, sizeof(name_))}; }
443 : :
444 : 7 : std::string_view Tensor::model_id() const noexcept { return {model_id_, ::strnlen(model_id_, sizeof(model_id_))}; }
445 : :
446 : 8 : std::string_view Tensor::layout() const noexcept { return {layout_, ::strnlen(layout_, sizeof(layout_))}; }
447 : :
448 : 0 : const uint32_t* Tensor::shape() const noexcept { return shape_; }
449 : :
450 : 19 : uint32_t Tensor::shape_at(uint8_t dim) const noexcept {
451 [ + + ]: 19 : if VUNLIKELY (dim >= kMaxRank) {
452 : 1 : return 0;
453 : : }
454 : :
455 : 18 : return shape_[dim];
456 : : }
457 : :
458 : 0 : const uint32_t* Tensor::strides() const noexcept { return strides_; }
459 : :
460 : 12 : uint32_t Tensor::stride_at(uint8_t dim) const noexcept {
461 [ + + ]: 12 : if VUNLIKELY (dim >= kMaxRank) {
462 : 1 : return 0;
463 : : }
464 : :
465 : 11 : return strides_[dim];
466 : : }
467 : :
468 : 3 : uint32_t Tensor::channel() const noexcept { return channel_; }
469 : :
470 : 3 : uint32_t Tensor::freq() const noexcept { return freq_; }
471 : :
472 : 5 : uint32_t Tensor::batch_size() const noexcept { return batch_size_; }
473 : :
474 : 3 : float Tensor::quant_scale() const noexcept { return quant_scale_; }
475 : :
476 : 3 : int32_t Tensor::quant_zero_point() const noexcept { return quant_zero_point_; }
477 : :
478 : 12 : Tensor::DataType Tensor::dtype() const noexcept { return dtype_; }
479 : :
480 : 17 : uint8_t Tensor::rank() const noexcept { return rank_; }
481 : :
482 : 8 : Tensor::Device Tensor::device() const noexcept { return device_; }
483 : :
484 : 12 : uint8_t Tensor::element_size() const noexcept { return element_size_; }
485 : :
486 : 38 : const uint8_t* Tensor::data() const noexcept { return data_; }
487 : :
488 : 28 : size_t Tensor::size() const noexcept { return size_; }
489 : :
490 : 22 : bool Tensor::is_owner() const noexcept { return is_owner_; }
491 : :
492 : 2 : void Tensor::set_update_time_ns(uint64_t update_time_ns) noexcept { update_time_ns_ = update_time_ns; }
493 : :
494 : 8 : void Tensor::set_name(std::string_view name) noexcept {
495 : 8 : std::memset(name_, 0, sizeof(name_));
496 : :
497 : 8 : size_t copy_size = name.size();
498 : :
499 [ + + ]: 8 : if (copy_size >= sizeof(name_)) {
500 : 1 : copy_size = sizeof(name_) - 1;
501 : : }
502 : :
503 [ + - ]: 8 : if VLIKELY (copy_size != 0) {
504 : 8 : std::memcpy(name_, name.data(), copy_size);
505 : : }
506 : 8 : }
507 : :
508 : 6 : void Tensor::set_model_id(std::string_view model_id) noexcept {
509 : 6 : std::memset(model_id_, 0, sizeof(model_id_));
510 : :
511 : 6 : size_t copy_size = model_id.size();
512 : :
513 [ + + ]: 6 : if (copy_size >= sizeof(model_id_)) {
514 : 1 : copy_size = sizeof(model_id_) - 1;
515 : : }
516 : :
517 [ + - ]: 6 : if VLIKELY (copy_size != 0) {
518 : 6 : std::memcpy(model_id_, model_id.data(), copy_size);
519 : : }
520 : 6 : }
521 : :
522 : 7 : void Tensor::set_layout(std::string_view layout) noexcept {
523 : 7 : std::memset(layout_, 0, sizeof(layout_));
524 : :
525 : 7 : size_t copy_size = layout.size();
526 : :
527 [ + + ]: 7 : if (copy_size >= sizeof(layout_)) {
528 : 1 : copy_size = sizeof(layout_) - 1;
529 : : }
530 : :
531 [ + - ]: 7 : if VLIKELY (copy_size != 0) {
532 : 7 : std::memcpy(layout_, layout.data(), copy_size);
533 : : }
534 : 7 : }
535 : :
536 : 18 : void Tensor::set_shape(const uint32_t* shape, uint8_t rank) noexcept {
537 : 18 : std::memset(shape_, 0, sizeof(shape_));
538 : 18 : std::memset(strides_, 0, sizeof(strides_));
539 : :
540 [ + + + + : 18 : if VUNLIKELY (rank == 0 || !shape) {
+ + ]
541 : 2 : rank_ = 0;
542 : 2 : num_elements_ = 0;
543 : 2 : batch_size_ = 0;
544 : 2 : return;
545 : : }
546 : :
547 [ + + ]: 16 : if (rank > kMaxRank) {
548 : 1 : rank = kMaxRank;
549 : : }
550 : :
551 : 16 : rank_ = rank;
552 : :
553 : 16 : uint64_t total = 1;
554 : :
555 [ + + ]: 60 : for (uint8_t i = 0; i < rank; ++i) {
556 : 44 : shape_[i] = shape[i];
557 : 44 : total *= shape[i];
558 : : }
559 : :
560 : 16 : num_elements_ = total;
561 : 16 : batch_size_ = shape_[0];
562 : :
563 : 16 : uint32_t running = 1;
564 : :
565 [ + + ]: 60 : for (uint8_t i = rank; i > 0; --i) {
566 : 44 : strides_[i - 1] = running;
567 : 44 : running *= shape_[i - 1];
568 : : }
569 : : }
570 : :
571 : 5 : void Tensor::set_shape_at(uint8_t dim, uint32_t value) noexcept {
572 [ + + ]: 5 : if VUNLIKELY (dim >= kMaxRank) {
573 : 1 : return;
574 : : }
575 : :
576 : 4 : shape_[dim] = value;
577 : : }
578 : :
579 : 5 : void Tensor::set_stride_at(uint8_t dim, uint32_t value) noexcept {
580 [ + + ]: 5 : if VUNLIKELY (dim >= kMaxRank) {
581 : 1 : return;
582 : : }
583 : :
584 : 4 : strides_[dim] = value;
585 : : }
586 : :
587 : 2 : void Tensor::set_channel(uint32_t channel) noexcept { channel_ = channel; }
588 : :
589 : 2 : void Tensor::set_freq(uint32_t freq) noexcept { freq_ = freq; }
590 : :
591 : 1 : void Tensor::set_batch_size(uint32_t batch_size) noexcept { batch_size_ = batch_size; }
592 : :
593 : 3 : void Tensor::set_quant_scale(float quant_scale) noexcept { quant_scale_ = quant_scale; }
594 : :
595 : 3 : void Tensor::set_quant_zero_point(int32_t quant_zero_point) noexcept { quant_zero_point_ = quant_zero_point; }
596 : :
597 : 15 : void Tensor::set_dtype(DataType dtype) noexcept {
598 : 15 : dtype_ = dtype;
599 : 15 : element_size_ = element_size_of(dtype);
600 : 15 : }
601 : :
602 : 7 : void Tensor::set_device(Device device) noexcept { device_ = device; }
603 : :
604 : 34 : uint8_t Tensor::element_size_of(DataType dtype) noexcept {
605 : 34 : uint8_t target_size = 0;
606 : :
607 [ + + + + : 34 : if (dtype == kBool || dtype == kInt8 || dtype == kUint8) {
+ + ]
608 : 4 : target_size = 1;
609 [ + + + + : 30 : } else if (dtype == kInt16 || dtype == kUint16 || dtype == kFloat16 || dtype == kBfloat16) {
+ + + + ]
610 : 5 : target_size = 2;
611 [ + + + + : 25 : } else if (dtype == kInt32 || dtype == kUint32 || dtype == kFloat32) {
+ + ]
612 : 16 : target_size = 4;
613 [ + + + + : 9 : } else if (dtype == kInt64 || dtype == kUint64 || dtype == kFloat64) {
+ + ]
614 : 6 : target_size = 8;
615 : : }
616 : :
617 : 34 : return target_size;
618 : : }
619 : :
620 : : } // namespace zerocopy
621 : :
622 : : } // namespace vlink
|
