Skip to content

Commit

Permalink
Use negative offsets for read_metadata_from(stream) function, drop …
Browse files Browse the repository at this point in the history 
…its overflow with `ov::AlignedBuffer`
  • Loading branch information
@MirceaDan99
MirceaDan99 committed Dec 18, 2024
1 parent f003969 commit 155df1b
Show file tree
Hide file tree
Showing 5 changed files with 28 additions and 87 deletions.
13 changes: 2 additions & 11 deletions src/plugins/intel_npu/src/plugin/include/metadata.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -102,8 +102,6 @@ struct MetadataBase {

virtual uint64_t get_blob_size() const = 0;

virtual size_t get_ov_header_offset() const = 0;

virtual ~MetadataBase() = default;
};

Expand All @@ -123,13 +121,10 @@ struct Metadata<METADATA_VERSION_1_0> : public MetadataBase {
private:
uint64_t _blobDataSize;
uint32_t _version;
size_t _ovHeaderOffset;
OpenvinoVersion _ovVersion;

public:
Metadata();

Metadata(size_t ovHeaderOffset, uint64_t blobDataSize);
Metadata(uint64_t blobDataSize);

void read(std::istream& stream) override;

Expand All @@ -154,8 +149,6 @@ struct Metadata<METADATA_VERSION_1_0> : public MetadataBase {
void set_ov_version(const OpenvinoVersion& newVersion);

uint64_t get_blob_size() const override;

size_t get_ov_header_offset() const override;
};

/**
Expand All @@ -164,7 +157,7 @@ struct Metadata<METADATA_VERSION_1_0> : public MetadataBase {
* @return Unique pointer to the created MetadataBase object if the major version is supported; otherwise, returns
* 'nullptr'.
*/
std::unique_ptr<MetadataBase> create_metadata(uint32_t version, size_t ovHeaderOffset, uint64_t blobDataSize);
std::unique_ptr<MetadataBase> create_metadata(uint32_t version, uint64_t blobDataSize);

/**
* @brief Reads metadata from a blob.
Expand All @@ -174,6 +167,4 @@ std::unique_ptr<MetadataBase> create_metadata(uint32_t version, size_t ovHeaderO
*/
std::unique_ptr<MetadataBase> read_metadata_from(std::istream& stream);

std::unique_ptr<MetadataBase> read_metadata_from(std::istream& stream, const std::shared_ptr<ov::AlignedBuffer>& modelBuffer);

} // namespace intel_npu
4 changes: 1 addition & 3 deletions src/plugins/intel_npu/src/plugin/src/compiled_model.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -75,10 +75,8 @@ void CompiledModel::export_model(std::ostream& stream) const {
_logger.debug("CompiledModel::export_model");
size_t blobSizeBeforeVersioning = _graph->export_blob(stream);

auto meta = Metadata<CURRENT_METADATA_VERSION>();
auto meta = Metadata<CURRENT_METADATA_VERSION>(blobSizeBeforeVersioning);
meta.write(stream);
stream.write(reinterpret_cast<const char*>(&blobSizeBeforeVersioning), sizeof(blobSizeBeforeVersioning));
stream.write(MAGIC_BYTES.data(), MAGIC_BYTES.size());
}

std::shared_ptr<const ov::Model> CompiledModel::get_runtime_model() const {
Expand Down
82 changes: 20 additions & 62 deletions src/plugins/intel_npu/src/plugin/src/metadata.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,15 +14,20 @@

namespace {

size_t getFileSize(std::istream& stream) {
std::streampos getFileSize(std::istream& stream) {
auto log = intel_npu::Logger::global().clone("getFileSize");
if (!stream) {
OPENVINO_THROW("Stream is in bad status! Please check the passed stream status!");
}

const size_t streamStart = stream.tellg();
if (stream.rdbuf()->in_avail() != 0) {
// ov::OwningSharedStreamBuffer scenario
return stream.rdbuf()->in_avail() + stream.tellg();
}

const std::streampos streamStart = stream.tellg();
stream.seekg(0, std::ios_base::end);
const size_t streamEnd = stream.tellg();
const std::streampos streamEnd = stream.tellg();
stream.seekg(streamStart, std::ios_base::beg);

log.debug("Read blob size: streamStart=%zu, streamEnd=%zu", streamStart, streamEnd);
Expand Down Expand Up @@ -52,16 +57,9 @@ void OpenvinoVersion::read(std::istream& stream) {
stream.read(_version.data(), _size);
}

Metadata<METADATA_VERSION_1_0>::Metadata()
: _blobDataSize{0},
_version{METADATA_VERSION_1_0},
_ovHeaderOffset{0},
_ovVersion{ov::get_openvino_version().buildNumber} {}

Metadata<METADATA_VERSION_1_0>::Metadata(size_t ovHeaderOffset, uint64_t blobDataSize)
Metadata<METADATA_VERSION_1_0>::Metadata(uint64_t blobDataSize)
: _blobDataSize{blobDataSize},
_version{METADATA_VERSION_1_0},
_ovHeaderOffset{ovHeaderOffset},
_ovVersion{ov::get_openvino_version().buildNumber} {}

void Metadata<METADATA_VERSION_1_0>::read(std::istream& stream) {
Expand All @@ -76,12 +74,14 @@ void Metadata<METADATA_VERSION_1_0>::write(std::ostream& stream) {
stream.write(reinterpret_cast<const char*>(&_version), sizeof(_version));
stream.write(reinterpret_cast<const char*>(&_ovVersion.get_size()), sizeof(_ovVersion.get_size()));
stream.write(_ovVersion.get_version().data(), _ovVersion.get_version().size());
stream.write(reinterpret_cast<const char*>(&_blobDataSize), sizeof(_blobDataSize));
stream.write(MAGIC_BYTES.data(), MAGIC_BYTES.size());
}

std::unique_ptr<MetadataBase> create_metadata(uint32_t version, size_t ovHeaderOffset, uint64_t blobDataSize) {
std::unique_ptr<MetadataBase> create_metadata(uint32_t version, uint64_t blobDataSize) {
switch (version) {
case METADATA_VERSION_1_0:
return std::make_unique<Metadata<METADATA_VERSION_1_0>>(ovHeaderOffset, blobDataSize);
return std::make_unique<Metadata<METADATA_VERSION_1_0>>(blobDataSize);

default:
return nullptr;
Expand Down Expand Up @@ -118,63 +118,24 @@ std::unique_ptr<MetadataBase> read_metadata_from(std::istream& stream) {
std::string blobMagicBytes;
blobMagicBytes.resize(magicBytesSize);

size_t currentStreamPos = stream.tellg();
size_t streamSize = getFileSize(stream);
stream.seekg(streamSize - magicBytesSize, std::ios::beg);
std::streampos currentStreamPos = stream.tellg();
std::streampos streamSize = getFileSize(stream);
stream.seekg(-currentStreamPos + streamSize - magicBytesSize, std::ios::cur);
stream.read(blobMagicBytes.data(), magicBytesSize);
if (MAGIC_BYTES != blobMagicBytes) {
logger.error("Blob is missing NPU metadata!");
return nullptr;
}

uint64_t blobDataSize;
stream.seekg(streamSize - magicBytesSize - sizeof(blobDataSize), std::ios::beg);
stream.seekg(-std::streampos(magicBytesSize) - sizeof(blobDataSize), std::ios::cur);
stream.read(reinterpret_cast<char*>(&blobDataSize), sizeof(blobDataSize));
stream.seekg(currentStreamPos + blobDataSize, std::ios::beg);
stream.seekg(-stream.tellg() + currentStreamPos + blobDataSize, std::ios::cur);

uint32_t metaVersion;
stream.read(reinterpret_cast<char*>(&metaVersion), sizeof(metaVersion));

auto storedMeta = create_metadata(metaVersion, currentStreamPos, blobDataSize);
if (storedMeta != nullptr) {
storedMeta->read(stream);
} else {
logger.warning("Imported blob metadata version: %d.%d, but the current version is: %d.%d",
get_major(metaVersion),
get_minor(metaVersion),
get_major(CURRENT_METADATA_VERSION),
get_minor(CURRENT_METADATA_VERSION));
}
stream.seekg(currentStreamPos, std::ios::beg);
return storedMeta;
}

std::unique_ptr<MetadataBase> read_metadata_from(std::istream& stream, const std::shared_ptr<ov::AlignedBuffer>& modelBuffer) {
Logger logger("NPUPlugin", Logger::global().level());
if (modelBuffer == nullptr) {
return read_metadata_from(stream);
}
size_t magicBytesSize = MAGIC_BYTES.size();
std::string blobMagicBytes;
blobMagicBytes.resize(magicBytesSize);

size_t currentStreamPos = stream.tellg();
size_t streamSize = modelBuffer->size();

blobMagicBytes.assign(reinterpret_cast<const char*>(modelBuffer->get_ptr(streamSize - magicBytesSize)), magicBytesSize);
if (MAGIC_BYTES != blobMagicBytes) {
logger.error("Blob is missing NPU metadata!");
return nullptr;
}

uint64_t blobDataSize;
blobDataSize = *reinterpret_cast<uint64_t*>(modelBuffer->get_ptr(streamSize - magicBytesSize - sizeof(blobDataSize)));

uint32_t metaVersion;
metaVersion = *reinterpret_cast<uint32_t*>(modelBuffer->get_ptr(currentStreamPos + blobDataSize));

auto storedMeta = create_metadata(metaVersion, currentStreamPos, blobDataSize);
stream.seekg(blobDataSize + sizeof(metaVersion), std::ios::cur);
auto storedMeta = create_metadata(metaVersion, blobDataSize);
if (storedMeta != nullptr) {
storedMeta->read(stream);
} else {
Expand All @@ -184,6 +145,7 @@ std::unique_ptr<MetadataBase> read_metadata_from(std::istream& stream, const std
get_major(CURRENT_METADATA_VERSION),
get_minor(CURRENT_METADATA_VERSION));
}
stream.seekg(-stream.tellg() + currentStreamPos, std::ios::cur);
return storedMeta;
}

Expand All @@ -199,8 +161,4 @@ uint64_t Metadata<METADATA_VERSION_1_0>::get_blob_size() const {
return _blobDataSize;
}

size_t Metadata<METADATA_VERSION_1_0>::get_ov_header_offset() const {
return _ovHeaderOffset;
}

} // namespace intel_npu
4 changes: 2 additions & 2 deletions src/plugins/intel_npu/src/plugin/src/plugin.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -760,7 +760,7 @@ std::shared_ptr<ov::ICompiledModel> Plugin::import_model(std::istream& stream, c
auto compiler = compilerAdapterFactory.getCompiler(_backends->getIEngineBackend(), localConfig);

std::unique_ptr<BlobContainer> blobPtr;
auto storedMeta = read_metadata_from(stream, modelBuffer);
auto storedMeta = read_metadata_from(stream);

if (storedMeta == nullptr) {
OPENVINO_THROW("Could not read metadata!");
Expand All @@ -779,7 +779,7 @@ std::shared_ptr<ov::ICompiledModel> Plugin::import_model(std::istream& stream, c

blobPtr = std::make_unique<BlobContainerVector>(std::move(blob));
} else {
blobPtr = std::make_unique<BlobContainerAlignedBuffer>(modelBuffer, storedMeta->get_ov_header_offset(), graphSize);
blobPtr = std::make_unique<BlobContainerAlignedBuffer>(modelBuffer, stream.tellg(), graphSize);
}

auto graph = compiler->parse(std::move(blobPtr), localConfig);
Expand Down
12 changes: 3 additions & 9 deletions src/plugins/intel_npu/tests/unit/npu/metadata_version.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,11 +21,9 @@ TEST_F(MetadataUnitTests, readUnversionedBlob) {
TEST_F(MetadataUnitTests, writeAndReadMetadataFromBlob) {
std::stringstream stream;
size_t blobSize = 0;
auto meta = Metadata<CURRENT_METADATA_VERSION>(stream.tellg(), blobSize);
auto meta = Metadata<CURRENT_METADATA_VERSION>(blobSize);

OV_ASSERT_NO_THROW(meta.write(stream));
OV_ASSERT_NO_THROW(stream.write(reinterpret_cast<const char*>(&blobSize), sizeof(blobSize)));
OV_ASSERT_NO_THROW(stream.write(MAGIC_BYTES.data(), MAGIC_BYTES.size()));

auto storedMeta = read_metadata_from(stream);
ASSERT_NE(storedMeta, nullptr);
Expand All @@ -35,14 +33,12 @@ TEST_F(MetadataUnitTests, writeAndReadMetadataFromBlob) {
TEST_F(MetadataUnitTests, writeAndReadInvalidOpenvinoVersion) {
size_t blobSize = 0;
std::stringstream stream;
auto meta = Metadata<CURRENT_METADATA_VERSION>(stream.tellg(), blobSize);
auto meta = Metadata<CURRENT_METADATA_VERSION>(blobSize);

OpenvinoVersion badOvVersion("just_some_wrong_ov_version");
meta.set_ov_version(badOvVersion);

OV_ASSERT_NO_THROW(meta.write(stream));
OV_ASSERT_NO_THROW(stream.write(reinterpret_cast<const char*>(&blobSize), sizeof(blobSize)));
OV_ASSERT_NO_THROW(stream.write(MAGIC_BYTES.data(), MAGIC_BYTES.size()));

auto storedMeta = read_metadata_from(stream);
ASSERT_NE(storedMeta, nullptr);
Expand All @@ -52,14 +48,12 @@ TEST_F(MetadataUnitTests, writeAndReadInvalidOpenvinoVersion) {
TEST_F(MetadataUnitTests, writeAndReadInvalidMetadataVersion) {
size_t blobSize = 0;
std::stringstream stream;
auto meta = Metadata<CURRENT_METADATA_VERSION>(stream.tellg(), blobSize);
auto meta = Metadata<CURRENT_METADATA_VERSION>(blobSize);

constexpr uint32_t dummy_version = make_version(0x00007E57, 0x0000AC3D);
meta.set_version(dummy_version);

OV_ASSERT_NO_THROW(meta.write(stream));
OV_ASSERT_NO_THROW(stream.write(reinterpret_cast<const char*>(&blobSize), sizeof(blobSize)));
OV_ASSERT_NO_THROW(stream.write(MAGIC_BYTES.data(), MAGIC_BYTES.size()));

auto storedMeta = read_metadata_from(stream);
ASSERT_EQ(storedMeta, nullptr);
Expand Down

0 comments on commit 155df1b

Please sign in to comment.