Fix SFP/NUQ for bf16 rounding in Highway

SFP: Avoid rounding twice, and more robust TestDot.
NUQ: also more robust SNR, minor touchups to header.

PiperOrigin-RevId: 618030096
This commit is contained in:
Jan Wassenberg 2024-03-21 19:05:44 -07:00 committed by Copybara-Service
parent a135bc1e47
commit 24add61dd9
7 changed files with 230 additions and 85 deletions

View File

@ -40,6 +40,7 @@ set(SOURCES
compression/nuq-inl.h compression/nuq-inl.h
compression/sfp.h compression/sfp.h
compression/sfp-inl.h compression/sfp-inl.h
compression/test_util.h
util/app.h util/app.h
util/args.h util/args.h
) )

View File

@ -24,20 +24,36 @@ cc_library(
], ],
) )
# Deprecated because it is also implemented in Highway; will be removed once
# that Highway version is sufficiently widespread.
cc_library( cc_library(
name = "stats", name = "stats",
srcs = [ srcs = ["stats.cc"],
"stats.cc", hdrs = ["stats.h"],
],
hdrs = [
"distortion.h",
"stats.h",
],
deps = [ deps = [
"@hwy//:hwy", "@hwy//:hwy",
], ],
) )
cc_library(
name = "distortion",
hdrs = ["distortion.h"],
deps = [
"@hwy//:hwy",
],
)
cc_library(
name = "test_util",
hdrs = ["test_util.h"],
deps = [
":distortion",
":stats",
"@hwy//:hwy",
"@hwy//:hwy_test_util",
],
)
cc_library( cc_library(
name = "sfp", name = "sfp",
hdrs = [ hdrs = [
@ -62,12 +78,11 @@ cc_test(
tags = ["hwy_ops_test"], tags = ["hwy_ops_test"],
deps = [ deps = [
":sfp", ":sfp",
":stats", ":test_util",
"@googletest//:gtest_main", "@googletest//:gtest_main",
"@hwy//:hwy", "@hwy//:hwy",
"@hwy//:hwy_test_util", "@hwy//:hwy_test_util",
"@hwy//:nanobenchmark", "@hwy//:nanobenchmark",
"@hwy//:thread_pool",
], ],
) )
@ -98,7 +113,7 @@ cc_test(
deps = [ deps = [
":nuq", ":nuq",
":sfp", ":sfp",
":stats", ":test_util",
"@googletest//:gtest_main", "@googletest//:gtest_main",
"@hwy//:hwy", "@hwy//:hwy",
"@hwy//:hwy_test_util", "@hwy//:hwy_test_util",
@ -118,6 +133,7 @@ cc_library(
], ],
deps = [ deps = [
":blob_store", ":blob_store",
":distortion",
":nuq", ":nuq",
":sfp", ":sfp",
":stats", ":stats",
@ -134,6 +150,7 @@ cc_library(
"analyze.h", "analyze.h",
], ],
deps = [ deps = [
":distortion",
":nuq", ":nuq",
":sfp", ":sfp",
":stats", ":stats",

View File

@ -42,6 +42,10 @@
#include "hwy/contrib/sort/vqsort-inl.h" #include "hwy/contrib/sort/vqsort-inl.h"
#include "hwy/highway.h" #include "hwy/highway.h"
#ifndef HWY_IF_CONSTEXPR
#define HWY_IF_CONSTEXPR if
#endif
HWY_BEFORE_NAMESPACE(); HWY_BEFORE_NAMESPACE();
namespace gcpp { namespace gcpp {
namespace HWY_NAMESPACE { namespace HWY_NAMESPACE {
@ -124,7 +128,7 @@ class NuqClustering {
} }
private: private:
// Float has enough precision for our relatively small kGroupSize (128). // Float has enough precision for our relatively small kGroupSize (256).
float cumsum_[kGroupSize + 1]; float cumsum_[kGroupSize + 1];
float cumsum2_[kGroupSize + 1]; float cumsum2_[kGroupSize + 1];
float inv_len_[kGroupSize + 1]; float inv_len_[kGroupSize + 1];
@ -168,8 +172,8 @@ class NuqClustering {
// `centers`; prior centers are zero-initialized. // `centers`; prior centers are zero-initialized.
// //
// O(kClusters * kGroupSize * kGroupSize), but the constant factors are so low // O(kClusters * kGroupSize * kGroupSize), but the constant factors are so low
// that this is about 10 times as fast as the O(kClusters * kGroupSize) SMAWK // that this is about 5 times as fast as the O(kClusters * kGroupSize) SMAWK
// as implemented in FAISS, for our kGroupSize <= 128. // as implemented in FAISS, for our kGroupSize of 256.
template <class DF> template <class DF>
static HWY_NOINLINE size_t ClusterExactL2(DF df, const float* x, static HWY_NOINLINE size_t ClusterExactL2(DF df, const float* x,
ClusterBuf& buf, ClusterBuf& buf,
@ -228,7 +232,7 @@ class NuqClustering {
// Center = mean, O(1) thanks to cumulative sums. // Center = mean, O(1) thanks to cumulative sums.
const float sum = cc.SumOfSorted(start, last); const float sum = cc.SumOfSorted(start, last);
const int size = static_cast<int>(last) - static_cast<int>(start) + 1; const int size = static_cast<int>(last) - static_cast<int>(start) + 1;
HWY_DASSERT(0 < size && size <= kGroupSize); HWY_DASSERT(0 < size && size <= static_cast<int>(kGroupSize));
centers[k] = sum / static_cast<float>(size); centers[k] = sum / static_cast<float>(size);
// We know the range inside sorted_and_i[]; translate to original indices, // We know the range inside sorted_and_i[]; translate to original indices,
@ -427,7 +431,7 @@ class NuqCodec {
// instead of TableLookupBytes, which requires extra interleaving of lo/hi. // instead of TableLookupBytes, which requires extra interleaving of lo/hi.
HWY_DASSERT(hn::Lanes(du) >= 8); HWY_DASSERT(hn::Lanes(du) >= 8);
if (NumTables(du) == 2) { HWY_IF_CONSTEXPR(NumTables(du) == 2) {
// Reduce cap for second half to avoid loading past the end of the table. // Reduce cap for second half to avoid loading past the end of the table.
const hn::CappedTag<hwy::bfloat16_t, kClusters / 2> d_table2; const hn::CappedTag<hwy::bfloat16_t, kClusters / 2> d_table2;
*tbl1 = hn::ResizeBitCast(du, hn::LoadU(d_table2, table + kClusters / 2)); *tbl1 = hn::ResizeBitCast(du, hn::LoadU(d_table2, table + kClusters / 2));
@ -449,11 +453,12 @@ class NuqCodec {
const auto indices0 = hn::IndicesFromVec(du, idx0); const auto indices0 = hn::IndicesFromVec(du, idx0);
const auto indices1 = hn::IndicesFromVec(du, idx1); const auto indices1 = hn::IndicesFromVec(du, idx1);
if (NumTables(du) == 1) { HWY_IF_CONSTEXPR(NumTables(du) == 1) {
(void)tbl1; (void)tbl1;
c0 = hn::TableLookupLanes(tbl0, indices0); c0 = hn::TableLookupLanes(tbl0, indices0);
c1 = hn::TableLookupLanes(tbl0, indices1); c1 = hn::TableLookupLanes(tbl0, indices1);
} else { }
HWY_IF_CONSTEXPR(NumTables(du) == 2) { // `else` is poorly formatted.
c0 = hn::TwoTablesLookupLanes(du, tbl0, tbl1, indices0); c0 = hn::TwoTablesLookupLanes(du, tbl0, tbl1, indices0);
c1 = hn::TwoTablesLookupLanes(du, tbl0, tbl1, indices1); c1 = hn::TwoTablesLookupLanes(du, tbl0, tbl1, indices1);
} }
@ -521,8 +526,8 @@ class NuqCodec {
// Decodes `num` values from the stream `in`, starting at the offset `in_ofs` // Decodes `num` values from the stream `in`, starting at the offset `in_ofs`
// (in units of values), to bf16 in `out`. `in_capacity`, `in_ofs` and `num` // (in units of values), to bf16 in `out`. `in_capacity`, `in_ofs` and `num`
// must all be multiples of `kGroupSize`. // must all be multiples of `kGroupSize`.
template <class DF, HWY_IF_BF16_D(DF)> template <class DBF, HWY_IF_BF16_D(DBF)>
static HWY_INLINE void Dec(DF dbf, const size_t in_capacity, static HWY_INLINE void Dec(DBF dbf, const size_t in_capacity,
const NuqStream* const in, const size_t in_ofs, const NuqStream* const in, const size_t in_ofs,
hwy::bfloat16_t* const out, const size_t num) { hwy::bfloat16_t* const out, const size_t num) {
const hn::RebindToUnsigned<decltype(dbf)> d16; const hn::RebindToUnsigned<decltype(dbf)> d16;

View File

@ -27,6 +27,7 @@
#include "hwy/aligned_allocator.h" #include "hwy/aligned_allocator.h"
#include "hwy/base.h" #include "hwy/base.h"
#include "hwy/timer.h"
// clang-format off // clang-format off
#undef HWY_TARGET_INCLUDE #undef HWY_TARGET_INCLUDE
@ -35,15 +36,14 @@
#include "hwy/foreach_target.h" // IWYU pragma: keep #include "hwy/foreach_target.h" // IWYU pragma: keep
// Other headers that include Highway must come after foreach_target.h // Other headers that include Highway must come after foreach_target.h
// copybara:import_next_line:gemma_cpp // copybara:import_next_line:gemma_cpp
#include "compression/distortion.h"
// copybara:import_next_line:gemma_cpp
#include "compression/nuq-inl.h" #include "compression/nuq-inl.h"
// copybara:import_next_line:gemma_cpp // copybara:import_next_line:gemma_cpp
#include "compression/nuq.h" #include "compression/nuq.h"
// copybara:import_next_line:gemma_cpp
#include "compression/test_util.h"
#include "hwy/highway.h" #include "hwy/highway.h"
#include "hwy/tests/hwy_gtest.h" #include "hwy/tests/hwy_gtest.h"
#include "hwy/tests/test_util-inl.h" #include "hwy/tests/test_util-inl.h"
#include "hwy/timer.h"
HWY_BEFORE_NAMESPACE(); HWY_BEFORE_NAMESPACE();
namespace gcpp { namespace gcpp {
@ -181,12 +181,14 @@ struct TestNormal {
auto in = hwy::AllocateAligned<float>(kGroupSize); auto in = hwy::AllocateAligned<float>(kGroupSize);
HWY_ASSERT(in); HWY_ASSERT(in);
std::mt19937 rng(123); hwy::RandomState rng;
std::normal_distribution<float> dist{0.001f, 0.3f}; Stats in_stats;
for (size_t i = 0; i < kGroupSize; ++i) { for (size_t i = 0; i < kGroupSize; ++i) {
in[i] = dist(rng); const double r = RandomGaussian(rng);
in_stats.Notify(r);
in[i] = hwy::ConvertScalarTo<T>(r);
} }
std::shuffle(in.get(), in.get() + kGroupSize, rng); VerifyGaussian(in_stats);
ClusterBuf buf; ClusterBuf buf;
float centers[kClusters]; float centers[kClusters];
@ -212,9 +214,9 @@ struct TestNormal {
const float snr = stats.GeomeanValueDivL1(); const float snr = stats.GeomeanValueDivL1();
fprintf(stderr, "p-norm %.3E snr %.2f @%zu = %.4E\n", pnorm, snr, fprintf(stderr, "p-norm %.3E snr %.2f @%zu = %.4E\n", pnorm, snr,
stats.MaxIndex(), stats.MaxL1()); stats.MaxIndex(), stats.MaxL1());
static_assert(kGroupSize == 128 || kGroupSize == 256, "Update expected"); static_assert(kGroupSize == 256, "Update expected");
const float expected_pnorm = kGroupSize == 128 ? 3E-2f : 3.4E-2f; const float expected_pnorm = 3.68E-2f;
const float expected_snr = kGroupSize == 128 ? 17.4f : 13.1f; const float expected_snr = 12.7f;
HWY_ASSERT(expected_pnorm <= pnorm && pnorm < 1.02f * expected_pnorm); HWY_ASSERT(expected_pnorm <= pnorm && pnorm < 1.02f * expected_pnorm);
HWY_ASSERT(expected_snr <= snr && snr < 1.01f * expected_snr); HWY_ASSERT(expected_snr <= snr && snr < 1.01f * expected_snr);
} }
@ -345,21 +347,27 @@ struct TestDot {
auto nuq = hwy::AllocateAligned<NuqStream>(NuqStream::PackedEnd(num)); auto nuq = hwy::AllocateAligned<NuqStream>(NuqStream::PackedEnd(num));
HWY_ASSERT(in && dec && vec && nuq); HWY_ASSERT(in && dec && vec && nuq);
std::mt19937 rng(123); // Generate inputs and verify their distribution.
std::normal_distribution<float> dist{0.001f, 0.3f}; hwy::RandomState rng;
Stats in_stats;
for (size_t i = 0; i < num; ++i) { for (size_t i = 0; i < num; ++i) {
in[i] = dist(rng); const float r = static_cast<float>(RandomGaussian(rng));
vec[i] = hwy::ConvertScalarTo<T>(dist(rng)); in_stats.Notify(r);
in[i] = r;
} }
// This changes the correlation between in and vec, which considerably for (size_t i = 0; i < num; ++i) {
// affects the error of the result. const float r = static_cast<float>(RandomGaussian(rng));
std::shuffle(in.get(), in.get() + num, rng); in_stats.Notify(r);
vec[i] = hwy::ConvertScalarTo<T>(r);
}
VerifyGaussian(in_stats);
ClusterBuf buf; ClusterBuf buf;
const size_t unused_clusters = const size_t unused_clusters =
NuqCodec::Enc(df, in.get(), num, buf, num, nuq.get(), 0); NuqCodec::Enc(df, in.get(), num, buf, num, nuq.get(), 0);
HWY_ASSERT(unused_clusters == 0); HWY_ASSERT(unused_clusters == 0);
// Compute dot product without decompression.
double actual = 0.0; double actual = 0.0;
double elapsed = hwy::HighestValue<double>(); double elapsed = hwy::HighestValue<double>();
for (size_t rep = 0; rep < 20; ++rep) { for (size_t rep = 0; rep < 20; ++rep) {
@ -380,24 +388,39 @@ struct TestDot {
fprintf(stderr, "Vec %zu Dec %.2f MB/s\n", Lanes(d) * sizeof(T), fprintf(stderr, "Vec %zu Dec %.2f MB/s\n", Lanes(d) * sizeof(T),
num * sizeof(in[0]) * 1E-6 / elapsed); num * sizeof(in[0]) * 1E-6 / elapsed);
double expected = 0.0; // using original input // Exact and decompressed dot products for comparison.
double expected2 = 0.0; // using decoded NUQ double exact = 0.0; // using original input
double expected = 0.0; // using decoded NUQ
DistortionStats dec_stats;
Stats ratios;
for (size_t i = 0; i < num; ++i) { for (size_t i = 0; i < num; ++i) {
expected += in[i] * hwy::ConvertScalarTo<double>(vec[i]); dec_stats.Notify(in[i], dec[i]);
expected2 += dec[i] * hwy::ConvertScalarTo<double>(vec[i]); const float v1 = hwy::ConvertScalarTo<float>(vec[i]);
exact += in[i] * v1;
expected += dec[i] * v1;
if (expected != 0.0f) {
ratios.Notify(exact / expected);
} }
const double l1 = hwy::ScalarAbs(expected - actual); }
const double snr = 1.0 + hwy::ScalarAbs(expected) / l1; const double dec_snr = dec_stats.GeomeanValueDivL1();
fprintf(stderr, "expected %.3f e2 %.4f actual %.4f l1 %E snr %.2f\n", const double dot_snr = 1.0 / hwy::ScalarAbs(1.0 - ratios.GeometricMean());
expected, expected2, actual, l1, snr); // exact and actual fluctuate due to the combination of NUQ imprecision,
HWY_ASSERT(hwy::ScalarAbs(expected2 - actual) < 1E-4); // and whether vec[i] is negative or positive, so this is quite loose.
static_assert(kGroupSize == 128 || kGroupSize == 256, "Update expected"); const float final_ratio = HWY_MIN(exact / actual, actual / exact);
const double expected_l1 = kGroupSize == 128 ? 7.3E-2 : 4.34E-2; fprintf(stderr, "ratios %s\n", ratios.ToString().c_str());
const double expected_snr = kGroupSize == 128 ? 9.7f fprintf(stderr,
: sizeof(T) == 2 ? 14.5f "exact %.3f e2 %.4f actual %.4f final_ratio %.3f dec_snr %.2f "
: 14.9f; "dot_snr %.2f\n",
HWY_ASSERT(expected_l1 <= l1 && l1 < 1.02f * expected_l1); exact, expected, actual, final_ratio, dec_snr, dot_snr);
HWY_ASSERT(expected_snr <= snr && snr < 1.01f * expected_snr); // Final values are not too far apart.
HWY_ASSERT(0.88f <= final_ratio && final_ratio <= 1.0f);
// Decompressed and uncompressed dot should match exactly.
HWY_ASSERT(hwy::ScalarAbs(expected - actual) < 1E-4f);
// dec[] is close to in[], but we already check that in TestStream.
HWY_ASSERT(dec_snr >= 13.0);
// Geomean of ratios for each i is an approximation of the actual SNR.
HWY_ASSERT(dot_snr >= (sizeof(T) == 2 ? 17.0 : 14.0));
static_assert(kGroupSize == 256, "Update expected*");
} }
}; };

View File

@ -449,6 +449,18 @@ class SfpCodec {
return Enc2U(d16, w0, w1); return Enc2U(d16, w0, w1);
} }
// Truncates two f32 to bf16, in lane order, without rounding (see Enc4F).
template <class DBF, class DF = hn::RepartitionToWide<DBF>>
static HWY_INLINE hn::Vec<DBF> Truncate2To(DBF dbf, hn::Vec<DF> f0,
hn::Vec<DF> f1) {
const hn::RebindToUnsigned<DBF> d16;
using V16 = hn::Vec<decltype(d16)>;
const V16 u0 = BitCast(d16, f0);
const V16 u1 = BitCast(d16, f1);
return BitCast(DBF(), HWY_IS_LITTLE_ENDIAN ? ConcatOdd(d16, u1, u0)
: ConcatEven(d16, u1, u0));
}
template <class DF, HWY_IF_F32_D(DF), template <class DF, HWY_IF_F32_D(DF),
class V8 = hn::Vec<hn::Repartition<uint8_t, DF>>> class V8 = hn::Vec<hn::Repartition<uint8_t, DF>>>
static HWY_INLINE V8 Enc4F(DF df, const float* HWY_RESTRICT in) { static HWY_INLINE V8 Enc4F(DF df, const float* HWY_RESTRICT in) {
@ -462,9 +474,10 @@ class SfpCodec {
const VF f1 = hn::LoadU(df, in + NF * 1); const VF f1 = hn::LoadU(df, in + NF * 1);
const VF f2 = hn::LoadU(df, in + NF * 2); const VF f2 = hn::LoadU(df, in + NF * 2);
const VF f3 = hn::LoadU(df, in + NF * 3); const VF f3 = hn::LoadU(df, in + NF * 3);
// Chop off the lower 16 bits; EncBytes still rounds properly. // Chop off the lower 16 bits instead of OrderedDemote2To, which rounds to
const V16 w0 = hn::BitCast(d16, hn::OrderedDemote2To(dbf, f0, f1)); // the nearest bf16, because EncBytes will round again.
const V16 w1 = hn::BitCast(d16, hn::OrderedDemote2To(dbf, f2, f3)); const V16 w0 = hn::BitCast(d16, Truncate2To(dbf, f0, f1));
const V16 w1 = hn::BitCast(d16, Truncate2To(dbf, f2, f3));
return Enc2U(d16, w0, w1); return Enc2U(d16, w0, w1);
} }

View File

@ -25,12 +25,11 @@
#include <stdint.h> #include <stdint.h>
#include <stdio.h> #include <stdio.h>
#include <algorithm>
#include <random>
#include <set> #include <set>
#include "hwy/aligned_allocator.h" #include "hwy/aligned_allocator.h"
#include "hwy/base.h" #include "hwy/base.h"
#include "hwy/timer.h"
// clang-format off // clang-format off
#undef HWY_TARGET_INCLUDE #undef HWY_TARGET_INCLUDE
@ -39,13 +38,12 @@
#include "hwy/foreach_target.h" // IWYU pragma: keep #include "hwy/foreach_target.h" // IWYU pragma: keep
// Any highway.h must come after foreach_target.h // Any highway.h must come after foreach_target.h
// copybara:import_next_line:gemma_cpp // copybara:import_next_line:gemma_cpp
#include "compression/distortion.h"
// copybara:import_next_line:gemma_cpp
#include "compression/sfp-inl.h" #include "compression/sfp-inl.h"
// copybara:import_next_line:gemma_cpp
#include "compression/test_util.h"
#include "hwy/highway.h" #include "hwy/highway.h"
#include "hwy/tests/hwy_gtest.h" #include "hwy/tests/hwy_gtest.h"
#include "hwy/tests/test_util-inl.h" #include "hwy/tests/test_util-inl.h"
#include "hwy/timer.h"
HWY_BEFORE_NAMESPACE(); HWY_BEFORE_NAMESPACE();
namespace gcpp { namespace gcpp {
@ -358,25 +356,31 @@ struct TestDot {
template <typename T, class D> template <typename T, class D>
HWY_INLINE void operator()(T /*unused*/, D d) { HWY_INLINE void operator()(T /*unused*/, D d) {
const hn::Repartition<float, D> df; const hn::Repartition<float, D> df;
const size_t num = 384; const size_t num = 1024; // not too many for GeometricMean overflow.
auto in = hwy::AllocateAligned<T>(num); auto in = hwy::AllocateAligned<T>(num);
auto dec = hwy::AllocateAligned<T>(num); auto dec = hwy::AllocateAligned<T>(num);
auto vec = hwy::AllocateAligned<T>(num); auto vec = hwy::AllocateAligned<T>(num);
auto sfp = hwy::AllocateAligned<SfpStream>(num); auto sfp = hwy::AllocateAligned<SfpStream>(num);
HWY_ASSERT(in && dec && vec && sfp); HWY_ASSERT(in && dec && vec && sfp);
std::mt19937 rng(123); // Generate inputs and verify their distribution.
std::normal_distribution<float> dist{0.001f, 0.3f}; hwy::RandomState rng;
Stats in_stats;
for (size_t i = 0; i < num; ++i) { for (size_t i = 0; i < num; ++i) {
in[i] = hwy::ConvertScalarTo<T>(dist(rng)); const float r = static_cast<float>(RandomGaussian(rng));
vec[i] = hwy::ConvertScalarTo<T>(dist(rng)); in_stats.Notify(r);
in[i] = hwy::ConvertScalarTo<T>(r);
} }
// This changes the correlation between in and vec, which considerably for (size_t i = 0; i < num; ++i) {
// affects the error of the result. const float r = static_cast<float>(RandomGaussian(rng));
std::shuffle(in.get(), in.get() + num, rng); in_stats.Notify(r);
vec[i] = hwy::ConvertScalarTo<T>(r);
}
VerifyGaussian(in_stats);
SfpCodec::Enc(d, in.get(), num, sfp.get()); SfpCodec::Enc(d, in.get(), num, sfp.get());
// Compute dot product without decompression.
double actual = 0.0; double actual = 0.0;
double elapsed = hwy::HighestValue<double>(); double elapsed = hwy::HighestValue<double>();
for (size_t rep = 0; rep < 200; ++rep) { for (size_t rep = 0; rep < 200; ++rep) {
@ -393,26 +397,44 @@ struct TestDot {
} }
SfpCodec::Dec(d, sfp.get(), num, dec.get()); SfpCodec::Dec(d, sfp.get(), num, dec.get());
fprintf(stderr, "Vec %zu Dot %.2f MB/s\n", Lanes(d) * sizeof(T), fprintf(stderr, "Vec %zu Dot %zu-bit %.2f MB/s\n", Lanes(d) * sizeof(T),
num * sizeof(T) * 1E-6 / elapsed); sizeof(T) * 8, num * sizeof(T) * 1E-6 / elapsed);
double expected = 0.0; // using original input // Exact and decompressed dot products for comparison.
double expected2 = 0.0; // using decoded SFP float exact = 0.0f; // using original input
float expected = 0.0f; // using decoded SFP
DistortionStats dec_stats;
Stats ratios;
for (size_t i = 0; i < num; ++i) { for (size_t i = 0; i < num; ++i) {
expected += hwy::ConvertScalarTo<double>(in[i]) * const float in1 = hwy::ConvertScalarTo<float>(in[i]);
hwy::ConvertScalarTo<double>(vec[i]); const float dec1 = hwy::ConvertScalarTo<float>(dec[i]);
expected2 += hwy::ConvertScalarTo<double>(dec[i]) * const float vec1 = hwy::ConvertScalarTo<float>(vec[i]);
hwy::ConvertScalarTo<double>(vec[i]); dec_stats.Notify(in1, dec1);
exact += in1 * vec1;
expected += dec1 * vec1;
if (expected != 0.0f) {
ratios.Notify(exact / expected);
} }
const double l1 = hwy::ScalarAbs(expected - actual); }
const double snr = 1.0 + hwy::ScalarAbs(expected) / l1; const double dec_snr = dec_stats.GeomeanValueDivL1();
fprintf(stderr, "expected %.3f e2 %.4f actual %.4f l1 %E snr %.2f\n", const double dot_snr = 1.0 / hwy::ScalarAbs(1.0 - ratios.GeometricMean());
expected, expected2, actual, l1, snr); // exact and actual fluctuate due to the combination of SFP imprecision,
HWY_ASSERT(hwy::ScalarAbs(expected2 - actual) < 1E-4); // and whether vec[i] is negative or positive, so this is quite loose.
const double expected_l1 = sizeof(T) == 2 ? 1.52E-2 : 1.15E-2; const float final_ratio = HWY_MIN(exact / actual, actual / exact);
const double expected_snr = sizeof(T) == 2 ? 80.1f : 104.9f; fprintf(stderr, "ratios %s\n", ratios.ToString().c_str());
HWY_ASSERT(expected_l1 <= l1 && l1 < 1.02f * expected_l1); fprintf(stderr,
HWY_ASSERT(expected_snr <= snr && snr < 1.01f * expected_snr); "exact %.3f e2 %.4f actual %.4f final_ratio %.3f dec_snr %.2f "
"dot_snr %.2f\n",
exact, expected, actual, final_ratio, dec_snr, dot_snr);
// Final values are not too far apart.
HWY_ASSERT(0.87f <= final_ratio && final_ratio <= 1.0f);
// Decompressed and uncompressed dot should match exactly.
HWY_ASSERT(hwy::ScalarAbs(expected - actual) < 1E-4f);
// dec[] is close to in[], but we already check that in TestEncDec.
HWY_ASSERT(dec_snr >= 50.0);
// Geomean of ratios for each i should be very close to one.
HWY_ASSERT(dot_snr >= (sizeof(T) == 2 ? 70.0 : 1000.0));
} }
}; };

64
compression/test_util.h Normal file
View File

@ -0,0 +1,64 @@
// Copyright 2023 Google LLC
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef THIRD_PARTY_GEMMA_CPP_COMPRESSION_TEST_UTIL_H_
#define THIRD_PARTY_GEMMA_CPP_COMPRESSION_TEST_UTIL_H_
#include <stddef.h>
#include <stdint.h>
#include <cmath>
#include "hwy/base.h"
// IWYU pragma: begin_exports
// copybara:import_next_line:gemma_cpp
#include "compression/distortion.h"
// copybara:import_next_line:gemma_cpp
#include "compression/stats.h"
#include "hwy/tests/test_util.h" // RandomState
// IWYU pragma: end_exports
namespace gcpp {
// Returns random Gaussian (mean=0, stddev=1/3 similar to expected weights)
// using the central limit theorem. Avoid std::normal_distribution for
// consistent cross-platform output.
HWY_INLINE double RandomGaussian(hwy::RandomState& rng) {
uint64_t sum = 0;
constexpr int kReps = 40;
for (int rep = 0; rep < kReps; ++rep) {
sum += hwy::Random32(&rng) & 0xFFFFF;
}
const double sum_f =
static_cast<double>(sum) / static_cast<double>(0xFFFFF * kReps);
HWY_ASSERT(0.0 <= sum_f && sum_f <= 1.0);
const double plus_minus_1 = 2.0 * sum_f - 1.0;
HWY_ASSERT(-1.0 <= plus_minus_1 && plus_minus_1 <= 1.0);
// Normalize by stddev of sum of uniform random scaled to [-1, 1].
return plus_minus_1 * std::sqrt(kReps / 3.0);
};
HWY_INLINE void VerifyGaussian(Stats& stats) {
const double stddev = stats.StandardDeviation();
HWY_ASSERT(-0.01 <= stats.Mean() && stats.Mean() <= 0.01);
HWY_ASSERT(0.30 <= stddev && stddev <= 0.35);
HWY_ASSERT(-1.1 <= stats.Min() && stats.Min() <= -0.9);
HWY_ASSERT(0.9 <= stats.Max() && stats.Max() <= 1.1);
}
} // namespace gcpp
#endif // THIRD_PARTY_GEMMA_CPP_COMPRESSION_TEST_UTIL_H_