Skip to content

Commit

Permalink
Addressing PR#19096 review comments
Browse files Browse the repository at this point in the history
  • Loading branch information
@sergey-kozub
sergey-kozub committed Nov 27, 2024
1 parent a132809 commit b20cfd2
Show file tree
Hide file tree
Showing 10 changed files with 246 additions and 112 deletions.
10 changes: 5 additions & 5 deletions xla/hlo/builder/lib/math_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -96,20 +96,20 @@ class MathTypedTest : public MathTest {

bool has_inf = std::numeric_limits<T>::has_infinity;
bool has_nan = std::numeric_limits<T>::has_quiet_NaN;
bool is_finite = !has_inf && !has_nan;
bool is_nan_only = !has_inf && has_nan;
bool has_finite = !has_inf && !has_nan;
bool has_nan_only = !has_inf && has_nan;

auto expected = LiteralUtil::MakeTupleOwned(
LiteralUtil::CreateR1<bool>({true, true, true, true, true, is_finite,
is_finite, is_finite, is_finite}),
LiteralUtil::CreateR1<bool>({true, true, true, true, true, has_finite,
has_finite, has_finite, has_finite}),
LiteralUtil::CreateR1<bool>({false, false, false, false, false, has_inf,
has_inf, false, false}),
LiteralUtil::CreateR1<bool>(
{false, false, false, false, false, has_inf, false, false, false}),
LiteralUtil::CreateR1<bool>(
{false, false, false, false, false, false, has_inf, false, false}),
LiteralUtil::CreateR1<bool>({false, false, false, false, false,
is_nan_only, is_nan_only, has_nan,
has_nan_only, has_nan_only, has_nan,
has_nan}));
ComputeAndCompareLiteral(&b, expected, {});
}
Expand Down
9 changes: 1 addition & 8 deletions xla/literal.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -593,13 +593,6 @@ class LiteralBase {
static_assert(8 % bits_per_element == 0);

constexpr int elements_per_byte = 8 / bits_per_element;
constexpr auto cast = [](NativeT x) -> uint8_t {
if constexpr (primitive_util::IsFloatingPointType(primitive_type)) {
return Eigen::numext::bit_cast<uint8_t>(x);
}
return static_cast<uint8_t>(x);
};

int64_t bytes = elements.size() / elements_per_byte;
for (int64_t i = 0; i < bytes; ++i) {
uint8_t byte = 0;
Expand Down Expand Up @@ -710,14 +703,14 @@ class LiteralBase {
static_assert(!primitive_util::IsComplexType(primitive_type));
static_assert(8 % bits_per_element == 0);

constexpr int elements_per_byte = 8 / bits_per_element;
constexpr auto cast = [](uint8_t x) -> NativeT {
if constexpr (primitive_util::IsFloatingPointType(primitive_type)) {
return Eigen::numext::bit_cast<NativeT>(x);
}
return static_cast<NativeT>(x);
};

constexpr int elements_per_byte = 8 / bits_per_element;
int64_t bytes = elements.size() / elements_per_byte;
for (int64_t i = 0; i < bytes; ++i) {
uint8_t byte;
Expand Down
4 changes: 2 additions & 2 deletions xla/literal_comparison_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ TYPED_TEST(LiteralComparisonTest, CompareNear_NotEqual_4ulps) {
auto actual = LiteralUtil::CreateR0<TypeParam>(TypeParam(1.0));
float expV = 1.5; // F8E4M3*
if (type == F8E5M2 || type == F8E5M2FNUZ)
expV = 1.75;
expV = 2.0;
else if (type == F8E3M4)
expV = 1.25;
else if (type == F4E2M1FN)
Expand All @@ -99,7 +99,7 @@ TYPED_TEST(LiteralComparisonTest, FloatUsingCompareNear_NotEqual_4ulps) {
auto actual = LiteralUtil::CreateR0<float>(1.0);
float expV = 1.51; // F8E4M3*
if (type == F8E5M2 || type == F8E5M2FNUZ)
expV = 1.76;
expV = 2.01;
else if (type == F8E3M4)
expV = 1.26;
else if (type == F4E2M1FN)
Expand Down
2 changes: 1 addition & 1 deletion xla/python/ifrt/dtype.proto
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,7 +81,7 @@ message DTypeProto {
// collision.
KIND_STRING = 99;

// Next: 31
// Next: 32
}
// LINT.ThenChange()
Kind kind = 1;
Expand Down
2 changes: 1 addition & 1 deletion xla/python/ifrt/dtype_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ TEST(DTypeTest, BitSize) {
{DType::kF8E3M4, 8}, {DType::kF8E4M3, 8},
{DType::kF8E4M3FN, 8}, {DType::kF8E4M3B11FNUZ, 8},
{DType::kF8E4M3FNUZ, 8}, {DType::kF8E5M2, 8},
{DType::kF8E5M2FNUZ, 8}, {DType::kF8E8M0FNU, 4},
{DType::kF8E5M2FNUZ, 8}, {DType::kF8E8M0FNU, 8},
{DType::kS16, 16}, {DType::kU16, 16},
{DType::kF16, 16}, {DType::kBF16, 16},
{DType::kS32, 32}, {DType::kU32, 32},
Expand Down
237 changes: 168 additions & 69 deletions xla/service/elemental_ir_emitter.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -795,118 +795,217 @@ llvm::Value* EmitF8e4m3b11fnuzToF16(llvm::Value* f8_value,

absl::StatusOr<llvm::Value*> EmitF16ToF4e2m1fn(llvm::Value* f16_value,
llvm::IRBuilder<>* b) {
auto i8_const = [&](int val) {
return llvm::ConstantInt::get(b->getInt8Ty(), val);
};
auto i16_const = [&](int val) {
return llvm::ConstantInt::get(b->getInt16Ty(), val);
};
constexpr int mantissa_diff = 9; // 10 for F16, 1 for F4
constexpr int bias_diff = 14; // 15 for F16, 1 for F4

// Cast the input value to an integer for bitwise manipulation.
// Get the absolute value of the input (discard the sign).
// f16_bits = bitcast(f16_value, int)
// f16_abs_bits = f16_bits & 0x7FFF
llvm::Value* f16_bits = b->CreateBitCast(f16_value, b->getInt16Ty());
llvm::Value* f16_abs_bits = b->CreateAnd(f16_bits, i16_const(0x7FFF));

// If the input absolute value is >= 7.0 or an infinity, the result saturates
// to max value (6.0). If (0.75 <= input < 1), the result is rounded to 1.0.
// If (0 <= input <= 0.25), the result is rounded to 0.0.
// If the input is NaN, the result is undefined (implemented as minus zero).
// The rest of the cases are handled by the "happy path".
// is_overflow = f16_abs_bits >= 0x1.Cp2
// is_one = f16_abs_bits >= 0x1.8p-1 (used only if exponent underflows)
// is_zero = f16_abs_bits <= 0x1p-2 (used only if exponent underflows)
// is_nan = f16_abs_bits > 0x7C00 (F16 NaN threshold)
llvm::Value* is_overflow =
b->CreateICmpUGE(f16_abs_bits, i16_const(0x4700)); // 7.0
llvm::Value* is_one =
b->CreateICmpUGE(f16_abs_bits, i16_const(0x3A00)); // 0.75
llvm::Value* is_zero =
b->CreateICmpULE(f16_abs_bits, i16_const(0x3400)); // 0.25
llvm::Value* is_nan =
b->CreateICmpUGT(f16_abs_bits, i16_const(0x7C00)); // inf

// Truncate the mantissa to 1 bit and the exponent to 3 bits (not 2 bits, as
// the type doesn't have Inf/NaN and can represent unbiased exponent 2).
// This case, as well as the denormal, is handled below.
TF_ASSIGN_OR_RETURN(
llvm::Value * reduced_precision,
EmitReducePrecisionIR(
/*src_ty=*/F16, f16_value,
/*dest_exponent_bits=*/primitive_util::ExponentWidth(F4E2M1FN) + 1,
/*dest_mantissa_bits=*/primitive_util::SignificandWidth(F4E2M1FN) - 1,
/*quiet_nans=*/false, b));

// Cast the reduced precision value to an integer for bitwise manipulation.
// Discard the least significant (9) mantissa bits leaving 1 bit.
// Truncate to
// as_int16 = bitcast(reduced_precision, int)
// as_int8 = as_int16 >> (f16_mantissa - f4_mantissa)
llvm::Value* as_int16 = b->CreateBitCast(reduced_precision, b->getInt16Ty());
llvm::Value* as_int8 =
b->CreateTrunc(b->CreateLShr(as_int16, 9), b->getInt8Ty());
b->CreateTrunc(b->CreateLShr(as_int16, mantissa_diff), b->getInt8Ty());

// Extract sign, exponent and mantissa from reduced precision value.
auto i8_const = [&](int val) {
return llvm::ConstantInt::get(b->getInt8Ty(), val);
};
// Get the sign (0 or 1).
// f4_sign = as_int8 >> 6
llvm::Value* f4_sign = b->CreateLShr(as_int8, 6);

// Get exponent and mantissa bits without the sign.
// Important: the mask is 0x3F (not 0x7F), discard bit #6.
// f4_bits = as_int8 & 0x3F
llvm::Value* f4_bits = b->CreateAnd(as_int8, i8_const(0x3F));
llvm::Value* f4_normal = b->CreateSub(f4_bits, i8_const(28));

// Special case for exponent overflow.
auto i16_const = [&](int val) {
return llvm::ConstantInt::get(b->getInt16Ty(), val);
};
llvm::Value* f16_bits = b->CreateAnd(
b->CreateBitCast(f16_value, b->getInt16Ty()), i16_const(0x7FFF));
llvm::Value* is_overflow =
b->CreateICmpUGE(f16_bits, i16_const(0x4700)); // 7.0
llvm::Value* is_nan = b->CreateICmpUGT(f16_bits, i16_const(0x7C00)); // inf
llvm::Value* max_or_nan =
b->CreateSelect(is_nan, i8_const(0x8), i8_const(0x7));
llvm::Value* f4_normal_or_overflow =
b->CreateSelect(is_overflow, max_or_nan, f4_normal);

// Special case for exponent underflow.
// Convert F16 exponent to F4 exponent by readjusting the exponent bias.
// This produces the "normal" result, i.e. not Inf or NaN or denormal.
// f4_normal = f4_bits - ((f16_bias - f4_bias) << f4_mantissa)
constexpr int f4_exponent_offset = bias_diff << 1;
llvm::Value* f4_normal = b->CreateSub(f4_bits, i8_const(f4_exponent_offset));

// If the rounding resulted in zero exponent, the value is incorrect.
// This happens when the input is < 1.0
// is_underflow = f4_normal <= 1
llvm::Value* is_underflow = b->CreateICmpSLE(f4_normal, i8_const(1));
llvm::Value* is_one = b->CreateICmpUGE(f16_bits, i16_const(0x3A00)); // 0.75
llvm::Value* is_zero = b->CreateICmpULE(f16_bits, i16_const(0x3400)); // 0.25
llvm::Value* denorm_or_zero =
b->CreateSelect(is_zero, i8_const(0x0), i8_const(0x1));
llvm::Value* f4_small =
b->CreateSelect(is_one, i8_const(0x2), denorm_or_zero);
llvm::Value* f4_result =
b->CreateSelect(is_underflow, f4_small, f4_normal_or_overflow);

// Chain of selects that handles the special cases.
// f4_result =
// is_underflow ? (is_one ? 1.0 : (is_zero ? 0.0 : 0.5)) :
// is_overflow ? (is_nan ? -0.0 : 6.0) :
// f4_normal
llvm::Value* f4_result = b->CreateSelect(
is_underflow,
// If underflow, the input is < 1.0; the result is either 0.0, 0.5 or 1.0
b->CreateSelect(is_one, i8_const(0x2),
b->CreateSelect(is_zero, i8_const(0x0), i8_const(0x1))),
// If overflow, the input is >= 7.0 or infinity or NaN.
b->CreateSelect(is_overflow,
b->CreateSelect(is_nan, i8_const(0x8), i8_const(0x7)),
f4_normal));

// Add sign to the resulting value.
// f4_signed_result = (f4_sign << 3) | f4_result
return b->CreateOr(f4_result, b->CreateShl(f4_sign, 3));
}

llvm::Value* EmitF4e2m1fnToF16(llvm::Value* f8_value, llvm::IRBuilder<>* b) {
llvm::Value* as_int16 = b->CreateZExt(f8_value, b->getInt16Ty());

// Extract sign, exponent and mantissa from reduced precision value.
auto i16_const = [&](int val) {
return llvm::ConstantInt::get(b->getInt16Ty(), val);
};
llvm::Value* sign = b->CreateLShr(as_int16, 3);
llvm::Value* sign_shifted = b->CreateShl(sign, 15);
llvm::Value* bits = b->CreateAnd(as_int16, i16_const(0x7));
llvm::Value* bits_shifted = b->CreateShl(bits, 9);

// Re-bias the exponent and handle denormals.
llvm::Value* f16_normal = b->CreateAdd(bits_shifted, i16_const(14 << 10));
llvm::Value* is_denorm_or_zero = b->CreateICmpULE(bits, i16_const(1));
llvm::Value* is_zero = b->CreateICmpEQ(bits, i16_const(0));
llvm::Value* denorm_or_zero =
b->CreateSelect(is_zero, i16_const(0x0000), i16_const(0x3800));
llvm::Value* f16_result =
b->CreateSelect(is_denorm_or_zero, denorm_or_zero, f16_normal);
constexpr int mantissa_diff = 9; // 10 for F16, 1 for F4
constexpr int bias_diff = 14; // 15 for F16, 1 for F4

// The input value is a 8-bit integer, extend it to 16-bit integer.
// as_int16 = bitcast(f8_value, int)
llvm::Value* as_int16 = b->CreateZExt(f8_value, b->getInt16Ty());

// Get the sign and shift it to F16 position.
// f4_sign = as_int16 >> 3
// f16_sign_bit = f4_sign << 15
llvm::Value* f4_sign = b->CreateLShr(as_int16, 3);
llvm::Value* f16_sign_bit = b->CreateShl(f4_sign, 15);

// Get exponent and mantissa bits without the sign.
// f4_bits = as_int16 & 0x7
// f16_bits = f4_bits << (f16_mantissa - f4_mantissa)
llvm::Value* f4_bits = b->CreateAnd(as_int16, i16_const(0x7));
llvm::Value* f16_bits = b->CreateShl(f4_bits, mantissa_diff);

// Convert F16 exponent to F4 exponent by readjusting the exponent bias.
// f4_normal = f4_bits - ((f16_bias - f4_bias) << f4_mantissa)
constexpr int f16_exponent_offset = bias_diff << 10;
llvm::Value* f16_normal =
b->CreateAdd(f16_bits, i16_const(f16_exponent_offset));

// For denormal and zero, the exponent is different. Handle these cases
// separately below.
// is_denorm_or_zero = f4_bits <= 1
// is_zero = f4_bits == 0
llvm::Value* is_denorm_or_zero = b->CreateICmpULE(f4_bits, i16_const(1));
llvm::Value* is_zero = b->CreateICmpEQ(f4_bits, i16_const(0));

// Chain of selects that handles the special cases.
// f16_result = is_denorm_or_zero ? (is_zero ? 0.0 : 0.5) : f16_normal
llvm::Value* f16_result = b->CreateSelect(
is_denorm_or_zero,
b->CreateSelect(is_zero, i16_const(0x0000), i16_const(0x3800)),
f16_normal);

// Add sign to the resulting value.
llvm::Value* f16_signed = b->CreateOr(f16_result, sign_shifted);
return b->CreateBitCast(f16_signed, b->getHalfTy());
// f16_signed_result = f16_sign_bit | f16_result
llvm::Value* f16_signed_result = b->CreateOr(f16_result, f16_sign_bit);
return b->CreateBitCast(f16_signed_result, b->getHalfTy());
}

llvm::Value* EmitF32ToF8e8m0fnu(llvm::Value* f32_value, llvm::IRBuilder<>* b) {
llvm::Value* as_int32 = b->CreateBitCast(f32_value, b->getInt32Ty());

// Result is NaN if input is zero, negative, infinity or NaN.
auto i32_const = [&](int val) {
return llvm::ConstantInt::get(b->getInt32Ty(), val);
};
llvm::Value* is_denorm = b->CreateICmpULE(as_int32, i32_const(0x400000));
llvm::Value* is_nan =
b->CreateOr(b->CreateICmpSLE(as_int32, i32_const(0)),
b->CreateICmpSGE(as_int32, i32_const(0x7F400000)));

// Round the value and extract exponent.
llvm::Value* rounded = b->CreateAdd(as_int32, i32_const(0x400000));
llvm::Value* shifted = b->CreateAShr(rounded, 23);
llvm::Value* finite = b->CreateSelect(is_denorm, i32_const(0x00), shifted);
llvm::Value* f32_result = b->CreateSelect(is_nan, i32_const(0xFF), finite);
// Cast the input value to an integer for bitwise manipulation.
// as_int32 = bitcast(f32_value, int)
llvm::Value* as_int32 = b->CreateBitCast(f32_value, b->getInt32Ty());

// Check if the input is zero, negative, overflow, infinity or NaN.
// All of these cases cannot be represented in the E8M0 format.
// is_zero_or_negative = as_int32 <= 0
// is_overflow_or_nan = as_int32 >= 0x1.8p127
// is_nan = is_zero_or_negative | is_overflow_or_nan
llvm::Value* is_zero_or_negative = b->CreateICmpSLE(as_int32, i32_const(0));
llvm::Value* is_overflow_or_nan =
b->CreateICmpSGE(as_int32, i32_const(0x7F400000)); // 1.5 * 2^127
llvm::Value* is_nan = b->CreateOr(is_zero_or_negative, is_overflow_or_nan);

// Check if the input is a denormal which should round to the minimum value
// (2^-127), as there is no zero value.
// is_denorm = as_int32 <= 0x1p-127
llvm::Value* is_denorm =
b->CreateICmpULE(as_int32, i32_const(0x400000)); // 1.0 * 2^-127

// Round the value (always up) and discard the mantissa.
// rounded = as_int32 + 0x1p-127
// f8_normal = as_int32 >> f32_mantissa
llvm::Value* rounded =
b->CreateAdd(as_int32, i32_const(0x400000)); // 1.0 * 2^-127
llvm::Value* f8_normal = b->CreateAShr(rounded, 23);

// Chain of selects that handles the special cases.
// f8_result = is_nan ? 0xFF : (is_denorm ? 0x00 : f8_normal)
llvm::Value* f8_result =
b->CreateSelect(is_nan, i32_const(0xFF),
b->CreateSelect(is_denorm, i32_const(0x00), f8_normal));

// Truncate to the result type.
return b->CreateTrunc(f32_result, b->getInt8Ty());
return b->CreateTrunc(f8_result, b->getInt8Ty());
}

llvm::Value* EmitF8e8m0fnuToF32(llvm::Value* f8_value, llvm::IRBuilder<>* b) {
// Shift exponent to the left for the normal case.
llvm::Value* as_int32 = b->CreateZExt(f8_value, b->getInt32Ty());
llvm::Value* shifted = b->CreateShl(as_int32, 23);

// Special values for 0x00 (denorm) and 0xFF (NaN).
auto i32_const = [&](int val) {
return llvm::ConstantInt::get(b->getInt32Ty(), val);
};

// The input value is a 8-bit integer, extend it to 32-bit integer.
// as_int32 = bitcast(f8_value, int)
llvm::Value* as_int32 = b->CreateZExt(f8_value, b->getInt32Ty());

// Check if the input is a denormal or NaN.
// is_zero = as_int32 == 0x00
// is_nan = as_int32 == 0xFF
llvm::Value* is_zero = b->CreateICmpEQ(as_int32, i32_const(0));
llvm::Value* is_nan = b->CreateICmpEQ(as_int32, i32_const(0xFF));
llvm::Value* denorm_or_shifted =
b->CreateSelect(is_zero, i32_const(0x00400000), shifted);
llvm::Value* f32_result =
b->CreateSelect(is_nan, i32_const(0x7FC00000), denorm_or_shifted);

// Shift exponent to the left for the normal case.
// f32_normal = as_int32 << mantissa_diff
llvm::Value* f32_normal = b->CreateShl(as_int32, 23);

// Chain of selects that handles the special cases.
// f32_result = is_nan ? 0x7FC00000 : (is_zero ? 0x1p-127 : f32_normal)
llvm::Value* f32_result = b->CreateSelect(
is_nan, i32_const(0x7FC00000),
b->CreateSelect(is_zero, i32_const(0x400000), f32_normal));

// Bitcast integer bits to the result type.
return b->CreateBitCast(f32_result, b->getFloatTy());
}

Expand Down
Loading

0 comments on commit b20cfd2

Please sign in to comment.