diff --git a/be/benchmark/benchmark_main.cpp b/be/benchmark/benchmark_main.cpp index efb2930f8364fb..d68621982bd525 100644 --- a/be/benchmark/benchmark_main.cpp +++ b/be/benchmark/benchmark_main.cpp @@ -29,6 +29,7 @@ #include "benchmark_fmod.hpp" #include "benchmark_hll_merge.hpp" #include "benchmark_hybrid_set.hpp" +#include "benchmark_pdep_unpack.hpp" #include "benchmark_string.hpp" #include "benchmark_string_replace.hpp" #include "benchmark_zone_map_index.hpp" diff --git a/be/benchmark/benchmark_pdep_unpack.hpp b/be/benchmark/benchmark_pdep_unpack.hpp new file mode 100644 index 00000000000000..9ad3794232a9dd --- /dev/null +++ b/be/benchmark/benchmark_pdep_unpack.hpp @@ -0,0 +1,208 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +#pragma once + +#if defined(__x86_64__) && (defined(__GNUC__) || defined(__clang__)) + +#include + +#include +#include +#include + +#include "util/bit_stream_utils.inline.h" +#include "util/pdep_unpack.h" + +namespace doris { +namespace { + +constexpr int kPdepUnpackBatchSize = 32; +constexpr int kPdepUnpackL1NumValues = 1 << 12; +constexpr int kPdepUnpackL2NumValues = 1 << 18; +constexpr int kPdepUnpackLargeNumValues = 1 << 20; + +template +void scalar_unpack(const uint8_t* input, int num_values, T* output) { + constexpr int bytes_per_batch = kPdepUnpackBatchSize * BIT_WIDTH / 8; + for (int i = 0; i < num_values; i += kPdepUnpackBatchSize) { + BitPacking::Unpack32Values(input, bytes_per_batch, output + i); + input += bytes_per_batch; + } +} + +template +void pdep_unpack(const uint8_t* input, int num_values, T* output) { + for (int i = 0; i < num_values; i += kPdepUnpackBatchSize) { + PdepUnpack::unpack32(input, output + i); + input += kPdepUnpackBatchSize * BIT_WIDTH / 8; + } +} + +template +void unpack_if_supported(const uint8_t* input, int num_values, T* output) { + if constexpr (PdepUnpack::is_supported_type()) { + if constexpr (USE_PDEP) { + pdep_unpack(input, num_values, output); + } else { + scalar_unpack(input, num_values, output); + } + } else { + __builtin_unreachable(); + } +} + +template +void unpack(int bit_width, const uint8_t* input, int num_values, T* output) { +#define UNPACK_CASE(width) \ + case width: \ + unpack_if_supported(input, num_values, output); \ + return + switch (bit_width) { + UNPACK_CASE(1); + UNPACK_CASE(2); + UNPACK_CASE(3); + UNPACK_CASE(4); + UNPACK_CASE(5); + UNPACK_CASE(6); + UNPACK_CASE(7); + UNPACK_CASE(8); + UNPACK_CASE(9); + UNPACK_CASE(10); + UNPACK_CASE(11); + UNPACK_CASE(12); + UNPACK_CASE(13); + UNPACK_CASE(14); + UNPACK_CASE(15); + UNPACK_CASE(16); + UNPACK_CASE(17); + UNPACK_CASE(18); + UNPACK_CASE(19); + UNPACK_CASE(20); + UNPACK_CASE(21); + UNPACK_CASE(22); + UNPACK_CASE(23); + UNPACK_CASE(24); + UNPACK_CASE(25); + UNPACK_CASE(26); + UNPACK_CASE(27); + UNPACK_CASE(28); + UNPACK_CASE(29); + UNPACK_CASE(30); + UNPACK_CASE(31); + UNPACK_CASE(32); + default: + __builtin_unreachable(); + } +#undef UNPACK_CASE +} + +template +struct PdepUnpackBenchmarkData { + PdepUnpackBenchmarkData(int bit_width, int num_values) + : input(num_values * bit_width / 8), + scalar_output(num_values), + pdep_output(num_values) { + std::mt19937_64 rng(0x17993); + for (auto& byte : input) { + byte = static_cast(rng()); + } + unpack(bit_width, input.data(), num_values, scalar_output.data()); + } + + std::vector input; + std::vector scalar_output; + std::vector pdep_output; +}; + +template +void BM_ScalarUnpack(benchmark::State& state) { + const int bit_width = static_cast(state.range(0)); + const int num_values = static_cast(state.range(1)); + PdepUnpackBenchmarkData data(bit_width, num_values); + for (auto _ : state) { + unpack(bit_width, data.input.data(), num_values, data.scalar_output.data()); + benchmark::ClobberMemory(); + } + state.SetItemsProcessed(state.iterations() * num_values); +} + +template +void BM_PdepUnpack(benchmark::State& state) { + const int bit_width = static_cast(state.range(0)); + const int num_values = static_cast(state.range(1)); + if (!PdepUnpack::is_supported()) { + state.SkipWithError("CPU does not support BMI2 and AVX2"); + return; + } + PdepUnpackBenchmarkData data(bit_width, num_values); + unpack(bit_width, data.input.data(), num_values, data.pdep_output.data()); + if (data.scalar_output != data.pdep_output) { + state.SkipWithError("PDEP+AVX2 output differs from scalar output"); + return; + } + for (auto _ : state) { + unpack(bit_width, data.input.data(), num_values, data.pdep_output.data()); + benchmark::ClobberMemory(); + } + state.SetItemsProcessed(state.iterations() * num_values); +} + +template +void BM_ActualUnpack(benchmark::State& state) { + const int bit_width = static_cast(state.range(0)); + const int num_values = static_cast(state.range(1)); + PdepUnpackBenchmarkData data(bit_width, num_values); + auto result = BitPacking::UnpackValues(bit_width, data.input.data(), data.input.size(), + num_values, data.pdep_output.data()); + if (data.scalar_output != data.pdep_output) { + state.SkipWithError("Actual output differs from scalar output"); + return; + } + for (auto _ : state) { + result = BitPacking::UnpackValues(bit_width, data.input.data(), data.input.size(), + num_values, data.pdep_output.data()); + benchmark::DoNotOptimize(result); + benchmark::ClobberMemory(); + } + state.SetItemsProcessed(state.iterations() * num_values); +} + +#define REGISTER_PDEP_UNPACK_BENCHMARK(type, max_bit_width) \ + BENCHMARK_TEMPLATE(BM_ScalarUnpack, type) \ + ->ArgsProduct({benchmark::CreateDenseRange(1, max_bit_width, 1), \ + {kPdepUnpackL1NumValues, kPdepUnpackL2NumValues, \ + kPdepUnpackLargeNumValues}}); \ + BENCHMARK_TEMPLATE(BM_PdepUnpack, type) \ + ->ArgsProduct({benchmark::CreateDenseRange(1, max_bit_width, 1), \ + {kPdepUnpackL1NumValues, kPdepUnpackL2NumValues, \ + kPdepUnpackLargeNumValues}}); \ + BENCHMARK_TEMPLATE(BM_ActualUnpack, type) \ + ->ArgsProduct( \ + {benchmark::CreateDenseRange(1, max_bit_width, 1), \ + {kPdepUnpackL1NumValues, kPdepUnpackL2NumValues, kPdepUnpackLargeNumValues}}) + +REGISTER_PDEP_UNPACK_BENCHMARK(uint8_t, 8); +REGISTER_PDEP_UNPACK_BENCHMARK(uint16_t, 16); +REGISTER_PDEP_UNPACK_BENCHMARK(uint32_t, 32); + +#undef REGISTER_PDEP_UNPACK_BENCHMARK + +} // namespace +} // namespace doris + +#endif diff --git a/be/src/common/config.cpp b/be/src/common/config.cpp index b2d51ecbb6d19a..32823fb06ffb1a 100644 --- a/be/src/common/config.cpp +++ b/be/src/common/config.cpp @@ -467,6 +467,10 @@ DEFINE_mDouble(sparse_column_compaction_threshold_percent, "0.05"); // Enable RLE batch Put optimization for compaction DEFINE_mBool(enable_rle_batch_put_optimization, "true"); +// Enable PDEP-based bit unpacking. Disable it on CPUs where PDEP is microcoded and slower than +// the scalar implementation, such as AMD Zen+ and Zen 2. +DEFINE_Bool(enable_bmi2_optimizations, "true"); + // If enabled, segments will be flushed column by column DEFINE_mBool(enable_vertical_segment_writer, "true"); diff --git a/be/src/common/config.h b/be/src/common/config.h index 25c7a8ef89d1e1..55f11615bd19d9 100644 --- a/be/src/common/config.h +++ b/be/src/common/config.h @@ -536,6 +536,7 @@ DECLARE_mInt64(vertical_compaction_max_segment_size); DECLARE_mDouble(sparse_column_compaction_threshold_percent); // Enable RLE batch Put optimization for compaction DECLARE_mBool(enable_rle_batch_put_optimization); +DECLARE_Bool(enable_bmi2_optimizations); // If enabled, segments will be flushed column by column DECLARE_mBool(enable_vertical_segment_writer); diff --git a/be/src/util/bit_packing.inline.h b/be/src/util/bit_packing.inline.h index b4e0805d8aebe0..61e4caebf72ecc 100644 --- a/be/src/util/bit_packing.inline.h +++ b/be/src/util/bit_packing.inline.h @@ -20,6 +20,9 @@ #include #include "util/bit_packing.h" +#if defined(__x86_64__) && (defined(__GNUC__) || defined(__clang__)) +#include "util/pdep_unpack.h" +#endif namespace doris { inline int64_t BitPacking::NumValuesToUnpack(int bit_width, int64_t in_bytes, int64_t num_values) { @@ -83,8 +86,24 @@ std::pair BitPacking::UnpackValues(const uint8_t* __res const uint8_t* in_pos = in; OutType* out_pos = out; +#if defined(__x86_64__) && (defined(__GNUC__) || defined(__clang__)) + int64_t batches_read = 0; + if constexpr (PdepUnpack::should_use()) { + if (batches_to_read > 0 && PdepUnpack::is_supported()) { + for (; batches_read < batches_to_read; ++batches_read) { + PdepUnpack::unpack32(in_pos, out_pos); + in_pos += (BATCH_SIZE * BIT_WIDTH) / CHAR_BIT; + out_pos += BATCH_SIZE; + in_bytes -= (BATCH_SIZE * BIT_WIDTH) / CHAR_BIT; + } + } + } +#else + constexpr int64_t batches_read = 0; +#endif + // First unpack as many full batches as possible. - for (int64_t i = 0; i < batches_to_read; ++i) { + for (int64_t i = batches_read; i < batches_to_read; ++i) { in_pos = Unpack32Values(in_pos, in_bytes, out_pos); out_pos += BATCH_SIZE; in_bytes -= (BATCH_SIZE * BIT_WIDTH) / CHAR_BIT; diff --git a/be/src/util/pdep_unpack.h b/be/src/util/pdep_unpack.h new file mode 100644 index 00000000000000..10403d94e4a158 --- /dev/null +++ b/be/src/util/pdep_unpack.h @@ -0,0 +1,169 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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. + +#pragma once + +#if defined(__x86_64__) && (defined(__GNUC__) || defined(__clang__)) + +#include + +#include +#include +#include +#include +#include +#include + +#include "common/config.h" + +namespace doris { + +class PdepUnpack { +public: + static bool is_supported() { + return config::enable_bmi2_optimizations && __builtin_cpu_supports("bmi2") && + __builtin_cpu_supports("avx2"); + } + + template + static constexpr bool is_supported_type() { + return BIT_WIDTH > 0 && BIT_WIDTH <= std::numeric_limits::digits && + (std::is_same_v || std::is_same_v || + std::is_same_v); + } + + template + static constexpr bool should_use() { + // Keep the generic implementation available for benchmarking all supported widths, but + // only select PDEP in the production path for widths below 16. These widths can use the + // byte/word deposit layouts and the AVX2 widening specializations below. At 16 bits and + // above, unpack32() falls back to multiple generic 64-bit PDEP groups per batch and + // competes with efficient scalar specializations, including copy-like 16- and 32-bit + // cases. Benchmarks with L1-, L2-, and larger working sets show non-monotonic results and + // repeatable regressions for multiple high widths. Because the profitable high widths are + // CPU- and working-set-dependent, an irregular per-width allowlist would not be portable; + // use the scalar implementation conservatively instead. + return is_supported_type() && BIT_WIDTH < 16; + } + + template + __attribute__((target("bmi2,avx2"))) static void unpack32(const uint8_t* input, T* output) { + static_assert(is_supported_type()); + if constexpr (std::is_same_v) { + unpack32_with_pdep(input, output); + } else if constexpr (std::is_same_v && BIT_WIDTH <= 8) { + for (int group = 0; group < 2; ++group) { + const int first_value = group * 16; + const uint8_t* group_input = input + group * 2 * BIT_WIDTH; + uint64_t expanded0 = _pdep_u64(load_packed_group<0, 8 * BIT_WIDTH>(group_input), + pdep_mask()); + uint64_t expanded1 = + _pdep_u64(load_packed_group<8 * BIT_WIDTH, 8 * BIT_WIDTH>(group_input), + pdep_mask()); + __m128i packed8 = _mm_set_epi64x(static_cast(expanded1), + static_cast(expanded0)); + __m256i unpacked16 = _mm256_cvtepu8_epi16(packed8); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(output + first_value), unpacked16); + } + } else if constexpr (std::is_same_v && BIT_WIDTH <= 8) { + for (int group = 0; group < 4; ++group) { + uint64_t expanded = + _pdep_u64(load_packed_group<0, 8 * BIT_WIDTH>(input + group * BIT_WIDTH), + pdep_mask()); + __m256i unpacked32 = _mm256_cvtepu8_epi32(_mm_cvtsi64_si128(expanded)); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(output + group * 8), unpacked32); + } + } else if constexpr (std::is_same_v && BIT_WIDTH <= 15) { + for (int group = 0; group < 4; ++group) { + const int first_value = group * 8; + const uint8_t* group_input = input + group * BIT_WIDTH; + uint64_t expanded0 = _pdep_u64(load_packed_group<0, 4 * BIT_WIDTH>(group_input), + pdep_mask()); + uint64_t expanded1 = + _pdep_u64(load_packed_group<4 * BIT_WIDTH, 4 * BIT_WIDTH>(group_input), + pdep_mask()); + __m128i packed16 = _mm_set_epi64x(static_cast(expanded1), + static_cast(expanded0)); + __m256i unpacked32 = _mm256_cvtepu16_epi32(packed16); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(output + first_value), unpacked32); + } + } else { + unpack32_with_pdep(input, output); + } + } + +private: + template + static constexpr uint64_t pdep_mask() { + constexpr int lane_bits = sizeof(T) * 8; + constexpr int lanes = 64 / lane_bits; + constexpr uint64_t value_mask = (1ULL << BIT_WIDTH) - 1; + uint64_t mask = 0; + for (int lane = 0; lane < lanes; ++lane) { + mask |= value_mask << (lane * lane_bits); + } + return mask; + } + + template + static uint64_t load_packed_group(const uint8_t* input) { + constexpr int byte_offset = BIT_OFFSET / 8; + constexpr int shift = BIT_OFFSET % 8; + constexpr int bytes_needed = (shift + PACKED_BITS + 7) / 8; + static_assert(PACKED_BITS <= 64); + static_assert(bytes_needed <= 9); + + uint64_t low = 0; + std::memcpy(&low, input + byte_offset, bytes_needed < 8 ? bytes_needed : 8); + if constexpr (bytes_needed <= 8) { + return low >> shift; + } else { + uint8_t high = input[byte_offset + 8]; + return static_cast((static_cast(high) << 64 | low) >> + shift); + } + } + + template + __attribute__((target("bmi2"))) static void unpack_group(const uint8_t* input, T* output) { + constexpr int lanes = 64 / (sizeof(T) * 8); + constexpr int first_value = GROUP * lanes; + constexpr int bit_offset = first_value * BIT_WIDTH; + constexpr int packed_bits = lanes * BIT_WIDTH; + uint64_t packed = load_packed_group(input); + uint64_t expanded = _pdep_u64(packed, pdep_mask()); + std::memcpy(output + first_value, &expanded, sizeof(expanded)); + } + + template + __attribute__((target("bmi2"))) static void unpack32_with_pdep_impl( + const uint8_t* input, T* output, std::index_sequence) { + (unpack_group(input, output), ...); + } + + template + __attribute__((target("bmi2"))) static void unpack32_with_pdep(const uint8_t* input, + T* output) { + constexpr int lanes = 64 / (sizeof(T) * 8); + unpack32_with_pdep_impl(input, output, + std::make_index_sequence<32 / lanes> {}); + } +}; + +} // namespace doris + +#endif diff --git a/be/test/util/bit_packing_test.cpp b/be/test/util/bit_packing_test.cpp new file mode 100644 index 00000000000000..9bab326c85f553 --- /dev/null +++ b/be/test/util/bit_packing_test.cpp @@ -0,0 +1,125 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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 +#include +#include +#include + +#include "util/bit_stream_utils.inline.h" +#include "util/faststring.h" + +namespace doris { +namespace { + +#if defined(__x86_64__) && (defined(__GNUC__) || defined(__clang__)) +static_assert(PdepUnpack::should_use()); +static_assert(PdepUnpack::should_use()); +static_assert(!PdepUnpack::should_use()); +static_assert(!PdepUnpack::should_use()); +static_assert(!PdepUnpack::should_use()); +#endif + +template +std::vector make_values(int bit_width, int num_values) { + const uint64_t mask = bit_width == std::numeric_limits::digits + ? std::numeric_limits::max() + : (1ULL << bit_width) - 1; + std::vector values(num_values); + for (int i = 0; i < num_values; ++i) { + values[i] = static_cast((0x9E3779B9ULL * i + 0x7F4A7C15ULL) & mask); + } + return values; +} + +template +void pack_values(const std::vector& values, int bit_width, faststring* packed) { + BitWriter writer(packed); + for (T value : values) { + writer.PutValue(value, bit_width); + } + writer.Flush(); +} + +template +void test_all_bit_widths() { + constexpr int num_values = 65; + for (int bit_width = 1; bit_width <= std::numeric_limits::digits; ++bit_width) { + std::vector expected = make_values(bit_width, num_values); + faststring packed; + pack_values(expected, bit_width, &packed); + std::vector actual(num_values); + + auto [end, values_read] = + BitPacking::UnpackValues(bit_width, reinterpret_cast(packed.data()), + packed.size(), num_values, actual.data()); + + EXPECT_EQ(values_read, num_values) << "bit_width=" << bit_width; + EXPECT_EQ(end, reinterpret_cast(packed.data()) + packed.size()) + << "bit_width=" << bit_width; + EXPECT_EQ(actual, expected) << "bit_width=" << bit_width; + } +} + +template +void test_truncated_input() { + constexpr int num_values = 65; + for (int bit_width = 1; bit_width <= std::numeric_limits::digits; ++bit_width) { + std::vector expected = make_values(bit_width, num_values); + faststring packed; + pack_values(expected, bit_width, &packed); + const int64_t input_bytes = packed.size() - 1; + const int64_t expected_values_read = input_bytes * 8 / bit_width; + const int64_t expected_bytes_read = (expected_values_read * bit_width + 7) / 8; + std::vector actual(num_values, std::numeric_limits::max()); + + auto [end, values_read] = + BitPacking::UnpackValues(bit_width, reinterpret_cast(packed.data()), + input_bytes, num_values, actual.data()); + + EXPECT_EQ(values_read, expected_values_read) << "bit_width=" << bit_width; + EXPECT_EQ(end, reinterpret_cast(packed.data()) + expected_bytes_read) + << "bit_width=" << bit_width; + EXPECT_TRUE(std::equal(expected.begin(), expected.begin() + expected_values_read, + actual.begin())) + << "bit_width=" << bit_width; + } +} + +TEST(BitPackingTest, PdepUnpackAllBitWidths) { + test_all_bit_widths(); + test_all_bit_widths(); + test_all_bit_widths(); +} + +TEST(BitPackingTest, PdepUnpackTruncatedInput) { + test_truncated_input(); + test_truncated_input(); + test_truncated_input(); +} + +#if defined(__x86_64__) && (defined(__GNUC__) || defined(__clang__)) +TEST(BitPackingTest, DisableBmi2Optimizations) { + const bool was_enabled = config::enable_bmi2_optimizations; + config::enable_bmi2_optimizations = false; + EXPECT_FALSE(PdepUnpack::is_supported()); + config::enable_bmi2_optimizations = was_enabled; +} +#endif + +} // namespace +} // namespace doris