mirror of
https://github.com/Tencent/ncnn.git
synced 2026-10-18 21:54:39 +08:00
201145bdd9
* fix crash reading past end of memory-mapped model file DataReaderFromMemory used by mapped model loading was unbounded: read()/reference() advanced past the mapped region, so a truncated or corrupt .bin crashed with an access violation inside ModelBinFromDataReader instead of failing cleanly. Add a size-aware DataReaderFromMemory ctor that clamps reads and fails references past the end, use it in the mapped branches of Net::load_model, and propagate load_model()'s return value which was previously dropped. * fix bounded memory reader scans and add loading regressions * f * reuse bounded memory reader scan buffer --------- Co-authored-by: Dr3ladwyn <Dr3ladwyn@users.noreply.github.com> Co-authored-by: nihui <nihuini@tencent.com>
230 lines
6.2 KiB
C++
230 lines
6.2 KiB
C++
// Copyright 2026 Tencent
|
|
// SPDX-License-Identifier: BSD-3-Clause
|
|
|
|
#include <stdio.h>
|
|
|
|
#include "datareader.h"
|
|
#include "modelbin.h"
|
|
|
|
static int test_mat_from_bfloat16()
|
|
{
|
|
const unsigned short data[] = {0x3f80, 0xc020, 0x3f00};
|
|
ncnn::Mat m = ncnn::Mat::from_bfloat16(data, 3);
|
|
|
|
if (m.empty() || m.elemsize != 4u || m[0] != 1.f || m[1] != -2.5f || m[2] != 0.5f)
|
|
{
|
|
fprintf(stderr, "test_mat_from_bfloat16 failed\n");
|
|
return -1;
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
static int test_modelbin_bfloat16()
|
|
{
|
|
const unsigned char model_data[] = {
|
|
0x83, 0x8b, 0x34, 0x01,
|
|
0x80, 0x3f, 0x20, 0xc0, 0x00, 0x3f, 0x00, 0x00,
|
|
0x00, 0x00, 0x00, 0x00,
|
|
0x00, 0x00, 0x80, 0x40
|
|
};
|
|
|
|
const unsigned char* mem = model_data;
|
|
ncnn::DataReaderFromMemory dr(mem);
|
|
ncnn::ModelBinFromDataReader mb(dr);
|
|
|
|
ncnn::Mat m = mb.load(3, 0);
|
|
if (m.empty() || m.elemsize != 4u || m[0] != 1.f || m[1] != -2.5f || m[2] != 0.5f)
|
|
{
|
|
fprintf(stderr, "test_modelbin_bfloat16 failed\n");
|
|
return -1;
|
|
}
|
|
|
|
ncnn::Mat m2 = mb.load(1, 0);
|
|
if (m2.empty() || m2[0] != 4.f)
|
|
{
|
|
fprintf(stderr, "test_modelbin_bfloat16 alignment failed\n");
|
|
return -1;
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
class DataReaderFromMemoryNoReference : public ncnn::DataReader
|
|
{
|
|
public:
|
|
DataReaderFromMemoryNoReference(const unsigned char* _data, size_t _size)
|
|
: data(_data), size(_size), offset(0)
|
|
{
|
|
}
|
|
|
|
virtual size_t read(void* buf, size_t read_size) const
|
|
{
|
|
if (offset >= size)
|
|
return 0;
|
|
|
|
size_t remain = size - offset;
|
|
if (read_size > remain)
|
|
read_size = remain;
|
|
|
|
memcpy(buf, data + offset, read_size);
|
|
offset += read_size;
|
|
return read_size;
|
|
}
|
|
|
|
private:
|
|
const unsigned char* data;
|
|
size_t size;
|
|
mutable size_t offset;
|
|
};
|
|
|
|
static int test_modelbin_int8(const ncnn::DataReader& dr, int w)
|
|
{
|
|
ncnn::ModelBinFromDataReader mb(dr);
|
|
|
|
ncnn::Mat m = mb.load(w, 3);
|
|
if (m.empty() || m.w != w || m.elemsize != 1u || m.elempack != 1)
|
|
{
|
|
fprintf(stderr, "test_modelbin_int8 shape failed w=%d\n", w);
|
|
return -1;
|
|
}
|
|
|
|
const signed char expected[] = {0, 1, -128, 127};
|
|
const signed char* p = m;
|
|
for (int i = 0; i < w; i++)
|
|
{
|
|
if (p[i] != expected[i])
|
|
{
|
|
fprintf(stderr, "test_modelbin_int8 value failed w=%d i=%d %d != %d\n", w, i, p[i], expected[i]);
|
|
return -1;
|
|
}
|
|
}
|
|
|
|
ncnn::Mat m2 = mb.load(1, 1);
|
|
if (m2.empty() || m2[0] != 4.f)
|
|
{
|
|
fprintf(stderr, "test_modelbin_int8 alignment failed w=%d\n", w);
|
|
return -1;
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
static int test_modelbin_int8(int w)
|
|
{
|
|
unsigned char model_data[] = {
|
|
0x00, 0x01, 0x80, 0x7f,
|
|
0x00, 0x00, 0x80, 0x40
|
|
};
|
|
for (int i = w; i < 4; i++)
|
|
model_data[i] = 0;
|
|
|
|
// reference path
|
|
const unsigned char* mem = model_data;
|
|
ncnn::DataReaderFromMemory dr(mem);
|
|
int ret = test_modelbin_int8(dr, w);
|
|
if (ret != 0)
|
|
return ret;
|
|
|
|
// copy path
|
|
// DataReaderFromMemoryNoReference inherits DataReader::reference(),
|
|
// which returns 0, so ModelBinFromDataReader falls back to read().
|
|
DataReaderFromMemoryNoReference dr2(model_data, sizeof(model_data));
|
|
ret = test_modelbin_int8(dr2, w);
|
|
if (ret != 0)
|
|
return ret;
|
|
|
|
return 0;
|
|
}
|
|
|
|
static int test_modelbin_int8_short_read()
|
|
{
|
|
// three int8 elements require four bytes in the model because raw
|
|
// int8 data is padded to a 32-bit boundary
|
|
const unsigned char model_data[] = {
|
|
0x00, 0x01, 0x00
|
|
};
|
|
|
|
DataReaderFromMemoryNoReference dr(model_data, sizeof(model_data));
|
|
ncnn::ModelBinFromDataReader mb(dr);
|
|
|
|
ncnn::Mat m = mb.load(3, 3);
|
|
if (!m.empty())
|
|
{
|
|
fprintf(stderr, "test_modelbin_int8_short_read failed\n");
|
|
return -1;
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
static int test_modelbin_int8()
|
|
{
|
|
return 0
|
|
|| test_modelbin_int8(1)
|
|
|| test_modelbin_int8(2)
|
|
|| test_modelbin_int8(3)
|
|
|| test_modelbin_int8(4)
|
|
|| test_modelbin_int8_short_read();
|
|
}
|
|
|
|
static int test_modelbin_bounded()
|
|
{
|
|
const unsigned char model_data[][16] = {
|
|
{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x80, 0x3f, 0x00, 0x00, 0x20, 0xc0, 0x00, 0x00, 0x00, 0x3f},
|
|
{0x47, 0x6b, 0x30, 0x01, 0x00, 0x3c, 0x00, 0xc1, 0x00, 0x38, 0x00, 0x00},
|
|
{0x83, 0x8b, 0x34, 0x01, 0x80, 0x3f, 0x20, 0xc0, 0x00, 0x3f, 0x00, 0x00},
|
|
{0x38, 0x4b, 0x0d, 0x00, 0x00, 0x01, 0x80, 0x00},
|
|
{0x00, 0x01, 0x80, 0x00},
|
|
{0x00, 0x00, 0x80, 0x3f, 0x00, 0x00, 0x20, 0xc0, 0x00, 0x00, 0x00, 0x3f}
|
|
};
|
|
const size_t sizes[] = {16, 12, 12, 8, 4, 12};
|
|
const int types[] = {0, 0, 0, 0, 3, 1};
|
|
|
|
for (size_t i = 0; i < sizeof(sizes) / sizeof(sizes[0]); i++)
|
|
{
|
|
// truncate at every byte, including format tags and alignment padding
|
|
for (size_t size = 0; size <= sizes[i]; size++)
|
|
{
|
|
const unsigned char* mem = model_data[i];
|
|
ncnn::DataReaderFromMemory dr(mem, size);
|
|
ncnn::ModelBinFromDataReader mb(dr);
|
|
ncnn::Mat m = mb.load(3, types[i]);
|
|
if (m.empty() != (size < sizes[i]) || mem != model_data[i] + size)
|
|
{
|
|
fprintf(stderr, "test_modelbin_bounded failed format=%zu size=%zu\n", i, size);
|
|
return -1;
|
|
}
|
|
|
|
if (size == sizes[i])
|
|
{
|
|
if (i == 3 || i == 4)
|
|
{
|
|
const signed char* p = m;
|
|
if (m.elemsize != 1u || p[0] != 0 || p[1] != 1 || p[2] != -128)
|
|
{
|
|
fprintf(stderr, "test_modelbin_bounded int8 values failed\n");
|
|
return -1;
|
|
}
|
|
}
|
|
else if (m.elemsize != 4u || m[0] != 1.f || m[1] != -2.5f || m[2] != 0.5f)
|
|
{
|
|
fprintf(stderr, "test_modelbin_bounded float values failed format=%zu\n", i);
|
|
return -1;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return 0;
|
|
}
|
|
|
|
int main()
|
|
{
|
|
return 0
|
|
|| test_mat_from_bfloat16()
|
|
|| test_modelbin_bfloat16()
|
|
|| test_modelbin_int8()
|
|
|| test_modelbin_bounded();
|
|
}
|