This is an automated email from the ASF dual-hosted git repository.
Mryange pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/master by this push:
new 6875132c3d2 [improvement](be) Optimize bit unpacking with PDEP (#65738)
6875132c3d2 is described below
commit 6875132c3d2b0416d95ef6849658fb5959331735
Author: HappenLee <[email protected]>
AuthorDate: Tue Jul 21 18:03:40 2026 +0800
[improvement](be) Optimize bit unpacking with PDEP (#65738)
### What problem does this PR solve?
Issue Number: None
Related PR: None
Problem Summary:
`BitPacking::UnpackValues` currently decodes every complete 32-value
batch with scalar shifts and masks. This PR adds an x86 PDEP
implementation for `uint8_t`, `uint16_t`, and `uint32_t`, with AVX2
widening for narrow `uint16_t` and `uint32_t` values.
The optimized functions are compiled with `target("bmi2,avx2")`. Runtime
dispatch uses `__builtin_cpu_supports("bmi2")` and
`__builtin_cpu_supports("avx2")`; unsupported CPUs, architectures,
output types, and remainder values continue to use the generic scalar
implementation. Doris therefore does not require a global `-mbmi2` build
flag.
Production dispatch deliberately uses PDEP only when `BIT_WIDTH < 16`.
The generic PDEP implementation remains available to the benchmark for
all supported widths, but repeated measurements found non-monotonic,
CPU- and working-set-dependent results at 16-32 bits. A detailed code
comment explains why the scalar implementation is retained for these
widths instead of adding an irregular CPU-specific allowlist.
The PR also adds unit tests for every supported bit width, including
full batches and truncated/remainder input, plus a reproducible
benchmark comparing:
- the generic scalar kernel;
- the direct PDEP kernel;
- the actual `BitPacking::UnpackValues` dispatch path.
The benchmark covers 4K, 256K, and 1M values and validates optimized
output against the scalar reference before measuring.
#### Benchmark results
Hardware: Intel Xeon Platinum 8457C, 48 KiB L1D and 2 MiB private L2 per
core. Release build pinned to one CPU. Values below are CPU-time medians
in microseconds.
L1-sized working set: 4K `uint32_t` values, measured with 9 randomized
repetitions. Each iteration touches a 16 KiB output and 1-7 KiB of
packed input; including the benchmark's reference output, all buffers
total 33-39 KiB and fit in the 48 KiB L1D.
| bit width | scalar | PDEP | speedup |
| ---: | ---: | ---: | ---: |
| 2 | 0.628 | 0.338 | 1.86x |
| 4 | 0.551 | 0.336 | 1.64x |
| 6 | 0.969 | 0.355 | 2.73x |
| 8 | 0.224 | 0.221 | 1.01x |
| 10 | 1.216 | 0.565 | 2.15x |
| 12 | 0.870 | 0.564 | 1.54x |
| 14 | 0.828 | 0.737 | 1.12x |
L2-sized working set: 256K `uint32_t` values. The output is 1 MiB and
the packed input is 64-448 KiB.
| bit width | scalar | PDEP | speedup |
| ---: | ---: | ---: | ---: |
| 2 | 41.0 | 27.9 | 1.47x |
| 4 | 38.5 | 29.0 | 1.33x |
| 6 | 61.6 | 28.6 | 2.15x |
| 8 | 32.5 | 31.5 | 1.03x |
| 10 | 77.6 | 36.0 | 2.16x |
| 12 | 51.1 | 35.9 | 1.42x |
| 14 | 50.6 | 47.1 | 1.07x |
The result is workload- and bit-width-dependent. PDEP is faster for all
measured widths while the working set remains in L1 or L2, although the
gains at widths 8 and 14 are marginal. At 1M values, the actual PDEP
path is faster than the scalar path at widths 10, 12, and 14, but slower
at widths 2, 4, 6, and 8.
#### High-width dispatch validation
Widths 16-32 were measured twice independently. The following table is
the second run with 9 randomized repetitions; `scalar/PDEP` greater than
1 means PDEP is faster.
| bit width | 256K scalar | 256K PDEP | scalar/PDEP | 1M scalar | 1M
PDEP | scalar/PDEP |
| ---: | ---: | ---: | ---: | ---: | ---: | ---: |
| 16 | 27.84 | 55.89 | 0.498x | 234.49 | 273.85 | 0.856x |
| 17 | 62.60 | 60.04 | 1.043x | 253.42 | 255.63 | 0.991x |
| 18 | 65.04 | 56.52 | 1.151x | 263.04 | 258.88 | 1.016x |
| 19 | 77.90 | 59.92 | 1.300x | 313.56 | 261.44 | 1.199x |
| 20 | 63.27 | 55.17 | 1.147x | 283.19 | 264.64 | 1.070x |
| 21 | 67.23 | 58.28 | 1.154x | 277.33 | 269.93 | 1.027x |
| 22 | 70.41 | 57.38 | 1.227x | 284.12 | 271.71 | 1.046x |
| 23 | 68.79 | 68.49 | 1.004x | 281.42 | 289.78 | 0.971x |
| 24 | 73.75 | 56.84 | 1.297x | 342.26 | 278.40 | 1.229x |
| 25 | 67.05 | 81.08 | 0.827x | 291.52 | 330.73 | 0.881x |
| 26 | 69.09 | 79.70 | 0.867x | 290.12 | 324.68 | 0.894x |
| 27 | 66.78 | 61.01 | 1.095x | 285.06 | 293.61 | 0.971x |
| 28 | 63.18 | 75.05 | 0.842x | 290.85 | 310.27 | 0.937x |
| 29 | 64.87 | 58.74 | 1.104x | 295.88 | 304.93 | 0.970x |
| 30 | 64.47 | 60.25 | 1.070x | 303.09 | 307.15 | 0.987x |
| 31 | 53.94 | 66.08 | 0.816x | 316.73 | 310.85 | 1.019x |
| 32 | 42.03 | 71.03 | 0.592x | 297.63 | 333.45 | 0.893x |
The first independent run showed the same material regressions at widths
16, 25, 26, 28, and 32. The `uint16_t` width-16 boundary was also slower
with PDEP at 4K, 256K, and 1M values. Some high widths improve, but the
profitable set changes with the working set and has no monotonic
boundary. The production path therefore conservatively retains scalar
unpacking for all widths at or above 16. After this change, the actual
entry point was benchmarked again: width 15 follows direct PDEP, while
widths 16-32 follow the scalar path and all outputs match the scalar
reference.
### Release note
Improve bit-packed integer decoding for bit widths below 16 on BMI2 and
AVX2 capable x86 CPUs.
### Check List (For Author)
- Test:
- [x] Unit Test: `./run-be-ut.sh -j 48 --run --filter=BitPackingTest.*`
- [x] Manual test: compiled and linked `benchmark_test`; ran scalar,
direct PDEP, and actual-path cases at 4K, 256K, and 1M values
- [x] `build-support/check-format.sh`
- [x] `build-support/run-clang-tidy.sh --build-dir be/build_Release
--files be/benchmark/benchmark_main.cpp
be/test/util/bit_packing_test.cpp`
- Behavior changed: Yes. Supported x86 CPUs use PDEP for complete
32-value batches only when `BIT_WIDTH < 16`; all other cases retain the
scalar path.
- Does this need documentation: No
---------
Co-authored-by: Dongyang Li <[email protected]>
---
be/benchmark/benchmark_main.cpp | 4 +
be/benchmark/benchmark_pdep_unpack.hpp | 207 +++++++++++++++++++++++++++++++++
be/src/common/config.cpp | 4 +
be/src/common/config.h | 1 +
be/src/util/bit_packing.inline.h | 22 +++-
be/src/util/pdep_unpack.h | 169 +++++++++++++++++++++++++++
be/test/util/bit_packing_test.cpp | 130 +++++++++++++++++++++
7 files changed, 536 insertions(+), 1 deletion(-)
diff --git a/be/benchmark/benchmark_main.cpp b/be/benchmark/benchmark_main.cpp
index efb2930f836..d6ac9c113f4 100644
--- a/be/benchmark/benchmark_main.cpp
+++ b/be/benchmark/benchmark_main.cpp
@@ -29,10 +29,12 @@
#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"
#include "binary_cast_benchmark.hpp"
+#include "common/config.h"
#include "core/block/block.h"
#include "core/column/column_string.h"
#include "core/data_type/data_type.h"
@@ -76,6 +78,8 @@ BENCHMARK(Example1);
// ThreadContext + mem tracker, otherwise the allocator throws E-7412. Mirrors
// the minimal subset of be/test/testutil/run_all_tests.cpp::main.
int main(int argc, char** argv) {
+ doris::config::enable_bmi2_optimizations = true;
+
SCOPED_INIT_THREAD_CONTEXT();
doris::ExecEnv::GetInstance()->init_mem_tracker();
doris::thread_context()->thread_mem_tracker_mgr->init();
diff --git a/be/benchmark/benchmark_pdep_unpack.hpp
b/be/benchmark/benchmark_pdep_unpack.hpp
new file mode 100644
index 00000000000..65899475917
--- /dev/null
+++ b/be/benchmark/benchmark_pdep_unpack.hpp
@@ -0,0 +1,207 @@
+// 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 <benchmark/benchmark.h>
+
+#include <cstdint>
+#include <random>
+#include <vector>
+
+#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;
+const std::vector<int64_t> kPdepUnpackNumValues = {
+ 32, 64, 128, kPdepUnpackL1NumValues, kPdepUnpackL2NumValues,
kPdepUnpackLargeNumValues};
+
+template <typename T, int BIT_WIDTH>
+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<T, BIT_WIDTH>(input, bytes_per_batch,
output + i);
+ input += bytes_per_batch;
+ }
+}
+
+template <typename T, int BIT_WIDTH>
+void pdep_unpack(const uint8_t* input, int num_values, T* output) {
+ for (int i = 0; i < num_values; i += kPdepUnpackBatchSize) {
+ PdepUnpack::unpack32<T, BIT_WIDTH>(input, output + i);
+ input += kPdepUnpackBatchSize * BIT_WIDTH / 8;
+ }
+}
+
+template <typename T, int BIT_WIDTH, bool USE_PDEP>
+void unpack_if_supported(const uint8_t* input, int num_values, T* output) {
+ if constexpr (PdepUnpack::is_supported_type<T, BIT_WIDTH>()) {
+ if constexpr (USE_PDEP) {
+ pdep_unpack<T, BIT_WIDTH>(input, num_values, output);
+ } else {
+ scalar_unpack<T, BIT_WIDTH>(input, num_values, output);
+ }
+ } else {
+ __builtin_unreachable();
+ }
+}
+
+template <typename T, bool USE_PDEP>
+void unpack(int bit_width, const uint8_t* input, int num_values, T* output) {
+#define UNPACK_CASE(width) \
+ case width: \
+ unpack_if_supported<T, width, USE_PDEP>(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 <typename T>
+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<uint8_t>(rng());
+ }
+ unpack<T, false>(bit_width, input.data(), num_values,
scalar_output.data());
+ }
+
+ std::vector<uint8_t> input;
+ std::vector<T> scalar_output;
+ std::vector<T> pdep_output;
+};
+
+template <typename T>
+void BM_ScalarUnpack(benchmark::State& state) {
+ const int bit_width = static_cast<int>(state.range(0));
+ const int num_values = static_cast<int>(state.range(1));
+ PdepUnpackBenchmarkData<T> data(bit_width, num_values);
+ for (auto _ : state) {
+ unpack<T, false>(bit_width, data.input.data(), num_values,
data.scalar_output.data());
+ benchmark::ClobberMemory();
+ }
+ state.SetItemsProcessed(state.iterations() * num_values);
+}
+
+template <typename T>
+void BM_PdepUnpack(benchmark::State& state) {
+ const int bit_width = static_cast<int>(state.range(0));
+ const int num_values = static_cast<int>(state.range(1));
+ if (!PdepUnpack::is_supported()) {
+ state.SkipWithError("CPU does not support BMI2 and AVX2");
+ return;
+ }
+ PdepUnpackBenchmarkData<T> data(bit_width, num_values);
+ unpack<T, true>(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<T, true>(bit_width, data.input.data(), num_values,
data.pdep_output.data());
+ benchmark::ClobberMemory();
+ }
+ state.SetItemsProcessed(state.iterations() * num_values);
+}
+
+template <typename T>
+void BM_ActualUnpack(benchmark::State& state) {
+ const int bit_width = static_cast<int>(state.range(0));
+ const int num_values = static_cast<int>(state.range(1));
+ PdepUnpackBenchmarkData<T> 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),
kPdepUnpackNumValues}); \
+ BENCHMARK_TEMPLATE(BM_PdepUnpack, type)
\
+ ->ArgsProduct(
\
+ {benchmark::CreateDenseRange(1, max_bit_width, 1),
kPdepUnpackNumValues}); \
+ BENCHMARK_TEMPLATE(BM_ActualUnpack, type)
\
+ ->ArgsProduct(
\
+ {benchmark::CreateDenseRange(1, max_bit_width, 1),
kPdepUnpackNumValues})
+
+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 9f697a2dcbe..ffb6cde5392 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 f5c72eb776d..85f6a2bf146 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 b4e0805d8ae..d0de57a836b 100644
--- a/be/src/util/bit_packing.inline.h
+++ b/be/src/util/bit_packing.inline.h
@@ -20,6 +20,9 @@
#include <boost/preprocessor/repetition/repeat_from_to.hpp>
#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,25 @@ std::pair<const uint8_t*, int64_t>
BitPacking::UnpackValues(const uint8_t* __res
const uint8_t* in_pos = in;
OutType* out_pos = out;
+#if defined(__x86_64__) && (defined(__GNUC__) || defined(__clang__))
+ constexpr int MIN_PDEP_BATCHES = 4;
+ int64_t batches_read = 0;
+ if constexpr (PdepUnpack::should_use<OutType, BIT_WIDTH>()) {
+ if (batches_to_read >= MIN_PDEP_BATCHES && PdepUnpack::is_supported())
{
+ for (; batches_read < batches_to_read; ++batches_read) {
+ PdepUnpack::unpack32<OutType, BIT_WIDTH>(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<OutType, BIT_WIDTH>(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 00000000000..10403d94e4a
--- /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 <immintrin.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <limits>
+#include <type_traits>
+#include <utility>
+
+#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 <typename T, int BIT_WIDTH>
+ static constexpr bool is_supported_type() {
+ return BIT_WIDTH > 0 && BIT_WIDTH <= std::numeric_limits<T>::digits &&
+ (std::is_same_v<T, uint8_t> || std::is_same_v<T, uint16_t> ||
+ std::is_same_v<T, uint32_t>);
+ }
+
+ template <typename T, int BIT_WIDTH>
+ 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<T, BIT_WIDTH>() && BIT_WIDTH < 16;
+ }
+
+ template <typename T, int BIT_WIDTH>
+ __attribute__((target("bmi2,avx2"))) static void unpack32(const uint8_t*
input, T* output) {
+ static_assert(is_supported_type<T, BIT_WIDTH>());
+ if constexpr (std::is_same_v<T, uint8_t>) {
+ unpack32_with_pdep<T, BIT_WIDTH>(input, output);
+ } else if constexpr (std::is_same_v<T, uint16_t> && 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<uint8_t,
BIT_WIDTH>());
+ uint64_t expanded1 =
+ _pdep_u64(load_packed_group<8 * BIT_WIDTH, 8 *
BIT_WIDTH>(group_input),
+ pdep_mask<uint8_t, BIT_WIDTH>());
+ __m128i packed8 =
_mm_set_epi64x(static_cast<int64_t>(expanded1),
+
static_cast<int64_t>(expanded0));
+ __m256i unpacked16 = _mm256_cvtepu8_epi16(packed8);
+ _mm256_storeu_si256(reinterpret_cast<__m256i*>(output +
first_value), unpacked16);
+ }
+ } else if constexpr (std::is_same_v<T, uint32_t> && 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<uint8_t, BIT_WIDTH>());
+ __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<T, uint32_t> && 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<uint16_t,
BIT_WIDTH>());
+ uint64_t expanded1 =
+ _pdep_u64(load_packed_group<4 * BIT_WIDTH, 4 *
BIT_WIDTH>(group_input),
+ pdep_mask<uint16_t, BIT_WIDTH>());
+ __m128i packed16 =
_mm_set_epi64x(static_cast<int64_t>(expanded1),
+
static_cast<int64_t>(expanded0));
+ __m256i unpacked32 = _mm256_cvtepu16_epi32(packed16);
+ _mm256_storeu_si256(reinterpret_cast<__m256i*>(output +
first_value), unpacked32);
+ }
+ } else {
+ unpack32_with_pdep<T, BIT_WIDTH>(input, output);
+ }
+ }
+
+private:
+ template <typename T, int BIT_WIDTH>
+ 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 <int BIT_OFFSET, int PACKED_BITS>
+ 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<uint64_t>((static_cast<unsigned __int128>(high)
<< 64 | low) >>
+ shift);
+ }
+ }
+
+ template <typename T, int BIT_WIDTH, std::size_t GROUP>
+ __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<bit_offset, packed_bits>(input);
+ uint64_t expanded = _pdep_u64(packed, pdep_mask<T, BIT_WIDTH>());
+ std::memcpy(output + first_value, &expanded, sizeof(expanded));
+ }
+
+ template <typename T, int BIT_WIDTH, std::size_t... GROUPS>
+ __attribute__((target("bmi2"))) static void unpack32_with_pdep_impl(
+ const uint8_t* input, T* output, std::index_sequence<GROUPS...>) {
+ (unpack_group<T, BIT_WIDTH, GROUPS>(input, output), ...);
+ }
+
+ template <typename T, int BIT_WIDTH>
+ __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<T, BIT_WIDTH>(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 00000000000..26d6708b63e
--- /dev/null
+++ b/be/test/util/bit_packing_test.cpp
@@ -0,0 +1,130 @@
+// 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 <algorithm>
+#include <cstdint>
+#include <limits>
+#include <vector>
+
+#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<uint16_t, 15>());
+static_assert(PdepUnpack::should_use<uint32_t, 15>());
+static_assert(!PdepUnpack::should_use<uint16_t, 16>());
+static_assert(!PdepUnpack::should_use<uint32_t, 16>());
+static_assert(!PdepUnpack::should_use<uint32_t, 32>());
+#endif
+
+template <typename T>
+std::vector<T> make_values(int bit_width, int num_values) {
+ const uint64_t mask = bit_width == std::numeric_limits<T>::digits
+ ? std::numeric_limits<T>::max()
+ : (1ULL << bit_width) - 1;
+ std::vector<T> values(num_values);
+ for (int i = 0; i < num_values; ++i) {
+ values[i] = static_cast<T>((0x9E3779B9ULL * i + 0x7F4A7C15ULL) & mask);
+ }
+ return values;
+}
+
+template <typename T>
+void pack_values(const std::vector<T>& values, int bit_width, faststring*
packed) {
+ BitWriter writer(packed);
+ for (T value : values) {
+ writer.PutValue(value, bit_width);
+ }
+ writer.Flush();
+}
+
+template <typename T>
+void test_all_bit_widths(int num_values) {
+ for (int bit_width = 1; bit_width <= std::numeric_limits<T>::digits;
++bit_width) {
+ std::vector<T> expected = make_values<T>(bit_width, num_values);
+ faststring packed;
+ pack_values(expected, bit_width, &packed);
+ std::vector<T> actual(num_values);
+
+ auto [end, values_read] =
+ BitPacking::UnpackValues(bit_width, reinterpret_cast<const
uint8_t*>(packed.data()),
+ packed.size(), num_values,
actual.data());
+
+ EXPECT_EQ(values_read, num_values) << "bit_width=" << bit_width;
+ EXPECT_EQ(end, reinterpret_cast<const uint8_t*>(packed.data()) +
packed.size())
+ << "bit_width=" << bit_width;
+ EXPECT_EQ(actual, expected) << "bit_width=" << bit_width;
+ }
+}
+
+template <typename T>
+void test_truncated_input() {
+ constexpr int num_values = 160;
+ constexpr T sentinel = std::numeric_limits<T>::max();
+ for (int bit_width = 1; bit_width <= std::numeric_limits<T>::digits;
++bit_width) {
+ std::vector<T> expected = make_values<T>(bit_width, num_values);
+ faststring packed;
+ pack_values(expected, bit_width, &packed);
+ const auto* packed_data = reinterpret_cast<const
uint8_t*>(packed.data());
+ std::vector<uint8_t> input(packed_data, packed_data + packed.size() -
1);
+ const int64_t input_bytes = input.size();
+ 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<T> actual(num_values, sentinel);
+
+ auto [end, values_read] = BitPacking::UnpackValues(bit_width,
input.data(), input_bytes,
+ num_values,
actual.data());
+
+ EXPECT_EQ(values_read, expected_values_read) << "bit_width=" <<
bit_width;
+ EXPECT_EQ(end, input.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;
+ EXPECT_TRUE(std::all_of(actual.begin() + expected_values_read,
actual.end(),
+ [](T value) { return value == sentinel; }))
+ << "bit_width=" << bit_width;
+ }
+}
+
+TEST(BitPackingTest, PdepUnpackAllBitWidths) {
+ for (int num_values : {127, 128, 129}) {
+ test_all_bit_widths<uint8_t>(num_values);
+ test_all_bit_widths<uint16_t>(num_values);
+ test_all_bit_widths<uint32_t>(num_values);
+ }
+}
+
+TEST(BitPackingTest, PdepUnpackTruncatedInput) {
+ test_truncated_input<uint8_t>();
+ test_truncated_input<uint16_t>();
+ test_truncated_input<uint32_t>();
+}
+
+#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
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
