Skip to content

Commit

Permalink
clearing bit utils (facebook#2301)
Browse files Browse the repository at this point in the history
Summary:
Pull Request resolved: facebook#2301

make_maskl
make_maskr
set_lzero
set_lone
set_rzero
set_lone

Simple utils that correctly handle corner cases, such as shift == 64.
I looked at the assembly a bit, probably that's ok.
For x86 I used bmi2 where was appropriate.

Differential Revision: D63329499
  • Loading branch information
DenisYaroshevskiy authored and facebook-github-bot committed Sep 27, 2024
1 parent 0d0e17a commit bef767c
Show file tree
Hide file tree
Showing 2 changed files with 350 additions and 0 deletions.
150 changes: 150 additions & 0 deletions folly/lang/Bits.h
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,10 @@
#include <folly/lang/CString.h>
#include <folly/portability/Builtins.h>

#ifdef __BMI2__
#include <immintrin.h>
#endif

#if __has_include(<bit>) && (__cplusplus >= 202002L || (defined(__cpp_lib_bit_cast) && __cpp_lib_bit_cast >= 201806L))
#include <bit>
#endif
Expand Down Expand Up @@ -106,6 +110,11 @@ constexpr std::make_unsigned_t<Dst> bits_to_unsigned(Src const s) {
static_assert(std::is_unsigned<Dst>::value, "signed type");
return static_cast<Dst>(to_unsigned(s));
}

template <typename T>
inline constexpr bool supported_in_bits_operations_v =
std::is_unsigned_v<T> && sizeof(T) <= 8;

} // namespace detail

/// findFirstSet
Expand Down Expand Up @@ -223,6 +232,147 @@ inline constexpr T strictPrevPowTwo(T const v) {
return v > 1 ? prevPowTwo(T(v - 1)) : T(0);
}

/// make_maskr
/// make_maskr_fn
///
/// Returns an unsigned integer of type T, where n
/// least significant (right) bits are set and others are not.
template <class T>
struct make_maskr_fn {
static_assert(detail::supported_in_bits_operations_v<T>, "");

FOLLY_NODISCARD constexpr T operator()(std::uint32_t n) const {
if (!folly::is_constant_evaluated_or(true)) {
compiler_may_unsafely_assume(n <= sizeof(T) * 8);

#ifdef __BMI2__
if constexpr (sizeof(T) <= 4) {
return static_cast<T>(_bzhi_u32(static_cast<std::uint32_t>(-1), n));
}
return static_cast<T>(_bzhi_u64(static_cast<std::uint64_t>(-1), n));
#endif
}

if (sizeof(T) == 8 && n == 64) {
return static_cast<T>(-1);
}
return static_cast<T>((std::uint64_t{1} << n) - 1);
}
};

template <class T>
inline constexpr make_maskr_fn<T> make_maskr;

/// make_maskl
/// make_maskl_fn
///
/// Returns an unsigned integer of type T, where n
/// most significant bits (left) are set and others are not.
template <class T>
struct make_maskl_fn {
static_assert(detail::supported_in_bits_operations_v<T>, "");

FOLLY_NODISCARD constexpr T operator()(std::uint32_t n) const {
if (!folly::is_constant_evaluated_or(true)) {
compiler_may_unsafely_assume(n <= sizeof(T) * 8);

#ifdef __BMI2__
// assembler looks smaller here, if we use bzhi from `set_lowest_n_bits`
if constexpr (sizeof(T) == 8) {
return static_cast<T>(~make_maskr<T>(64 - n));
}
#endif
}

if (sizeof(T) == 8 && n == 0) {
return 0;
}
n = sizeof(T) * 8 - n;

std::uint64_t ones = static_cast<T>(~0);
return static_cast<T>(ones << n);
}
};

template <class T>
inline constexpr make_maskl_fn<T> make_maskl;

/// set_rzero
/// set_rzero_fn
///
/// Clears n least significant (right) bits. Other bits stay the same.
struct set_rzero_fn {
template <typename T>
FOLLY_NODISCARD constexpr T operator()(T x, std::uint32_t n) const {
static_assert(detail::supported_in_bits_operations_v<T>, "");

// alternative is to do two shifts but that has
// a dependency between them, so is likely worse
return x & make_maskl<T>(sizeof(T) * 8 - n);
}
};

inline constexpr set_rzero_fn set_rzero;

/// set_rone
/// set_rone_fn
///
/// Sets n least significant (right) bits. Other bits stay the same.
struct set_rone_fn {
template <typename T>
FOLLY_NODISCARD constexpr T operator()(T x, std::uint32_t n) const {
static_assert(detail::supported_in_bits_operations_v<T>, "");

// alternative is to do two shifts but that has
// a dependency between them, so is likely worse
return x | make_maskr<T>(n);
}
};

inline constexpr set_rone_fn set_rone;

/// set_lzero
/// set_lzero_fn
///
/// Clears n most significant (left) bits. Other bits stay the same.
struct set_lzero_fn {
template <typename T>
FOLLY_NODISCARD constexpr T operator()(T x, std::uint32_t n) const {
static_assert(detail::supported_in_bits_operations_v<T>, "");

if (!folly::is_constant_evaluated_or(true)) {
compiler_may_unsafely_assume(n <= sizeof(T) * 8);

#ifdef __BMI2__
if constexpr (sizeof(T) <= 4) {
return static_cast<T>(_bzhi_u32(x, sizeof(T) * 8 - n));
}
return static_cast<T>(_bzhi_u64(x, sizeof(T) * 8 - n));
#endif
}

// alternative is to do two shifts but that has
// a dependency between them, so is likely worse
return x & make_maskr<T>(sizeof(T) * 8 - n);
}
};

inline constexpr set_lzero_fn set_lzero;

/// set_lone
/// set_lone_fn
///
/// Sets n most significant (left) bits. Other bits stay the same.
struct set_lone_fn {
template <typename T>
FOLLY_NODISCARD constexpr T operator()(T x, std::uint32_t n) const {
static_assert(detail::supported_in_bits_operations_v<T>, "");
return x | make_maskl<T>(n);
}
};

inline constexpr set_lone_fn set_lone;

/**
* Endianness detection and manipulation primitives.
*/
Expand Down
200 changes: 200 additions & 0 deletions folly/lang/test/BitsTest.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,14 @@ void testEFS() {
}
}

template <typename T>
struct BitsAllUintsTest : ::testing::Test {};

using UintsToTest =
::testing::Types<std::uint8_t, std::uint16_t, std::uint32_t, std::uint64_t>;

TYPED_TEST_SUITE(BitsAllUintsTest, UintsToTest);

} // namespace

TEST(Bits, FindFirstSet) {
Expand Down Expand Up @@ -350,4 +358,196 @@ TEST(Bits, LoadUnalignedUB) {
EXPECT_EQ(0, x);
}

TYPED_TEST(BitsAllUintsTest, MakeMaskR) {
using T = TypeParam;

static_assert(make_maskr<T>(0) == 0b0, "");
static_assert(make_maskr<T>(1) == 0b1, "");
static_assert(make_maskr<T>(2) == 0b11, "");
static_assert(make_maskr<T>(3) == 0b111, "");
static_assert(make_maskr<T>(4) == 0b1111, "");

auto test = [] {
for (std::uint32_t i = 0; i <= std::min(sizeof(T) * 8, 63UL); ++i) {
std::uint64_t expected = (std::uint64_t{1} << i) - 1;
T actual = make_maskr<T>(i);
if (expected != actual) {
EXPECT_EQ(expected, actual) << i;
return false;
}
if (std::countr_one(expected) != static_cast<int>(i)) {
EXPECT_EQ(i, std::countr_one(expected)) << i;
return false;
}
}

if (sizeof(T) == 8) {
std::uint64_t expected = std::numeric_limits<std::uint64_t>::max();
T actual = make_maskr<T>(64);
if (expected != actual) {
EXPECT_EQ(expected, actual) << 64;
return false;
}
}

return true;
};

static_assert(test(), "");

// runtime can use a different implementation
EXPECT_TRUE(test());
}

TYPED_TEST(BitsAllUintsTest, MakeMaskL) {
using T = TypeParam;

constexpr std::size_t kBitSize = sizeof(T) * 8;

static_assert(make_maskl<T>(kBitSize) == static_cast<T>(~0b0), "");
static_assert(make_maskl<T>(kBitSize - 1) == static_cast<T>(~0b1), "");
static_assert(make_maskl<T>(kBitSize - 2) == static_cast<T>(~0b11), "");
static_assert(make_maskl<T>(kBitSize - 3) == static_cast<T>(~0b111), "");
static_assert(make_maskl<T>(kBitSize - 4) == static_cast<T>(~0b1111), "");

auto test = [] {
for (std::uint32_t i = 0; i <= kBitSize; ++i) {
T expected = ~make_maskr<T>(kBitSize - i);
T actual = make_maskl<T>(i);
if (expected != actual) {
EXPECT_EQ(expected, actual) << i;
return false;
}
if (std::countl_one(expected) != static_cast<int>(i)) {
EXPECT_EQ(i, std::countl_one(expected)) << i;
return false;
}
}
return true;
};

static_assert(test(), "");

// runtime can use a different implementation
EXPECT_TRUE(test());
}

TYPED_TEST(BitsAllUintsTest, SetRzero) {
using T = TypeParam;

constexpr std::size_t kBitSize = sizeof(T) * 8;

static_assert(set_rzero(T{0b11U}, 1U) == 0b10U, "");
static_assert(set_rzero(T{0b101U}, 1U) == 0b100U, "");

auto test = [] {
for (std::uint32_t i = 0; i <= kBitSize; ++i) {
T expected = make_maskl<T>(kBitSize - i);
T actual = set_rzero(static_cast<T>(-1), i);
if (expected != actual) {
EXPECT_EQ(expected, actual) << i;
return false;
}
if (std::countr_zero(expected) != static_cast<int>(i)) {
EXPECT_EQ(i, std::countr_zero(expected)) << i;
return false;
}
}
return true;
};
static_assert(test(), "");

// runtime can use a different implementation
EXPECT_TRUE(test());
}

TYPED_TEST(BitsAllUintsTest, SetRone) {
using T = TypeParam;

constexpr std::size_t kBitSize = sizeof(T) * 8;

static_assert(set_rone(T{0b10U}, 1U) == 0b11U, "");
static_assert(set_rone(T{0b100U}, 1U) == 0b101U, "");
static_assert(set_rone(T{0b100U}, 2U) == 0b111U, "");

auto test = [] {
for (std::uint32_t i = 0; i <= kBitSize; ++i) {
T expected = make_maskr<T>(i);
T actual = set_rone(T{}, i);
if (expected != actual) {
EXPECT_EQ(expected, actual) << i;
return false;
}
if (std::countr_one(expected) != static_cast<int>(i)) {
EXPECT_EQ(i, std::countr_one(expected)) << i;
return false;
}
}
return true;
};
static_assert(test(), "");

// runtime can use a different implementation
EXPECT_TRUE(test());
}

TYPED_TEST(BitsAllUintsTest, SetLzero) {
using T = TypeParam;

constexpr std::size_t kBitSize = sizeof(T) * 8;

static_assert(set_lzero(T{0b101U}, kBitSize - 1) == 0b1U, "");
static_assert(set_lzero(T{0b1100U}, kBitSize - 3) == 0b100U, "");

auto test = [] {
for (std::uint32_t i = 0; i <= kBitSize; ++i) {
T expected = make_maskr<T>(kBitSize - i);
T actual = set_lzero(static_cast<T>(-1), i);
if (expected != actual) {
EXPECT_EQ(expected, actual) << i;
return false;
}
if (std::countl_zero(expected) != static_cast<int>(i)) {
EXPECT_EQ(i, std::countl_zero(expected)) << i;
return false;
}
}
return true;
};
static_assert(test(), "");

// runtime can use a different implementation
EXPECT_TRUE(test());
}

TYPED_TEST(BitsAllUintsTest, SetLone) {
using T = TypeParam;

constexpr std::size_t kBitSize = sizeof(T) * 8;

static_assert(set_lone(T{0b1}, kBitSize - 2) == static_cast<T>(~0b10), "");
static_assert(
set_lone(T{0b1100U}, kBitSize - 3) == static_cast<T>(~0b11), "");

auto test = [] {
for (std::uint32_t i = 0; i <= kBitSize; ++i) {
T expected = make_maskl<T>(i);
T actual = set_lone(static_cast<T>(0), i);
if (expected != actual) {
EXPECT_EQ(expected, actual) << i;
return false;
}
if (std::countl_one(expected) != static_cast<int>(i)) {
EXPECT_EQ(i, std::countl_one(expected)) << i;
return false;
}
}
return true;
};
static_assert(test(), "");

// runtime can use a different implementation
EXPECT_TRUE(test());
}

} // namespace folly

0 comments on commit bef767c

Please sign in to comment.