File size: 4,622 Bytes
9422383 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | #include "model_store.h"
#include <cstring>
namespace {
template <typename T>
bool ReadScalar(const std::vector<uint8_t>& bytes, size_t offset, T* value) {
if (offset > bytes.size() || sizeof(T) > bytes.size() - offset) return false;
std::memcpy(value, bytes.data() + offset, sizeof(T));
return true;
}
bool BytesAreZero(const std::vector<uint8_t>& bytes, size_t offset, size_t size) {
if (offset > bytes.size() || size > bytes.size() - offset) return false;
for (size_t index = offset; index < offset + size; ++index) {
if (bytes[index] != 0) return false;
}
return true;
}
uint64_t Align(uint64_t value) {
const uint64_t alignment = yolo26::kModelAlignment;
return ((value + alignment - 1) / alignment) * alignment;
}
struct ModelHeader {
uint32_t version;
uint32_t resolution;
uint32_t class_count;
uint32_t slot_count;
uint64_t payload_offset;
};
bool ReadHeader(const std::vector<uint8_t>& bytes, ModelHeader* header) {
// model.bin is little-endian, as is the sole Linux AArch64 target. Scalar
// reads use memcpy to avoid alignment and aliasing assumptions.
return bytes.size() >= yolo26::kModelHeaderBytes &&
std::memcmp(bytes.data(), yolo26::kModelMagic.data(), yolo26::kModelMagic.size()) == 0 &&
ReadScalar(bytes, 8, &header->version) && ReadScalar(bytes, 12, &header->resolution) &&
ReadScalar(bytes, 16, &header->class_count) &&
ReadScalar(bytes, 20, &header->slot_count) &&
ReadScalar(bytes, 24, &header->payload_offset) &&
BytesAreZero(bytes, 32, yolo26::kModelHeaderBytes - 32);
}
bool HeaderMatchesRuntime(const std::vector<uint8_t>& bytes, const ModelHeader& header) {
const uint64_t expected_payload_offset =
Align(yolo26::kModelDirectoryOffset +
uint64_t{yolo26::kModelEntryBytes} * yolo26::kRuntimeWeightSlotCount);
return header.version == yolo26::kModelFormatVersion &&
header.slot_count == yolo26::kRuntimeWeightSlotCount &&
header.payload_offset == expected_payload_offset &&
header.payload_offset <= bytes.size() &&
yolo26::IsSupportedResolution(header.resolution) &&
yolo26::ClassDependentWeightBytesFit(header.class_count);
}
} // namespace
void ModelStore::Reset() {
bytes_.clear();
slots_ = {};
resolution_ = 0;
class_count_ = 0;
}
bool ModelStore::Load(const void* data, size_t size) {
// Replace the complete store rather than retaining any state from a previous
// model. On success, bytes_ owns the file data and slots_ indexes each weight
// payload inside it. Every failure below resets both so callers cannot use
// stale metadata or pointers.
Reset();
if (!data || size < yolo26::kModelHeaderBytes) return false;
// Keep one stable copy so Slot() can return pointers into the model without
// retaining the caller's buffer or copying each weight separately.
const auto* first = static_cast<const uint8_t*>(data);
bytes_.assign(first, first + size);
// Reject files whose header does not match the one model format and set of
// resolutions supported by this runtime.
ModelHeader header{};
if (!ReadHeader(bytes_, &header) || !HeaderMatchesRuntime(bytes_, header)) {
Reset();
return false;
}
resolution_ = header.resolution;
class_count_ = header.class_count;
// The directory contains one offset and byte count for every weight expected
// by the fixed graph. Entries must be nonempty, aligned, ordered without
// overlap, and wholly contained in the file before their locations are kept.
uint64_t previous_end = header.payload_offset;
for (uint32_t index = 0; index < yolo26::kRuntimeWeightSlotCount; ++index) {
const size_t entry = yolo26::kModelDirectoryOffset + size_t{index} * yolo26::kModelEntryBytes;
uint64_t offset = 0;
uint64_t size = 0;
if (!ReadScalar(bytes_, entry, &offset) || !ReadScalar(bytes_, entry + 8, &size) || size == 0 ||
offset % yolo26::kModelAlignment != 0 || offset < previous_end || offset > bytes_.size() ||
size > bytes_.size() - offset) {
Reset();
return false;
}
slots_[index] = {offset, size};
previous_end = offset + size;
}
// The final payload must end with the file; otherwise unaccounted trailing
// data would sit outside the validated slot directory.
if (previous_end != bytes_.size()) {
Reset();
return false;
}
return true;
}
const void* ModelStore::Slot(uint32_t slot, size_t expected_size) const {
if (slot >= slots_.size() || slots_[slot].size != expected_size) return nullptr;
return bytes_.data() + slots_[slot].offset;
}
|