happyz
/
gemma.cpp
mirror of https://github.com/google/gemma.cpp.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix compress-inl bf16->f32 overrun 

Caught by Arm hwasan but not x86 asan.

PiperOrigin-RevId: 677779421
This commit is contained in:
Jan Wassenberg 2024-09-23 07:09:41 -07:00 committed by Copybara-Service
parent 35fdf848c7
commit cdbfebb10f
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
compression/compress-inl.h
View File

@ -245,7 +245,7 @@ struct CompressTraits<BF16> {
size_t i = 0; size_t i = 0;
if (num >= N16) { if (num >= N16) {
for (i = 0; i <= num - N16; i += N16) { for (; i <= num - N16; i += N16) {
const VBF packed0 = hn::LoadU(dbf, packed.ptr + packed_ofs + i); const VBF packed0 = hn::LoadU(dbf, packed.ptr + packed_ofs + i);
hn::StoreU(packed0, dbf, raw + i); hn::StoreU(packed0, dbf, raw + i);
} }
@ -271,7 +271,7 @@ struct CompressTraits<BF16> {
size_t i = 0; size_t i = 0;
if (num >= 2 * NF) { if (num >= 2 * NF) {
for (i = 0; i <= num - 2 * NF; i += 2 * NF) { for (; i <= num - 2 * NF; i += 2 * NF) {
VF raw0, raw1; VF raw0, raw1;
Load2(df, packed, packed_ofs + i, raw0, raw1); Load2(df, packed, packed_ofs + i, raw0, raw1);
hn::StoreU(raw0, df, raw + i); hn::StoreU(raw0, df, raw + i);
@ -289,7 +289,7 @@ struct CompressTraits<BF16> {
// If at most one vector, the first store adds zero padding. Check before // If at most one vector, the first store adds zero padding. Check before
// storing the second, because callers only pad to one vector. // storing the second, because callers only pad to one vector.
hn::StoreU(raw0, df, raw + i); hn::StoreU(raw0, df, raw + i);
if (remaining >= NF) hn::StoreU(raw1, df, raw + i + NF); if (remaining > NF) hn::StoreU(raw1, df, raw + i + NF);
} }
} }
}; };