341 lines
13 KiB
C++
341 lines
13 KiB
C++
/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
|
|
|
|
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.
|
|
==============================================================================*/
|
|
#include "tensorflow/lite/experimental/shlo/bf16.h"
|
|
|
|
#include <cmath>
|
|
#include <cstdint>
|
|
#include <cstring>
|
|
#include <limits>
|
|
|
|
#include <gmock/gmock.h>
|
|
#include <gtest/gtest.h>
|
|
#include "absl/base/casts.h"
|
|
|
|
namespace shlo_ref {
|
|
namespace {
|
|
|
|
::testing::Matcher<BF16> MatchesBits(uint16_t bits) {
|
|
return ::testing::ResultOf([](BF16 y) { return absl::bit_cast<uint16_t>(y); },
|
|
::testing::Eq(bits));
|
|
}
|
|
|
|
::testing::Matcher<float> NearFloat(float x, float relative_error = 1e-3) {
|
|
return ::testing::FloatNear(x, std::abs(x) * relative_error);
|
|
}
|
|
|
|
float BinaryToFloat(uint32_t sign, uint32_t exponent, uint32_t high_mantissa,
|
|
uint32_t low_mantissa) {
|
|
float dest;
|
|
uint32_t src =
|
|
(sign << 31) + (exponent << 23) + (high_mantissa << 16) + low_mantissa;
|
|
memcpy(static_cast<void*>(&dest), static_cast<const void*>(&src),
|
|
sizeof(dest));
|
|
return dest;
|
|
}
|
|
|
|
template <typename T>
|
|
void TestRoundtrips() {
|
|
for (T value : {
|
|
-std::numeric_limits<T>::infinity(),
|
|
std::numeric_limits<T>::infinity(),
|
|
T(-1.0),
|
|
T(-0.5),
|
|
T(-0.0),
|
|
T(1.0),
|
|
T(0.5),
|
|
T(0.0),
|
|
}) {
|
|
EXPECT_EQ(value, static_cast<T>(static_cast<BF16>(value)));
|
|
}
|
|
}
|
|
|
|
TEST(BF16Test, FloatRoundtrips) { TestRoundtrips<float>(); }
|
|
|
|
TEST(BF16Test, DoubleRoundtrips) { TestRoundtrips<double>(); }
|
|
|
|
TEST(BF16Test, Float16Roundtrips) { TestRoundtrips<BF16>(); }
|
|
|
|
TEST(BF16Test, ConversionFromFloat) {
|
|
EXPECT_THAT(BF16(1.0f), MatchesBits(0x3f80));
|
|
EXPECT_THAT(BF16(0.5f), MatchesBits(0x3f00));
|
|
EXPECT_THAT(BF16(0.33333f), MatchesBits(0x3eab));
|
|
EXPECT_THAT(BF16(3.38e38f), MatchesBits(0x7f7e));
|
|
EXPECT_THAT(BF16(3.40e38f), MatchesBits(0x7f80)); // Becomes infinity.
|
|
}
|
|
|
|
TEST(BF16Test, RoundToNearestEven) {
|
|
float val1 = static_cast<float>(absl::bit_cast<BF16>(uint16_t{0x3c00}));
|
|
float val2 = static_cast<float>(absl::bit_cast<BF16>(uint16_t{0x3c01}));
|
|
float val3 = static_cast<float>(absl::bit_cast<BF16>(uint16_t{0x3c02}));
|
|
EXPECT_THAT(BF16(0.5f * (val1 + val2)), MatchesBits(0x3c00));
|
|
EXPECT_THAT(BF16(0.5f * (val2 + val3)), MatchesBits(0x3c02));
|
|
}
|
|
|
|
TEST(BF16Test, ConversionFromInt) {
|
|
EXPECT_THAT(BF16(-1), MatchesBits(0xbf80));
|
|
EXPECT_THAT(BF16(0), MatchesBits(0x0000));
|
|
EXPECT_THAT(BF16(1), MatchesBits(0x3f80));
|
|
EXPECT_THAT(BF16(2), MatchesBits(0x4000));
|
|
EXPECT_THAT(BF16(3), MatchesBits(0x4040));
|
|
EXPECT_THAT(BF16(12), MatchesBits(0x4140));
|
|
}
|
|
|
|
TEST(BF16Test, ConversionFromBool) {
|
|
EXPECT_THAT(BF16(false), MatchesBits(0x0000));
|
|
EXPECT_THAT(BF16(true), MatchesBits(0x3f80));
|
|
}
|
|
|
|
TEST(BF16Test, ConversionToBool) {
|
|
EXPECT_EQ(static_cast<bool>(BF16(3)), true);
|
|
EXPECT_EQ(static_cast<bool>(BF16(0.33333f)), true);
|
|
EXPECT_EQ(BF16(-0.0), false);
|
|
EXPECT_EQ(static_cast<bool>(BF16(0.0)), false);
|
|
}
|
|
|
|
TEST(BF16Test, ExplicitConversionToFloat) {
|
|
EXPECT_EQ(static_cast<float>(absl::bit_cast<BF16, uint16_t>(0x0000)), 0.0f);
|
|
EXPECT_EQ(static_cast<float>(absl::bit_cast<BF16, uint16_t>(0x3f80)), 1.0f);
|
|
}
|
|
|
|
TEST(BF16Test, ImplicitConversionToFloat) {
|
|
EXPECT_EQ((absl::bit_cast<BF16, uint16_t>(0x0000)), 0.0f);
|
|
EXPECT_EQ((absl::bit_cast<BF16, uint16_t>(0x3f80)), 1.0f);
|
|
}
|
|
|
|
TEST(BF16Test, Zero) {
|
|
EXPECT_EQ(BF16(0.0f), BF16(0.0f));
|
|
EXPECT_EQ(BF16(-0.0f), BF16(0.0f));
|
|
EXPECT_EQ(BF16(-0.0f), BF16(-0.0f));
|
|
EXPECT_THAT(BF16(0.0f), MatchesBits(0x0000));
|
|
EXPECT_THAT(BF16(-0.0f), MatchesBits(0x8000));
|
|
}
|
|
|
|
TEST(BF16Test, DefaultConstruct) {
|
|
EXPECT_EQ(static_cast<float>(BF16()), 0.0f);
|
|
}
|
|
|
|
TEST(BF16Test, Conversion) {
|
|
for (int i = 0; i < 100; ++i) {
|
|
float a = i + 1.25;
|
|
BF16 b = static_cast<BF16>(a);
|
|
float c = static_cast<float>(b);
|
|
EXPECT_LE(std::abs(c - a), a / 128);
|
|
}
|
|
}
|
|
|
|
TEST(BF16Test, Epsilon) {
|
|
EXPECT_LE(1.0f, static_cast<float>(std::numeric_limits<BF16>::epsilon() +
|
|
BF16(1.0f)));
|
|
EXPECT_EQ(1.0f, static_cast<float>(std::numeric_limits<BF16>::epsilon() /
|
|
BF16(2.0f) +
|
|
BF16(1.0f)));
|
|
}
|
|
|
|
TEST(BF16Test, Negate) {
|
|
EXPECT_EQ(static_cast<float>(-BF16(3.0f)), -3.0f);
|
|
EXPECT_EQ(static_cast<float>(-BF16(-4.5f)), 4.5f);
|
|
}
|
|
|
|
TEST(BF16Test, DivisionByZero) {
|
|
EXPECT_TRUE(std::isnan(static_cast<float>(BF16(0.0 / 0.0))));
|
|
EXPECT_TRUE(std::isinf(static_cast<float>(BF16(1.0 / 0.0))));
|
|
EXPECT_TRUE(std::isinf(static_cast<float>(BF16(-1.0 / 0.0))));
|
|
|
|
EXPECT_TRUE(std::isnan(BF16(0.0 / 0.0)));
|
|
EXPECT_TRUE(std::isinf(BF16(1.0 / 0.0)));
|
|
EXPECT_TRUE(std::isinf(BF16(-1.0 / 0.0)));
|
|
}
|
|
|
|
TEST(BF16Test, NonFinite) {
|
|
EXPECT_FALSE(std::isinf(
|
|
static_cast<float>(BF16(3.38e38f)))); // Largest finite number.
|
|
EXPECT_FALSE(std::isnan(static_cast<float>(BF16(0.0f))));
|
|
EXPECT_TRUE(
|
|
std::isinf(static_cast<float>(absl::bit_cast<BF16, uint16_t>(0xff80))));
|
|
EXPECT_TRUE(
|
|
std::isnan(static_cast<float>(absl::bit_cast<BF16, uint16_t>(0xffc0))));
|
|
EXPECT_TRUE(
|
|
std::isinf(static_cast<float>(absl::bit_cast<BF16, uint16_t>(0x7f80))));
|
|
EXPECT_TRUE(
|
|
std::isnan(static_cast<float>(absl::bit_cast<BF16, uint16_t>(0x7fc0))));
|
|
|
|
// Exactly same checks as above, just directly on the BF16 representation.
|
|
EXPECT_FALSE(isinf(absl::bit_cast<BF16, uint16_t>(0x7bff)));
|
|
EXPECT_FALSE(isnan(absl::bit_cast<BF16, uint16_t>(0x0000)));
|
|
EXPECT_TRUE(isinf(absl::bit_cast<BF16, uint16_t>(0xff80)));
|
|
EXPECT_TRUE(isnan(absl::bit_cast<BF16, uint16_t>(0xffc0)));
|
|
EXPECT_TRUE(isinf(absl::bit_cast<BF16, uint16_t>(0x7f80)));
|
|
EXPECT_TRUE(isnan(absl::bit_cast<BF16, uint16_t>(0x7fc0)));
|
|
|
|
EXPECT_THAT(BF16(BinaryToFloat(0x0, 0xff, 0x40, 0x0)), // +nan
|
|
MatchesBits(0x7fe0));
|
|
EXPECT_THAT(BF16(BinaryToFloat(0x1, 0xff, 0x40, 0x0)), // -nan
|
|
MatchesBits(0xffe0));
|
|
}
|
|
|
|
TEST(BF16Test, NumericLimits) {
|
|
static_assert(std::numeric_limits<BF16>::is_signed);
|
|
|
|
EXPECT_EQ(
|
|
absl::bit_cast<uint16_t>(std::numeric_limits<BF16>::infinity()),
|
|
absl::bit_cast<uint16_t>(BF16(std::numeric_limits<float>::infinity())));
|
|
// There is no guarantee that casting a 32-bit NaN to bfloat16 has a precise
|
|
// bit pattern. We test that it is in fact a NaN, then test the signaling
|
|
// bit (msb of significand is 1 for quiet, 0 for signaling).
|
|
constexpr uint16_t BFLOAT16_QUIET_BIT = 0x0040;
|
|
EXPECT_TRUE(isnan(std::numeric_limits<BF16>::quiet_NaN()));
|
|
EXPECT_TRUE(isnan(BF16(std::numeric_limits<float>::quiet_NaN())));
|
|
EXPECT_GT((absl::bit_cast<uint16_t>(std::numeric_limits<BF16>::quiet_NaN()) &
|
|
BFLOAT16_QUIET_BIT),
|
|
0);
|
|
EXPECT_GT(
|
|
(absl::bit_cast<uint16_t>(BF16(std::numeric_limits<float>::quiet_NaN())) &
|
|
BFLOAT16_QUIET_BIT),
|
|
0);
|
|
|
|
EXPECT_TRUE(isnan(std::numeric_limits<BF16>::signaling_NaN()));
|
|
EXPECT_TRUE(isnan(BF16(std::numeric_limits<float>::signaling_NaN())));
|
|
EXPECT_EQ(
|
|
0, (absl::bit_cast<uint16_t>(std::numeric_limits<BF16>::signaling_NaN()) &
|
|
BFLOAT16_QUIET_BIT));
|
|
EXPECT_EQ(0, (absl::bit_cast<uint16_t>(
|
|
BF16(std::numeric_limits<float>::signaling_NaN())) &
|
|
BFLOAT16_QUIET_BIT));
|
|
|
|
EXPECT_GT(std::numeric_limits<BF16>::min(), BF16(0.f));
|
|
EXPECT_GT(std::numeric_limits<BF16>::denorm_min(), BF16(0.f));
|
|
EXPECT_EQ(std::numeric_limits<BF16>::denorm_min() / BF16(2), BF16(0.f));
|
|
}
|
|
|
|
TEST(BF16Test, Arithmetic) {
|
|
EXPECT_EQ(static_cast<float>(BF16(2) + BF16(2)), 4);
|
|
EXPECT_EQ(static_cast<float>(BF16(2) + BF16(-2)), 0);
|
|
EXPECT_THAT(static_cast<float>(BF16(0.33333f) + BF16(0.66667f)),
|
|
NearFloat(1.0f));
|
|
EXPECT_EQ(static_cast<float>(BF16(2.0f) * BF16(-5.5f)), -11.0f);
|
|
EXPECT_THAT(static_cast<float>(BF16(1.0f) / BF16(3.0f)), NearFloat(0.3339f));
|
|
EXPECT_EQ(static_cast<float>(-BF16(4096.0f)), -4096.0f);
|
|
EXPECT_EQ(static_cast<float>(-BF16(-4096.0f)), 4096.0f);
|
|
}
|
|
|
|
TEST(BF16Test, Comparison) {
|
|
EXPECT_TRUE(BF16(1.0f) > BF16(0.5f));
|
|
EXPECT_TRUE(BF16(0.5f) < BF16(1.0f));
|
|
EXPECT_FALSE((BF16(1.0f) < BF16(0.5f)));
|
|
EXPECT_FALSE((BF16(0.5f) > BF16(1.0f)));
|
|
|
|
EXPECT_FALSE((BF16(4.0f) > BF16(4.0f)));
|
|
EXPECT_FALSE((BF16(4.0f) < BF16(4.0f)));
|
|
|
|
EXPECT_FALSE((BF16(0.0f) < BF16(-0.0f)));
|
|
EXPECT_FALSE((BF16(-0.0f) < BF16(0.0f)));
|
|
EXPECT_FALSE((BF16(0.0f) > BF16(-0.0f)));
|
|
EXPECT_FALSE((BF16(-0.0f) > BF16(0.0f)));
|
|
|
|
EXPECT_TRUE(BF16(0.2f) > BF16(-1.0f));
|
|
EXPECT_TRUE(BF16(-1.0f) < BF16(0.2f));
|
|
EXPECT_TRUE(BF16(-16.0f) < BF16(-15.0f));
|
|
|
|
EXPECT_TRUE(BF16(1.0f) == BF16(1.0f));
|
|
EXPECT_TRUE(BF16(1.0f) != BF16(2.0f));
|
|
|
|
EXPECT_FALSE((BF16(0.0 / 0.0) == BF16(0.0 / 0.0)));
|
|
EXPECT_TRUE(BF16(0.0 / 0.0) != BF16(0.0 / 0.0));
|
|
|
|
EXPECT_FALSE((BF16(1.0) == BF16(0.0 / 0.0)));
|
|
EXPECT_FALSE((BF16(1.0) < BF16(0.0 / 0.0)));
|
|
EXPECT_FALSE((BF16(1.0) > BF16(0.0 / 0.0)));
|
|
EXPECT_TRUE(BF16(1.0) != BF16(0.0 / 0.0));
|
|
|
|
EXPECT_TRUE(BF16(1.0) < BF16(1.0 / 0.0));
|
|
EXPECT_TRUE(BF16(1.0) > BF16(-1.0 / 0.0));
|
|
}
|
|
|
|
constexpr float PI = 3.14159265358979323846f;
|
|
|
|
TEST(BF16Test, BasicFunctions) {
|
|
// These calls should be found via ADL.
|
|
EXPECT_EQ(static_cast<float>(abs(BF16(3.5f))), 3.5f);
|
|
EXPECT_EQ(static_cast<float>(abs(BF16(3.5f))), 3.5f);
|
|
EXPECT_EQ(static_cast<float>(abs(BF16(-3.5f))), 3.5f);
|
|
EXPECT_EQ(static_cast<float>(abs(BF16(-3.5f))), 3.5f);
|
|
|
|
EXPECT_EQ(static_cast<float>(floor(BF16(3.5f))), 3.0f);
|
|
EXPECT_EQ(static_cast<float>(floor(BF16(3.5f))), 3.0f);
|
|
EXPECT_EQ(static_cast<float>(floor(BF16(-3.5f))), -4.0f);
|
|
EXPECT_EQ(static_cast<float>(floor(BF16(-3.5f))), -4.0f);
|
|
|
|
EXPECT_EQ(static_cast<float>(ceil(BF16(3.5f))), 4.0f);
|
|
EXPECT_EQ(static_cast<float>(ceil(BF16(3.5f))), 4.0f);
|
|
EXPECT_EQ(static_cast<float>(ceil(BF16(-3.5f))), -3.0f);
|
|
EXPECT_EQ(static_cast<float>(ceil(BF16(-3.5f))), -3.0f);
|
|
|
|
EXPECT_FLOAT_EQ(static_cast<float>(sqrt(BF16(0.0f))), 0.0f);
|
|
EXPECT_FLOAT_EQ(static_cast<float>(sqrt(BF16(0.0f))), 0.0f);
|
|
EXPECT_FLOAT_EQ(static_cast<float>(sqrt(BF16(4.0f))), 2.0f);
|
|
EXPECT_FLOAT_EQ(static_cast<float>(sqrt(BF16(4.0f))), 2.0f);
|
|
|
|
EXPECT_FLOAT_EQ(static_cast<float>(pow(BF16(0.0f), BF16(1.0f))), 0.0f);
|
|
EXPECT_FLOAT_EQ(static_cast<float>(pow(BF16(0.0f), BF16(1.0f))), 0.0f);
|
|
EXPECT_FLOAT_EQ(static_cast<float>(pow(BF16(2.0f), BF16(2.0f))), 4.0f);
|
|
EXPECT_FLOAT_EQ(static_cast<float>(pow(BF16(2.0f), BF16(2.0f))), 4.0f);
|
|
|
|
EXPECT_EQ(static_cast<float>(exp(BF16(0.0f))), 1.0f);
|
|
EXPECT_EQ(static_cast<float>(exp(BF16(0.0f))), 1.0f);
|
|
EXPECT_THAT(static_cast<float>(exp(BF16(PI))),
|
|
NearFloat(20.f + static_cast<float>(PI)));
|
|
EXPECT_THAT(static_cast<float>(exp(BF16(PI))),
|
|
NearFloat(20.f + static_cast<float>(PI)));
|
|
|
|
EXPECT_EQ(static_cast<float>(expm1(BF16(0.0f))), 0.0f);
|
|
EXPECT_EQ(static_cast<float>(expm1(BF16(0.0f))), 0.0f);
|
|
EXPECT_THAT(static_cast<float>(expm1(BF16(2.0f))), NearFloat(6.375f));
|
|
EXPECT_THAT(static_cast<float>(expm1(BF16(2.0f))), NearFloat(6.375f));
|
|
|
|
EXPECT_EQ(static_cast<float>(log(BF16(1.0f))), 0.0f);
|
|
EXPECT_EQ(static_cast<float>(log(BF16(1.0f))), 0.0f);
|
|
EXPECT_THAT(static_cast<float>(log(BF16(10.0f))), NearFloat(2.296875f));
|
|
EXPECT_THAT(static_cast<float>(log(BF16(10.0f))), NearFloat(2.296875f));
|
|
|
|
EXPECT_EQ(static_cast<float>(log1p(BF16(0.0f))), 0.0f);
|
|
EXPECT_EQ(static_cast<float>(log1p(BF16(0.0f))), 0.0f);
|
|
EXPECT_THAT(static_cast<float>(log1p(BF16(10.0f))), NearFloat(2.390625f));
|
|
EXPECT_THAT(static_cast<float>(log1p(BF16(10.0f))), NearFloat(2.390625f));
|
|
}
|
|
|
|
TEST(BF16Test, TrigonometricFunctions) {
|
|
EXPECT_THAT(cos(BF16(0.0f)), NearFloat(BF16(std::cos(0.0f))));
|
|
EXPECT_THAT(cos(BF16(0.0f)), NearFloat(BF16(std::cos(0.0f))));
|
|
EXPECT_FLOAT_EQ(cos(BF16(PI)), BF16(std::cos(PI)));
|
|
EXPECT_NEAR(cos(BF16(PI / 2)), BF16(std::cos(PI / 2)), 1e-3);
|
|
EXPECT_NEAR(cos(BF16(3 * PI / 2)), BF16(std::cos(3 * PI / 2)), 1e-2);
|
|
EXPECT_THAT(cos(BF16(3.5f)), NearFloat(BF16(std::cos(3.5f))));
|
|
|
|
EXPECT_FLOAT_EQ(sin(BF16(0.0f)), BF16(std::sin(0.0f)));
|
|
EXPECT_FLOAT_EQ(sin(BF16(0.0f)), BF16(std::sin(0.0f)));
|
|
EXPECT_NEAR(sin(BF16(PI)), BF16(std::sin(PI)), 1e-3);
|
|
EXPECT_THAT(sin(BF16(PI / 2)), NearFloat(BF16(std::sin(PI / 2))));
|
|
EXPECT_THAT(sin(BF16(3 * PI / 2)), NearFloat(BF16(std::sin(3 * PI / 2))));
|
|
EXPECT_THAT(sin(BF16(3.5f)), NearFloat(BF16(std::sin(3.5f))));
|
|
|
|
EXPECT_FLOAT_EQ(tan(BF16(0.0f)), BF16(std::tan(0.0f)));
|
|
EXPECT_FLOAT_EQ(tan(BF16(0.0f)), BF16(std::tan(0.0f)));
|
|
EXPECT_NEAR(tan(BF16(PI)), BF16(std::tan(PI)), 1e-3);
|
|
EXPECT_THAT(tan(BF16(3.5f)), NearFloat(BF16(std::tan(3.5f))));
|
|
}
|
|
|
|
} // namespace
|
|
} // namespace shlo_ref
|