From 3dac92f88932fd2fa592147230dddd5efc8aca47 Mon Sep 17 00:00:00 2001 From: mwish Date: Tue, 10 Nov 2026 01:27:57 +0800 Subject: [PATCH 1/2] POC: [C++] Use xsimd for SpacedExpandRightward on 4/8 byte values Dispatch SpacedExpandRightward for 4 and 8 byte trivially copyable types to a 128-bit SIMD implementation (NEON, SSE4.2), keeping the run based implementation as the scalar fallback and for other widths. The bitmap is walked one 64-bit word at a time: words with few nulls or a single valid/null transition go through a run path (memmove/memset, runs kept open across words), sparse words are scattered, and the rest are expanded 8 values at a time with an xsimd byte swizzle driven by an Arrow-owned lookup table. Every table index stays in [0, 16) and null slots are cleared with a keep mask, so only defined swizzle semantics are used and the bundled xsimd 14.2 is enough. Behavior change: the SIMD path zeroes every null slot, the scalar path only the trailing null_count slots. Co-Authored-By: Claude --- cpp/src/arrow/CMakeLists.txt | 4 + cpp/src/arrow/meson.build | 2 + cpp/src/arrow/util/CMakeLists.txt | 2 + cpp/src/arrow/util/meson.build | 2 + cpp/src/arrow/util/spaced.cc | 65 ++++ cpp/src/arrow/util/spaced_benchmark.cc | 125 ++++++++ cpp/src/arrow/util/spaced_internal.h | 59 +++- cpp/src/arrow/util/spaced_simd_128.cc | 372 ++++++++++++++++++++++ cpp/src/arrow/util/spaced_simd_internal.h | 52 +++ cpp/src/arrow/util/spaced_test.cc | 209 ++++++++++++ 10 files changed, 881 insertions(+), 11 deletions(-) create mode 100644 cpp/src/arrow/util/spaced.cc create mode 100644 cpp/src/arrow/util/spaced_benchmark.cc create mode 100644 cpp/src/arrow/util/spaced_simd_128.cc create mode 100644 cpp/src/arrow/util/spaced_simd_internal.h create mode 100644 cpp/src/arrow/util/spaced_test.cc diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt index e71d92fdb880..1fb5cfaedaf3 100644 --- a/cpp/src/arrow/CMakeLists.txt +++ b/cpp/src/arrow/CMakeLists.txt @@ -571,6 +571,8 @@ set(ARROW_UTIL_SRCS util/mutex.cc util/ree_util.cc util/secure_string.cc + util/spaced.cc + util/spaced_simd_128.cc util/string.cc util/string_util.cc util/task_group.cc @@ -597,6 +599,8 @@ append_runtime_avx512_src(ARROW_UTIL_SRCS util/bpacking_simd_avx512.cc) if(ARROW_CPU_FLAG STREQUAL "x86" AND ARROW_HAVE_RUNTIME_SSE4_2) set_source_files_properties(util/bpacking_simd_128.cc PROPERTIES COMPILE_OPTIONS "${ARROW_SSE4_2_FLAG}") + set_source_files_properties(util/spaced_simd_128.cc PROPERTIES COMPILE_OPTIONS + "${ARROW_SSE4_2_FLAG}") endif() append_runtime_sve256_src(ARROW_UTIL_SRCS util/bpacking_simd_256.cc) diff --git a/cpp/src/arrow/meson.build b/cpp/src/arrow/meson.build index ae77a3d99e21..d3896c729828 100644 --- a/cpp/src/arrow/meson.build +++ b/cpp/src/arrow/meson.build @@ -209,6 +209,8 @@ arrow_util_srcs = [ 'util/mutex.cc', 'util/ree_util.cc', 'util/secure_string.cc', + 'util/spaced.cc', + 'util/spaced_simd_128.cc', 'util/string.cc', 'util/string_util.cc', 'util/task_group.cc', diff --git a/cpp/src/arrow/util/CMakeLists.txt b/cpp/src/arrow/util/CMakeLists.txt index 5bce8549cf39..557a212a0de4 100644 --- a/cpp/src/arrow/util/CMakeLists.txt +++ b/cpp/src/arrow/util/CMakeLists.txt @@ -109,6 +109,7 @@ add_arrow_test(bit-utility-test bpacking_test.cc rle_bitmap_test.cc rle_encoding_test.cc + spaced_test.cc test_common.cc) add_arrow_test(encoding-utility-test SOURCES alp/alp_test.cc) @@ -141,6 +142,7 @@ add_arrow_benchmark(machine_benchmark) add_arrow_benchmark(queue_benchmark) add_arrow_benchmark(range_benchmark) add_arrow_benchmark(small_vector_benchmark) +add_arrow_benchmark(spaced_benchmark) add_arrow_benchmark(tdigest_benchmark) add_arrow_benchmark(thread_pool_benchmark) add_arrow_benchmark(trie_benchmark) diff --git a/cpp/src/arrow/util/meson.build b/cpp/src/arrow/util/meson.build index 7dd37922090b..d83f71a152dc 100644 --- a/cpp/src/arrow/util/meson.build +++ b/cpp/src/arrow/util/meson.build @@ -261,6 +261,7 @@ util_tests = { 'bitmap_test.cc', 'bpacking_test.cc', 'rle_encoding_test.cc', + 'spaced_test.cc', 'test_common.cc', ], }, @@ -302,6 +303,7 @@ util_benchmarks = [ 'queue', 'range', 'small_vector', + 'spaced', 'tdigest', 'thread_pool', 'trie', diff --git a/cpp/src/arrow/util/spaced.cc b/cpp/src/arrow/util/spaced.cc new file mode 100644 index 000000000000..90c713b4ace4 --- /dev/null +++ b/cpp/src/arrow/util/spaced.cc @@ -0,0 +1,65 @@ +// 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 "arrow/util/dispatch_internal.h" +#include "arrow/util/spaced_internal.h" +#include "arrow/util/spaced_simd_internal.h" + +namespace arrow::util::internal { + +namespace { + +template +struct SpacedExpandRightwardDynamicFunction { + using FunctionType = decltype(&SpacedExpandRightwardScalar); + + static constexpr auto targets() { + return std::array{ + ARROW_DISPATCH_TARGET_NONE(&SpacedExpandRightwardScalar) // + ARROW_DISPATCH_TARGET_NEON(&SpacedExpandRightwardNeon) // + ARROW_DISPATCH_TARGET_SSE4_2(&SpacedExpandRightwardSse4_2) // + }; + } +}; + +template +void SpacedExpandRightwardDispatch(void* buffer, int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset) { + static const ::arrow::internal::DynamicDispatch< + SpacedExpandRightwardDynamicFunction> + dispatch; + dispatch(static_cast(buffer), num_values, null_count, valid_bits, + valid_bits_offset); +} + +} // namespace + +void SpacedExpandRightward32(void* buffer, int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset) { + SpacedExpandRightwardDispatch(buffer, num_values, null_count, valid_bits, + valid_bits_offset); +} + +void SpacedExpandRightward64(void* buffer, int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset) { + SpacedExpandRightwardDispatch(buffer, num_values, null_count, valid_bits, + valid_bits_offset); +} + +} // namespace arrow::util::internal diff --git a/cpp/src/arrow/util/spaced_benchmark.cc b/cpp/src/arrow/util/spaced_benchmark.cc new file mode 100644 index 000000000000..24a3c8706d11 --- /dev/null +++ b/cpp/src/arrow/util/spaced_benchmark.cc @@ -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 + +#include + +#include "arrow/util/bit_util.h" +#include "arrow/util/spaced_internal.h" +#include "arrow/util/spaced_simd_internal.h" + +namespace arrow::util::internal { +namespace { + +template +using ExpandFunc = void (*)(Uint*, int, int, const uint8_t*, int64_t); + +// Either independent slots valid with probability `p`, or (p < 0) runs of +// valid and null slots with the given mean lengths. +struct ValidityPattern { + double p; + int mean_valid; + int mean_null; +}; + +constexpr ValidityPattern Bernoulli(double p) { return {p, 0, 0}; } +constexpr ValidityPattern Runs(int mean_valid, int mean_null) { + return {-1, mean_valid, mean_null}; +} + +std::vector MakeBitmap(const ValidityPattern& pattern, int n) { + std::mt19937_64 rng(42); + std::vector bitmap(bit_util::BytesForBits(n)); + if (pattern.p >= 0) { + std::bernoulli_distribution valid(pattern.p); + for (int i = 0; i < n; ++i) bit_util::SetBitTo(bitmap.data(), i, valid(rng)); + } else { + std::geometric_distribution valid_len(1.0 / pattern.mean_valid); + std::geometric_distribution null_len(1.0 / pattern.mean_null); + bool valid = true; + for (int i = 0; i < n;) { + const int len = 1 + (valid ? valid_len(rng) : null_len(rng)); + for (int end = std::min(n, i + len); i < end; ++i) { + bit_util::SetBitTo(bitmap.data(), i, valid); + } + valid = !valid; + } + } + return bitmap; +} + +// The timing includes copying the dense values back in place before each +// expansion. +template +void BM_SpacedExpandRightward(benchmark::State& state, ExpandFunc expand, + ValidityPattern pattern) { + constexpr int kNumValues = 2048; + const auto bitmap = MakeBitmap(pattern, kNumValues); + int num_valid = 0; + for (int i = 0; i < kNumValues; ++i) { + num_valid += bit_util::GetBit(bitmap.data(), i); + } + const int null_count = kNumValues - num_valid; + + std::vector dense(kNumValues); + for (int i = 0; i < kNumValues; ++i) dense[i] = static_cast(i + 1); + std::vector buffer(kNumValues); + + for (auto _ : state) { + std::memcpy(buffer.data(), dense.data(), sizeof(Uint) * num_valid); + expand(buffer.data(), kNumValues, null_count, bitmap.data(), 0); + benchmark::ClobberMemory(); + } + state.SetItemsProcessed(state.iterations() * kNumValues); +} + +#define BENCHMARK_SPACED_EXPAND_PATTERNS(NAME, FUNC) \ + BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_p0_02, FUNC, Bernoulli(0.02)); \ + BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_p0_10, FUNC, Bernoulli(0.10)); \ + BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_p0_50, FUNC, Bernoulli(0.50)); \ + BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_p0_90, FUNC, Bernoulli(0.90)); \ + BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_p0_95, FUNC, Bernoulli(0.95)); \ + BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_p0_99, FUNC, Bernoulli(0.99)); \ + BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_runs_v2_n2, FUNC, Runs(2, 2)); \ + BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_runs_v8_n8, FUNC, Runs(8, 8)); \ + BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_runs_v32_n32, FUNC, Runs(32, 32)); \ + BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_runs_v64_n2, FUNC, Runs(64, 2)); \ + BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_runs_v128_n8, FUNC, Runs(128, 8)); \ + BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_runs_v256_n256, FUNC, \ + Runs(256, 256)); \ + BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_runs_v1000_n1, FUNC, Runs(1000, 1)) + +BENCHMARK_SPACED_EXPAND_PATTERNS(Scalar32, &SpacedExpandRightwardScalar); +BENCHMARK_SPACED_EXPAND_PATTERNS(Scalar64, &SpacedExpandRightwardScalar); + +#if defined(ARROW_HAVE_NEON) +BENCHMARK_SPACED_EXPAND_PATTERNS(Neon32, &SpacedExpandRightwardNeon); +BENCHMARK_SPACED_EXPAND_PATTERNS(Neon64, &SpacedExpandRightwardNeon); +#elif defined(ARROW_HAVE_SSE4_2) +BENCHMARK_SPACED_EXPAND_PATTERNS(Sse4_2_32, &SpacedExpandRightwardSse4_2); +BENCHMARK_SPACED_EXPAND_PATTERNS(Sse4_2_64, &SpacedExpandRightwardSse4_2); +#endif + +#undef BENCHMARK_SPACED_EXPAND_PATTERNS + +} // namespace +} // namespace arrow::util::internal diff --git a/cpp/src/arrow/util/spaced_internal.h b/cpp/src/arrow/util/spaced_internal.h index b43cbdd751ef..f0f3c9b23911 100644 --- a/cpp/src/arrow/util/spaced_internal.h +++ b/cpp/src/arrow/util/spaced_internal.h @@ -20,8 +20,10 @@ #include #include #include +#include #include "arrow/util/bit_run_reader.h" +#include "arrow/util/visibility.h" namespace arrow::util::internal { @@ -53,18 +55,13 @@ inline int SpacedCompress(const T* src, int num_values, const uint8_t* valid_bit /// \brief Relocate values according to a validity bitmap, to the right /// -/// Non-null values should initially be densely packed at the left of the buffer. -/// This method spreads the values out according to the given validity bitmap. -/// Null entries are zero-initialized. -/// -/// \param[in, out] buffer the in-place buffer -/// \param[in] num_values total size of buffer including null slots -/// \param[in] null_count number of null slots -/// \param[in] valid_bits bitmap data indicating position of valid slots -/// \param[in] valid_bits_offset offset into valid_bits +/// Run based implementation, see SpacedExpandRightward. Only the trailing +/// `null_count` slots of the buffer are zeroed, the other null slots keep +/// whatever value they held. template -inline void SpacedExpandRightward(T* buffer, int num_values, int null_count, - const uint8_t* valid_bits, int64_t valid_bits_offset) { +inline void SpacedExpandRightwardScalar(T* buffer, int num_values, int null_count, + const uint8_t* valid_bits, + int64_t valid_bits_offset) { // Point to end as we add the spacing from the back. int idx_decode = num_values - null_count; @@ -98,6 +95,46 @@ inline void SpacedExpandRightward(T* buffer, int num_values, int null_count, assert(idx_decode == 0); } +/// \brief SpacedExpandRightward for 4 byte values, dispatched on the CPU +ARROW_EXPORT void SpacedExpandRightward32(void* buffer, int num_values, int null_count, + const uint8_t* valid_bits, + int64_t valid_bits_offset); + +/// \brief SpacedExpandRightward for 8 byte values, dispatched on the CPU +ARROW_EXPORT void SpacedExpandRightward64(void* buffer, int num_values, int null_count, + const uint8_t* valid_bits, + int64_t valid_bits_offset); + +/// \brief Relocate values according to a validity bitmap, to the right +/// +/// Non-null values should initially be densely packed at the left of the buffer. +/// This method spreads the values out according to the given validity bitmap. +/// Null entries are zero-initialized. +/// +/// 4 and 8 byte values go through a SIMD implementation when the CPU has one, +/// which zeroes every null slot. Other widths, and CPUs without SIMD support, +/// use SpacedExpandRightwardScalar. +/// +/// \param[in, out] buffer the in-place buffer +/// \param[in] num_values total size of buffer including null slots +/// \param[in] null_count number of null slots +/// \param[in] valid_bits bitmap data indicating position of valid slots +/// \param[in] valid_bits_offset offset into valid_bits +template +inline void SpacedExpandRightward(T* buffer, int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset) { + if constexpr (std::is_trivially_copyable_v && sizeof(T) == 4) { + SpacedExpandRightward32(buffer, num_values, null_count, valid_bits, + valid_bits_offset); + } else if constexpr (std::is_trivially_copyable_v && sizeof(T) == 8) { + SpacedExpandRightward64(buffer, num_values, null_count, valid_bits, + valid_bits_offset); + } else { + SpacedExpandRightwardScalar(buffer, num_values, null_count, valid_bits, + valid_bits_offset); + } +} + /// \brief Relocate values according to a validity bitmap, to the left /// /// Non-null values should initially be densely packed at the right of the buffer. diff --git a/cpp/src/arrow/util/spaced_simd_128.cc b/cpp/src/arrow/util/spaced_simd_128.cc new file mode 100644 index 000000000000..cf6c1262c840 --- /dev/null +++ b/cpp/src/arrow/util/spaced_simd_128.cc @@ -0,0 +1,372 @@ +// 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. + +#if defined(ARROW_HAVE_NEON) +# define SPACED_EXPAND_PLATFORM SpacedExpandRightwardNeon +#elif defined(ARROW_HAVE_SSE4_2) || defined(ARROW_HAVE_RUNTIME_SSE4_2) +# define SPACED_EXPAND_PLATFORM SpacedExpandRightwardSse4_2 +#endif + +#if defined(SPACED_EXPAND_PLATFORM) + +# include +# include +# include +# include + +# include + +# include "arrow/util/bit_util.h" +# include "arrow/util/endian.h" +# include "arrow/util/macros.h" +# include "arrow/util/spaced_simd_internal.h" +# include "arrow/util/ubsan.h" + +namespace arrow::util::internal { + +namespace { + +using Arch = xsimd::default_arch; + +// --------------------------------------------------------------------------- +// Lookup tables +// +// AArch64 without FEAT_CSSC has no scalar popcount: std::popcount is a round +// trip through a SIMD register (fmov/cnt/addv/fmov), so the hot paths look +// small popcounts up instead. +// --------------------------------------------------------------------------- + +struct PopCountTable { + uint8_t count[256] = {}; + constexpr PopCountTable() { + for (unsigned m = 0; m < 256; ++m) { + count[m] = static_cast(std::popcount(m)); + } + } +}; +inline constexpr PopCountTable kPopCount{}; + +// For an 8 bit chunk mask split in sub-batches of kSubLanes lanes, byte j of +// row[mask] is the number of valid lanes in sub-batches 0..j. +template +struct PrefixTable { + uint32_t row[256] = {}; + constexpr PrefixTable() { + for (unsigned m = 0; m < 256; ++m) { + uint32_t r = 0; + for (int j = 0; j < 4 && (j + 1) * kSubLanes <= 8; ++j) { + const unsigned low = (1u << ((j + 1) * kSubLanes)) - 1u; + r |= static_cast(std::popcount(m & low)) << (8 * j); + } + row[m] = r; + } + } +}; +template +inline constexpr PrefixTable kPrefix{}; + +// Byte shuffle expanding the low lanes of a 16 byte register according to a +// kLanes bit mask. Null lanes read byte 0 so that every index stays in +// [0, 16), the only range where the semantics of a dynamic swizzle are +// defined on every architecture; `keep` then zeroes them. +template +struct ExpandTable { + struct alignas(32) Row { + uint8_t shuffle[16]; + uint8_t keep[16]; + }; + Row rows[1 << kLanes] = {}; + constexpr ExpandTable() { + constexpr int kValueBytes = 16 / kLanes; + for (int m = 0; m < (1 << kLanes); ++m) { + int prefix = 0; + for (int lane = 0; lane < kLanes; ++lane) { + const bool valid = ((m >> lane) & 1) != 0; + for (int b = 0; b < kValueBytes; ++b) { + rows[m].shuffle[lane * kValueBytes + b] = + valid ? static_cast(prefix * kValueBytes + b) : uint8_t{0}; + rows[m].keep[lane * kValueBytes + b] = valid ? uint8_t{0xFF} : uint8_t{0}; + } + if (valid) ++prefix; + } + } + } +}; +template +inline constexpr ExpandTable kExpand{}; + +// --------------------------------------------------------------------------- +// Chunk kernel: expand the low popcount(mask) values of `src` to the 8 slots +// of `dst` according to `mask`, zeroing the null slots. +// +// Reads at most 8 values from `src`: sub-batch j reads kSubLanes values at +// the number of valid lanes below it, which is at most 8 - kSubLanes. +// Everything is loaded before anything is stored, so `dst` may overlap `src`. +// --------------------------------------------------------------------------- + +template +struct ExpandChunk { + using Batch = xsimd::batch; + using Bytes = xsimd::batch; + static_assert(Bytes::size == 16, "128 bit SIMD only"); + + static constexpr int kLanes = 8; + static constexpr int kSubLanes = static_cast(Batch::size); + static constexpr int kSubs = kLanes / kSubLanes; + + static void Expand(Uint* dst, const Uint* src, unsigned mask) { + const uint32_t prefixes = kPrefix.row[mask]; + Batch sub[kSubs]; + unsigned prefix = 0; + for (int j = 0; j < kSubs; ++j) { + const unsigned m = (mask >> (j * kSubLanes)) & ((1u << kSubLanes) - 1u); + const auto& row = kExpand.rows[m]; + const auto bytes = + xsimd::bitwise_cast(Batch::load_unaligned(src + prefix)); + const auto expanded = xsimd::swizzle(bytes, Bytes::load_aligned(row.shuffle)) & + Bytes::load_aligned(row.keep); + sub[j] = xsimd::bitwise_cast(expanded); + prefix = (prefixes >> (8 * j)) & 0xFF; + } + for (int j = 0; j < kSubs; ++j) { + sub[j].store_unaligned(dst + j * kSubLanes); + } + } +}; + +// --------------------------------------------------------------------------- +// Word driver +// +// The bitmap is walked one 64 bit word at a time from the end of the buffer, +// and every word goes to one of three paths: +// * runs (at most kDenseNulls nulls, or a single valid / null transition): +// memmove whole runs of valid values and memset runs of nulls, keeping a +// run open across words so that a long run is moved only once; +// * sparse (few valid values, far from their destination): zero the word +// and scatter the values one by one; +// * otherwise: 8 chunks of the SIMD kernel. +// --------------------------------------------------------------------------- + +// Bits [pos, pos + 64) of the bitmap. Only reads the bytes holding those bits. +inline uint64_t LoadWord(const uint8_t* bitmap, int64_t pos) { + const uint8_t* p = bitmap + (pos >> 3); + const int shift = static_cast(pos & 7); + const uint64_t lo = bit_util::FromLittleEndian(SafeLoadAs(p)); + if (shift == 0) return lo; + return (lo >> shift) | (static_cast(p[8]) << (64 - shift)); +} + +// Whether the word goes to the run path: at most kDenseNulls nulls, or a +// single transition between a valid run and a null run (including empty +// words). Neither needs a popcount. +template +inline bool IsRunWord(uint64_t bits) { + uint64_t z = ~bits; + for (int k = 0; k < kDenseNulls; ++k) z &= z - 1; + return z == 0 || (bits & (bits + 1)) == 0 || (~bits & (~bits + 1)) == 0; +} + +// Handle run words downwards from the one ending at `pos`, whose bits are +// `bits`. `nulls_below` is the number of nulls below `pos` on entry, and +// below the returned position on return. Kept out of line: having the open +// run state in the main loop slows every path down. +template +ARROW_NOINLINE int ExpandRunWords(Uint* out, int pos, int& nulls_below, uint64_t bits, + const uint8_t* valid_bits, int64_t valid_bits_offset) { + // The open run: values [run_hi - x, run_hi) come from run_shift slots lower. + bool pending = false; + int run_hi = 0; + int run_shift = 0; + auto flush = [&](int run_lo) { + if (pending && run_hi > run_lo) { + std::memmove(out + run_lo, out + run_lo - run_shift, + sizeof(Uint) * static_cast(run_hi - run_lo)); + } + pending = false; + }; + while (true) { + const int lo = pos - 64; + if (bits == 0) { + // Empty words close the open run and are zeroed all at once. + flush(pos); + const int hi = pos; + pos = lo; + while (pos >= 64 && LoadWord(valid_bits, valid_bits_offset + pos - 64) == 0) { + pos -= 64; + } + std::memset(out + pos, 0, sizeof(Uint) * static_cast(hi - pos)); + nulls_below -= hi - pos; + } else { + if (!pending && (bits >> 63) != 0) { + pending = true; + run_hi = pos; + run_shift = nulls_below; + } + // One step per run of nulls [null_lo, null_hi], from the top. Zeroing + // them is safe once the run above is moved: everything below still + // comes from lower slots. + uint64_t z = ~bits; + while (z != 0) { + const int null_hi = 63 - std::countl_zero(z); + const uint64_t valid_below = bits & ((uint64_t{1} << null_hi) - 1); + const int null_lo = valid_below == 0 ? 0 : 64 - std::countl_zero(valid_below); + const int len = null_hi - null_lo + 1; + nulls_below -= len; + flush(lo + null_hi + 1); + if (len == 1) { + out[lo + null_lo] = 0; + } else { + std::memset(out + lo + null_lo, 0, sizeof(Uint) * static_cast(len)); + } + z &= (uint64_t{1} << null_lo) - 1; + if (null_lo > 0) { + pending = true; + run_hi = lo + null_lo; + run_shift = nulls_below; + } + } + pos = lo; + } + // No nulls below: the open run (shift 0) and everything below it is + // already in place. + if (nulls_below == 0) return pos; + if (pos < 64) break; + bits = LoadWord(valid_bits, valid_bits_offset + pos - 64); + if (bits == ~uint64_t{0}) { + // Full words only extend the open run: skip them two at a time. + if (!pending) { + pending = true; + run_hi = pos; + run_shift = nulls_below; + } + pos -= 64; + while (pos >= 128) { + const uint64_t a = LoadWord(valid_bits, valid_bits_offset + pos - 64); + const uint64_t b = LoadWord(valid_bits, valid_bits_offset + pos - 128); + if ((a & b) != ~uint64_t{0}) break; + pos -= 128; + } + if (pos < 64) break; + bits = LoadWord(valid_bits, valid_bits_offset + pos - 64); + } + if (!IsRunWord(bits)) break; + } + flush(pos); + return pos; +} + +template +void ExpandRightward(Uint* out, int num_values, int null_count, const uint8_t* valid_bits, + int64_t valid_bits_offset) { + using Kernel = ExpandChunk; + constexpr int kLanes = Kernel::kLanes; + constexpr int kChunks = 64 / kLanes; + // Words with at most that many valid values, far enough from their + // destination, are scattered one by one. + constexpr int kSparseValues = 8; + + if (null_count == 0) return; + int idx_decode = num_values - null_count; + if (idx_decode == 0) { + std::memset(out, 0, sizeof(Uint) * static_cast(num_values)); + return; + } + int pos = num_values; + while (pos >= 64) { + const int lo = pos - 64; + const uint64_t bits = LoadWord(valid_bits, valid_bits_offset + lo); + if (IsRunWord(bits)) { + int nulls_below = pos - idx_decode; + pos = ExpandRunWords(out, pos, nulls_below, bits, valid_bits, + valid_bits_offset); + if (nulls_below == 0) return; + idx_decode = pos - nulls_below; + continue; + } + const int pc = std::popcount(bits); + idx_decode -= pc; + const int src_begin = idx_decode; + const int shift = lo - src_begin; + if (pc <= kSparseValues && shift >= 64) { + // Source and destination do not overlap. + std::memset(out + lo, 0, 64 * sizeof(Uint)); + int q = src_begin + pc; + uint64_t b = bits; + while (b != 0) { + const int i = 63 - std::countl_zero(b); + b ^= uint64_t{1} << i; + out[lo + i] = out[--q]; + } + } else { + // The kernel reads up to kLanes values from its source, which can go + // past the valid values of the top word only. Stage those in a buffer. + Uint staged[64 + kLanes]; + const Uint* base = out; + int rem = src_begin + pc; + if (rem + kLanes > num_values) { + std::memcpy(staged, out + src_begin, sizeof(Uint) * static_cast(pc)); + base = staged; + rem = pc; + } + auto expand_chunk = [&](int c) { + const unsigned mask = static_cast((bits >> (c * kLanes)) & 0xFF); + const int pc_c = kPopCount.count[mask]; + if (pc_c == 0) { + std::memset(out + lo + c * kLanes, 0, sizeof(Uint) * kLanes); + return; + } + const int src = rem - pc_c; + Kernel::Expand(out + lo + c * kLanes, base + src, mask); + rem = src; + }; + // Unrolled by hand: Arrow builds with -O2, which keeps the loop. + [&](std::integer_sequence) { + (expand_chunk(kChunks - 1 - C), ...); + }(std::make_integer_sequence{}); + } + pos = lo; + if (shift == 0) return; + } + for (int i = pos - 1; i >= 0; --i) { + if (bit_util::GetBit(valid_bits, valid_bits_offset + i)) { + out[i] = out[--idx_decode]; + } else { + out[i] = 0; + } + } +} + +} // namespace + +template +void SPACED_EXPAND_PLATFORM(Uint* buffer, int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset) { + // Tuned on Apple M2: words with more nulls are faster through the kernel. + constexpr int kDenseNulls = sizeof(Uint) == 4 ? 1 : 2; + ExpandRightward(buffer, num_values, null_count, valid_bits, + valid_bits_offset); +} + +template void SPACED_EXPAND_PLATFORM(uint32_t*, int, int, const uint8_t*, + int64_t); +template void SPACED_EXPAND_PLATFORM(uint64_t*, int, int, const uint8_t*, + int64_t); + +} // namespace arrow::util::internal + +# undef SPACED_EXPAND_PLATFORM +#endif // SPACED_EXPAND_PLATFORM diff --git a/cpp/src/arrow/util/spaced_simd_internal.h b/cpp/src/arrow/util/spaced_simd_internal.h new file mode 100644 index 000000000000..f9ba5a0cc38f --- /dev/null +++ b/cpp/src/arrow/util/spaced_simd_internal.h @@ -0,0 +1,52 @@ +// 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 + +#include + +#include "arrow/util/visibility.h" + +namespace arrow::util::internal { + +#if defined(ARROW_HAVE_NEON) +# define SPACED_EXPAND_ARCH128 SpacedExpandRightwardNeon +#elif defined(ARROW_HAVE_SSE4_2) || defined(ARROW_HAVE_RUNTIME_SSE4_2) +# define SPACED_EXPAND_ARCH128 SpacedExpandRightwardSse4_2 +#endif + +#if defined(SPACED_EXPAND_ARCH128) + +/// SpacedExpandRightward on unsigned 4 or 8 byte lanes, with 128 bit SIMD. +/// Every null slot is zeroed. +template +ARROW_EXPORT void SPACED_EXPAND_ARCH128(Uint* buffer, int num_values, int null_count, + const uint8_t* valid_bits, + int64_t valid_bits_offset); + +extern template ARROW_TEMPLATE_EXPORT void SPACED_EXPAND_ARCH128( + uint32_t* buffer, int num_values, int null_count, const uint8_t* valid_bits, + int64_t valid_bits_offset); + +extern template ARROW_TEMPLATE_EXPORT void SPACED_EXPAND_ARCH128( + uint64_t* buffer, int num_values, int null_count, const uint8_t* valid_bits, + int64_t valid_bits_offset); + +#endif // SPACED_EXPAND_ARCH128 +#undef SPACED_EXPAND_ARCH128 + +} // namespace arrow::util::internal diff --git a/cpp/src/arrow/util/spaced_test.cc b/cpp/src/arrow/util/spaced_test.cc new file mode 100644 index 000000000000..a57bff32fde8 --- /dev/null +++ b/cpp/src/arrow/util/spaced_test.cc @@ -0,0 +1,209 @@ +// 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 +#include +#include + +#include + +#include "arrow/testing/gtest_util.h" +#include "arrow/util/bit_util.h" +#include "arrow/util/spaced_internal.h" +#include "arrow/util/spaced_simd_internal.h" + +#if defined(ARROW_HAVE_SSE4_2) || defined(ARROW_HAVE_RUNTIME_SSE4_2) +# include "arrow/util/cpu_info.h" +#endif + +namespace arrow::util::internal { + +namespace { + +struct Pattern { + enum Kind { kBernoulli, kRuns, kTopNulls }; + + std::string name; + Kind kind; + // kBernoulli: probability of a valid slot; kTopNulls: number of nulls + double p = 0; + // kRuns: mean lengths of the runs of valid and null slots + int mean_valid = 0; + int mean_null = 0; +}; + +std::vector AllPatterns() { + std::vector patterns; + for (double p : {0.0, 0.02, 0.1, 0.5, 0.9, 0.95, 0.97, 0.99, 0.995, 0.999, 1.0}) { + patterns.push_back({"p" + std::to_string(p), Pattern::kBernoulli, p}); + } + for (auto [valid, null] : std::vector>{{2, 2}, + {8, 8}, + {32, 32}, + {256, 256}, + {64, 2}, + {128, 1}, + {128, 8}, + {1000, 1}}) { + patterns.push_back({"runs_v" + std::to_string(valid) + "_n" + std::to_string(null), + Pattern::kRuns, 0, valid, null}); + } + // All valid except a few nulls in the last word: exercises the chunk + // kernel next to the end of the buffer. + for (int k : {1, 2, 3, 5, 7, 9, 12}) { + patterns.push_back({"top" + std::to_string(k), Pattern::kTopNulls, double(k)}); + } + return patterns; +} + +// Random bits outside of [offset, offset + n) so that the offset matters. +std::vector MakeBitmap(const Pattern& pattern, int n, int64_t offset, + std::mt19937_64& rng) { + std::vector bitmap(std::max(1, bit_util::BytesForBits(offset + n))); + for (auto& byte : bitmap) byte = static_cast(rng()); + auto set = [&](int i, bool valid) { + bit_util::SetBitTo(bitmap.data(), offset + i, valid); + }; + switch (pattern.kind) { + case Pattern::kBernoulli: { + std::bernoulli_distribution valid(pattern.p); + for (int i = 0; i < n; ++i) set(i, valid(rng)); + break; + } + case Pattern::kRuns: { + std::geometric_distribution valid_len(1.0 / pattern.mean_valid); + std::geometric_distribution null_len(1.0 / pattern.mean_null); + bool valid = true; + for (int i = 0; i < n;) { + const int len = 1 + (valid ? valid_len(rng) : null_len(rng)); + for (int end = std::min(n, i + len); i < end; ++i) set(i, valid); + valid = !valid; + } + break; + } + case Pattern::kTopNulls: { + for (int i = 0; i < n; ++i) set(i, true); + const int top = std::min(n, 64); + for (int k = 0; k < static_cast(pattern.p) && top > 0; ++k) { + set(n - 1 - static_cast(rng() % top), false); + } + break; + } + } + return bitmap; +} + +template +T MakeValue(int k) { + // Never 0, so that a value landing in a null slot is caught. + return static_cast(k % 30011 + 1); +} + +template +bool BitwiseEqual(T a, T b) { + return std::memcmp(&a, &b, sizeof(T)) == 0; +} + +// Check `expand` against a straightforward implementation. Null slots must +// be zero if `zeros_nulls`, and are not checked otherwise. +template +void CheckExpand(Expand&& expand, bool zeros_nulls) { + constexpr int kGuard = 16; + std::mt19937_64 rng(42); + for (int n : {0, 1, 7, 63, 64, 65, 127, 128, 129, 900, 2048}) { + for (int64_t offset : {0, 1, 5, 63}) { + for (const auto& pattern : AllPatterns()) { + ARROW_SCOPED_TRACE("n = ", n, ", offset = ", offset, + ", pattern = ", pattern.name); + const auto bitmap = MakeBitmap(pattern, n, offset, rng); + + std::vector expected(n); + int num_valid = 0; + for (int i = 0; i < n; ++i) { + if (bit_util::GetBit(bitmap.data(), offset + i)) { + expected[i] = MakeValue(num_valid++); + } + } + const int null_count = n - num_valid; + + // Leftover bytes in the null slots and after the buffer + std::vector buffer(n + kGuard); + std::memset(static_cast(buffer.data()), 0xA5, sizeof(T) * buffer.size()); + const std::vector guard(buffer.end() - kGuard, buffer.end()); + for (int k = 0; k < num_valid; ++k) buffer[k] = MakeValue(k); + + expand(buffer.data(), n, null_count, bitmap.data(), offset); + + for (int i = 0; i < n; ++i) { + const bool valid = bit_util::GetBit(bitmap.data(), offset + i); + if (valid || zeros_nulls) { + ASSERT_TRUE(BitwiseEqual(buffer[i], expected[i])) << "at index " << i; + } + } + for (int i = 0; i < kGuard; ++i) { + ASSERT_TRUE(BitwiseEqual(buffer[n + i], guard[i])) << "write past the end"; + } + } + } + } +} + +template +class TestSpacedExpandRightward : public ::testing::Test {}; + +using UintTypes = ::testing::Types; +TYPED_TEST_SUITE(TestSpacedExpandRightward, UintTypes); + +TYPED_TEST(TestSpacedExpandRightward, Scalar) { + CheckExpand(&SpacedExpandRightwardScalar, /*zeros_nulls=*/false); +} + +#if defined(ARROW_HAVE_NEON) +TYPED_TEST(TestSpacedExpandRightward, Neon) { + CheckExpand(&SpacedExpandRightwardNeon, /*zeros_nulls=*/true); +} +#endif + +#if defined(ARROW_HAVE_SSE4_2) || defined(ARROW_HAVE_RUNTIME_SSE4_2) +TYPED_TEST(TestSpacedExpandRightward, Sse4_2) { + if (!::arrow::internal::CpuInfo::GetInstance()->IsSupported( + ::arrow::internal::CpuInfo::SSE4_2)) { + GTEST_SKIP() << "Test requires SSE4.2"; + } + CheckExpand(&SpacedExpandRightwardSse4_2, /*zeros_nulls=*/true); +} +#endif + +// Through the public template: 4 and 8 byte types are dispatched (possibly +// to the scalar implementation), other widths use the scalar implementation. +template +class TestSpacedExpandRightwardTyped : public ::testing::Test {}; + +using ValueTypes = ::testing::Types; +TYPED_TEST_SUITE(TestSpacedExpandRightwardTyped, ValueTypes); + +TYPED_TEST(TestSpacedExpandRightwardTyped, Dispatch) { + CheckExpand(&SpacedExpandRightward, /*zeros_nulls=*/false); +} + +} // namespace + +} // namespace arrow::util::internal From 1bc728efb227b0c998e4a42c02c50234457a7faa Mon Sep 17 00:00:00 2001 From: mwish Date: Wed, 11 Nov 2026 02:02:35 +0800 Subject: [PATCH 2/2] POC: Sync spaced expansion kernels and direct PLAIN decoding --- cpp/src/arrow/CMakeLists.txt | 3 + cpp/src/arrow/util/spaced.cc | 113 ++++- cpp/src/arrow/util/spaced_benchmark.cc | 139 +++++-- cpp/src/arrow/util/spaced_internal.h | 71 +++- cpp/src/arrow/util/spaced_simd_128.cc | 299 ++------------ cpp/src/arrow/util/spaced_simd_256.cc | 134 ++++++ cpp/src/arrow/util/spaced_simd_avx512.cc | 97 +++++ cpp/src/arrow/util/spaced_simd_internal.h | 54 ++- .../arrow/util/spaced_simd_kernel_internal.h | 386 ++++++++++++++++++ cpp/src/arrow/util/spaced_test.cc | 85 +++- cpp/src/parquet/decoder.cc | 36 ++ 11 files changed, 1068 insertions(+), 349 deletions(-) create mode 100644 cpp/src/arrow/util/spaced_simd_256.cc create mode 100644 cpp/src/arrow/util/spaced_simd_avx512.cc create mode 100644 cpp/src/arrow/util/spaced_simd_kernel_internal.h diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt index 1fb5cfaedaf3..0a896a82263a 100644 --- a/cpp/src/arrow/CMakeLists.txt +++ b/cpp/src/arrow/CMakeLists.txt @@ -595,6 +595,9 @@ append_runtime_sse4_2_src(ARROW_UTIL_SRCS util/byte_stream_split_internal_sse4_2 append_runtime_avx2_src(ARROW_UTIL_SRCS util/bpacking_simd_256.cc) append_runtime_avx512_src(ARROW_UTIL_SRCS util/bpacking_simd_avx512.cc) +append_runtime_avx2_src(ARROW_UTIL_SRCS util/spaced_simd_256.cc) +append_runtime_avx512_src(ARROW_UTIL_SRCS util/spaced_simd_avx512.cc) + # also provides the NEON kernels on aarch64, so it stays in ARROW_UTIL_SRCS if(ARROW_CPU_FLAG STREQUAL "x86" AND ARROW_HAVE_RUNTIME_SSE4_2) set_source_files_properties(util/bpacking_simd_128.cc PROPERTIES COMPILE_OPTIONS diff --git a/cpp/src/arrow/util/spaced.cc b/cpp/src/arrow/util/spaced.cc index 90c713b4ace4..9da6b50cd322 100644 --- a/cpp/src/arrow/util/spaced.cc +++ b/cpp/src/arrow/util/spaced.cc @@ -20,20 +20,107 @@ #include "arrow/util/dispatch_internal.h" #include "arrow/util/spaced_internal.h" #include "arrow/util/spaced_simd_internal.h" +#include "arrow/util/spaced_simd_kernel_internal.h" +#include "arrow/util/ubsan.h" namespace arrow::util::internal { namespace { +// For an 8 bit chunk mask, row[mask][j] is the number of valid lanes below +// lane j if lane j is valid, 0 otherwise. +struct SourceIndexTable { + uint8_t row[256][8] = {}; + constexpr SourceIndexTable() { + for (unsigned m = 0; m < 256; ++m) { + for (int j = 0, prefix = 0; j < 8; ++j) { + if ((m >> j) & 1) row[m][j] = static_cast(prefix++); + } + } + } +}; +constexpr SourceIndexTable kSourceIndex{}; + +// Chunk kernel (see spaced::ChunkedWordKernel) on 8 values without SIMD. +// Every lane finds its source through the table and is cleared with a mask, +// so that the lanes are independent and there is no branch. +template +struct ExpandChunkScalar { + static constexpr int kLanes = 8; + + static void Expand(Uint* dst, const Uint* src, unsigned mask) { + const uint8_t* index = kSourceIndex.row[mask]; + Uint values[kLanes]; + for (int j = 0; j < kLanes; ++j) { + values[j] = SafeLoad(src + index[j]); + } + for (int j = 0; j < kLanes; ++j) { + dst[j] = values[j] & (Uint{0} - static_cast((mask >> j) & 1)); + } + } +}; + +// Words with at most that many valid values are scattered one by one. +constexpr int kSparseValues = 8; + +template +using Kernel = spaced::ChunkedWordKernel>; + +} // namespace + +template +void SpacedExpandRightwardNone(Uint* buffer, int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset) { + constexpr int kDenseNulls = 4; + spaced::InPlaceExpander, kDenseNulls, kSparseValues>::Expand( + buffer, num_values, null_count, valid_bits, valid_bits_offset); +} + +template +void SpacedExpandNone(const Uint* src, int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset, + Uint* output) { + spaced::OutOfPlaceExpander, kSparseValues>::Expand( + src, num_values, null_count, valid_bits, valid_bits_offset, output); +} + +template void SpacedExpandRightwardNone(uint32_t*, int, int, const uint8_t*, + int64_t); +template void SpacedExpandRightwardNone(uint64_t*, int, int, const uint8_t*, + int64_t); +template void SpacedExpandNone(const uint32_t*, int, int, const uint8_t*, + int64_t, uint32_t*); +template void SpacedExpandNone(const uint64_t*, int, int, const uint8_t*, + int64_t, uint64_t*); + +namespace { + template struct SpacedExpandRightwardDynamicFunction { - using FunctionType = decltype(&SpacedExpandRightwardScalar); + using FunctionType = decltype(&SpacedExpandRightwardNone); static constexpr auto targets() { return std::array{ - ARROW_DISPATCH_TARGET_NONE(&SpacedExpandRightwardScalar) // + ARROW_DISPATCH_TARGET_NONE(&SpacedExpandRightwardNone) // ARROW_DISPATCH_TARGET_NEON(&SpacedExpandRightwardNeon) // ARROW_DISPATCH_TARGET_SSE4_2(&SpacedExpandRightwardSse4_2) // + ARROW_DISPATCH_TARGET_AVX2(&SpacedExpandRightwardAvx2) // + ARROW_DISPATCH_TARGET_AVX512(&SpacedExpandRightwardAvx512) // + }; + } +}; + +template +struct SpacedExpandDynamicFunction { + using FunctionType = decltype(&SpacedExpandNone); + + static constexpr auto targets() { + return std::array{ + ARROW_DISPATCH_TARGET_NONE(&SpacedExpandNone) // + ARROW_DISPATCH_TARGET_NEON(&SpacedExpandNeon) // + ARROW_DISPATCH_TARGET_SSE4_2(&SpacedExpandSse4_2) // + ARROW_DISPATCH_TARGET_AVX2(&SpacedExpandAvx2) // + ARROW_DISPATCH_TARGET_AVX512(&SpacedExpandAvx512) // }; } }; @@ -48,6 +135,16 @@ void SpacedExpandRightwardDispatch(void* buffer, int num_values, int null_count, valid_bits_offset); } +template +void SpacedExpandDispatch(const void* src, int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset, + void* output) { + static const ::arrow::internal::DynamicDispatch> + dispatch; + dispatch(static_cast(src), num_values, null_count, valid_bits, + valid_bits_offset, static_cast(output)); +} + } // namespace void SpacedExpandRightward32(void* buffer, int num_values, int null_count, @@ -62,4 +159,16 @@ void SpacedExpandRightward64(void* buffer, int num_values, int null_count, valid_bits_offset); } +void SpacedExpand32(const void* src, int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset, void* output) { + SpacedExpandDispatch(src, num_values, null_count, valid_bits, + valid_bits_offset, output); +} + +void SpacedExpand64(const void* src, int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset, void* output) { + SpacedExpandDispatch(src, num_values, null_count, valid_bits, + valid_bits_offset, output); +} + } // namespace arrow::util::internal diff --git a/cpp/src/arrow/util/spaced_benchmark.cc b/cpp/src/arrow/util/spaced_benchmark.cc index 24a3c8706d11..0d077741687c 100644 --- a/cpp/src/arrow/util/spaced_benchmark.cc +++ b/cpp/src/arrow/util/spaced_benchmark.cc @@ -24,6 +24,7 @@ #include #include "arrow/util/bit_util.h" +#include "arrow/util/cpu_info.h" #include "arrow/util/spaced_internal.h" #include "arrow/util/spaced_simd_internal.h" @@ -31,7 +32,10 @@ namespace arrow::util::internal { namespace { template -using ExpandFunc = void (*)(Uint*, int, int, const uint8_t*, int64_t); +using ExpandRightwardFunc = void (*)(Uint*, int, int, const uint8_t*, int64_t); + +template +using ExpandFunc = void (*)(const Uint*, int, int, const uint8_t*, int64_t, Uint*); // Either independent slots valid with probability `p`, or (p < 0) runs of // valid and null slots with the given mean lengths. @@ -67,58 +71,115 @@ std::vector MakeBitmap(const ValidityPattern& pattern, int n) { return bitmap; } -// The timing includes copying the dense values back in place before each -// expansion. +constexpr int kNumValues = 2048; + template -void BM_SpacedExpandRightward(benchmark::State& state, ExpandFunc expand, - ValidityPattern pattern) { - constexpr int kNumValues = 2048; - const auto bitmap = MakeBitmap(pattern, kNumValues); - int num_valid = 0; - for (int i = 0; i < kNumValues; ++i) { - num_valid += bit_util::GetBit(bitmap.data(), i); +struct ExpandInput { + explicit ExpandInput(ValidityPattern pattern) + : bitmap(MakeBitmap(pattern, kNumValues)), storage(2 * kNumValues + kGap) { + int num_valid = 0; + for (int i = 0; i < kNumValues; ++i) { + num_valid += bit_util::GetBit(bitmap.data(), i); + } + null_count = kNumValues - num_valid; + for (int i = 0; i < kNumValues; ++i) dense()[i] = static_cast(i + 1); } - const int null_count = kNumValues - num_valid; - std::vector dense(kNumValues); - for (int i = 0; i < kNumValues; ++i) dense[i] = static_cast(i + 1); - std::vector buffer(kNumValues); + // The dense values and the output in a single allocation, 1 KiB apart modulo + // 4 KiB: the copies are much slower when the source and the destination are + // a few bytes apart modulo 4 KiB (4K aliasing), which separate allocations + // of the same size can easily be. + static constexpr int kGap = 1024 / sizeof(Uint); + const Uint* dense() const { return storage.data(); } + Uint* dense() { return storage.data(); } + Uint* output() { return storage.data() + kNumValues + kGap; } + + std::vector bitmap; + std::vector storage; + int null_count; +}; + +bool SkipUnsupported(benchmark::State& state, int64_t cpu_flags) { + if (::arrow::internal::CpuInfo::GetInstance()->IsSupported(cpu_flags)) return false; + state.SkipWithMessage("CPU not supported"); + return true; +} +// In place, as after decoding the dense values: the timing includes copying +// them to the buffer before each expansion. +template +void BM_SpacedExpandRightward(benchmark::State& state, ExpandRightwardFunc expand, + int64_t cpu_flags, ValidityPattern pattern) { + if (SkipUnsupported(state, cpu_flags)) return; + ExpandInput input(pattern); for (auto _ : state) { - std::memcpy(buffer.data(), dense.data(), sizeof(Uint) * num_valid); - expand(buffer.data(), kNumValues, null_count, bitmap.data(), 0); + std::memcpy(input.output(), input.dense(), + sizeof(Uint) * (kNumValues - input.null_count)); + expand(input.output(), kNumValues, input.null_count, input.bitmap.data(), 0); benchmark::ClobberMemory(); } state.SetItemsProcessed(state.iterations() * kNumValues); } -#define BENCHMARK_SPACED_EXPAND_PATTERNS(NAME, FUNC) \ - BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_p0_02, FUNC, Bernoulli(0.02)); \ - BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_p0_10, FUNC, Bernoulli(0.10)); \ - BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_p0_50, FUNC, Bernoulli(0.50)); \ - BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_p0_90, FUNC, Bernoulli(0.90)); \ - BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_p0_95, FUNC, Bernoulli(0.95)); \ - BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_p0_99, FUNC, Bernoulli(0.99)); \ - BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_runs_v2_n2, FUNC, Runs(2, 2)); \ - BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_runs_v8_n8, FUNC, Runs(8, 8)); \ - BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_runs_v32_n32, FUNC, Runs(32, 32)); \ - BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_runs_v64_n2, FUNC, Runs(64, 2)); \ - BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_runs_v128_n8, FUNC, Runs(128, 8)); \ - BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_runs_v256_n256, FUNC, \ - Runs(256, 256)); \ - BENCHMARK_CAPTURE(BM_SpacedExpandRightward, NAME##_runs_v1000_n1, FUNC, Runs(1000, 1)) - -BENCHMARK_SPACED_EXPAND_PATTERNS(Scalar32, &SpacedExpandRightwardScalar); -BENCHMARK_SPACED_EXPAND_PATTERNS(Scalar64, &SpacedExpandRightwardScalar); +// Out of place, straight from the dense values. +template +void BM_SpacedExpand(benchmark::State& state, ExpandFunc expand, int64_t cpu_flags, + ValidityPattern pattern) { + if (SkipUnsupported(state, cpu_flags)) return; + ExpandInput input(pattern); + for (auto _ : state) { + expand(input.dense(), kNumValues, input.null_count, input.bitmap.data(), 0, + input.output()); + benchmark::ClobberMemory(); + } + state.SetItemsProcessed(state.iterations() * kNumValues); +} + +#define BENCHMARK_SPACED_EXPAND_PATTERNS(BM, NAME, FUNC, FLAGS) \ + BENCHMARK_CAPTURE(BM, NAME##_p0_02, FUNC, FLAGS, Bernoulli(0.02)); \ + BENCHMARK_CAPTURE(BM, NAME##_p0_10, FUNC, FLAGS, Bernoulli(0.10)); \ + BENCHMARK_CAPTURE(BM, NAME##_p0_50, FUNC, FLAGS, Bernoulli(0.50)); \ + BENCHMARK_CAPTURE(BM, NAME##_p0_90, FUNC, FLAGS, Bernoulli(0.90)); \ + BENCHMARK_CAPTURE(BM, NAME##_p0_95, FUNC, FLAGS, Bernoulli(0.95)); \ + BENCHMARK_CAPTURE(BM, NAME##_p0_99, FUNC, FLAGS, Bernoulli(0.99)); \ + BENCHMARK_CAPTURE(BM, NAME##_runs_v2_n2, FUNC, FLAGS, Runs(2, 2)); \ + BENCHMARK_CAPTURE(BM, NAME##_runs_v8_n8, FUNC, FLAGS, Runs(8, 8)); \ + BENCHMARK_CAPTURE(BM, NAME##_runs_v32_n32, FUNC, FLAGS, Runs(32, 32)); \ + BENCHMARK_CAPTURE(BM, NAME##_runs_v64_n2, FUNC, FLAGS, Runs(64, 2)); \ + BENCHMARK_CAPTURE(BM, NAME##_runs_v128_n8, FUNC, FLAGS, Runs(128, 8)); \ + BENCHMARK_CAPTURE(BM, NAME##_runs_v256_n256, FUNC, FLAGS, Runs(256, 256)); \ + BENCHMARK_CAPTURE(BM, NAME##_runs_v1000_n1, FUNC, FLAGS, Runs(1000, 1)) + +// Both directions, both widths +#define BENCHMARK_SPACED_EXPAND_TARGET(TARGET, FLAGS) \ + BENCHMARK_SPACED_EXPAND_PATTERNS(BM_SpacedExpandRightward, TARGET##_32, \ + &SpacedExpandRightward##TARGET, FLAGS); \ + BENCHMARK_SPACED_EXPAND_PATTERNS(BM_SpacedExpandRightward, TARGET##_64, \ + &SpacedExpandRightward##TARGET, FLAGS); \ + BENCHMARK_SPACED_EXPAND_PATTERNS(BM_SpacedExpand, TARGET##_32, \ + &SpacedExpand##TARGET, FLAGS); \ + BENCHMARK_SPACED_EXPAND_PATTERNS(BM_SpacedExpand, TARGET##_64, \ + &SpacedExpand##TARGET, FLAGS) + +// Scalar: run based, None: word driver without SIMD +BENCHMARK_SPACED_EXPAND_TARGET(Scalar, 0); +BENCHMARK_SPACED_EXPAND_TARGET(None, 0); #if defined(ARROW_HAVE_NEON) -BENCHMARK_SPACED_EXPAND_PATTERNS(Neon32, &SpacedExpandRightwardNeon); -BENCHMARK_SPACED_EXPAND_PATTERNS(Neon64, &SpacedExpandRightwardNeon); -#elif defined(ARROW_HAVE_SSE4_2) -BENCHMARK_SPACED_EXPAND_PATTERNS(Sse4_2_32, &SpacedExpandRightwardSse4_2); -BENCHMARK_SPACED_EXPAND_PATTERNS(Sse4_2_64, &SpacedExpandRightwardSse4_2); +BENCHMARK_SPACED_EXPAND_TARGET(Neon, 0); +#elif defined(ARROW_HAVE_SSE4_2) || defined(ARROW_HAVE_RUNTIME_SSE4_2) +BENCHMARK_SPACED_EXPAND_TARGET(Sse4_2, ::arrow::internal::CpuInfo::SSE4_2); +#endif + +#if defined(ARROW_HAVE_AVX2) || defined(ARROW_HAVE_RUNTIME_AVX2) +BENCHMARK_SPACED_EXPAND_TARGET(Avx2, ::arrow::internal::CpuInfo::AVX2); +#endif + +#if defined(ARROW_HAVE_AVX512) || defined(ARROW_HAVE_RUNTIME_AVX512) +BENCHMARK_SPACED_EXPAND_TARGET(Avx512, ::arrow::internal::CpuInfo::AVX512); #endif +#undef BENCHMARK_SPACED_EXPAND_TARGET #undef BENCHMARK_SPACED_EXPAND_PATTERNS } // namespace diff --git a/cpp/src/arrow/util/spaced_internal.h b/cpp/src/arrow/util/spaced_internal.h index f0f3c9b23911..2caf93714daa 100644 --- a/cpp/src/arrow/util/spaced_internal.h +++ b/cpp/src/arrow/util/spaced_internal.h @@ -109,11 +109,10 @@ ARROW_EXPORT void SpacedExpandRightward64(void* buffer, int num_values, int null /// /// Non-null values should initially be densely packed at the left of the buffer. /// This method spreads the values out according to the given validity bitmap. -/// Null entries are zero-initialized. /// -/// 4 and 8 byte values go through a SIMD implementation when the CPU has one, -/// which zeroes every null slot. Other widths, and CPUs without SIMD support, -/// use SpacedExpandRightwardScalar. +/// 4 and 8 byte values go through an implementation dispatched on the CPU, +/// which zeroes every null slot. Other widths use SpacedExpandRightwardScalar, +/// which only zeroes the trailing `null_count` slots. /// /// \param[in, out] buffer the in-place buffer /// \param[in] num_values total size of buffer including null slots @@ -135,6 +134,70 @@ inline void SpacedExpandRightward(T* buffer, int num_values, int null_count, } } +/// \brief Spread values out according to a validity bitmap +/// +/// Run based implementation of SpacedExpand. +template +inline void SpacedExpandScalar(const T* src, int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset, + T* output) { + // End of the last run: the slots between it and the next run are nulls. + int64_t nulls_begin = 0; + int64_t idx_src = 0; + arrow::internal::SetBitRunReader reader(valid_bits, valid_bits_offset, num_values); + while (true) { + const auto run = reader.NextRun(); + if (run.length == 0) { + break; + } + memset(static_cast(output + nulls_begin), 0, + (run.position - nulls_begin) * sizeof(T)); + memcpy(static_cast(output + run.position), src + idx_src, + run.length * sizeof(T)); + idx_src += run.length; + nulls_begin = run.position + run.length; + } + assert(idx_src == num_values - null_count); + memset(static_cast(output + nulls_begin), 0, + (num_values - nulls_begin) * sizeof(T)); +} + +/// \brief SpacedExpand for 4 byte values, dispatched on the CPU +ARROW_EXPORT void SpacedExpand32(const void* src, int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset, + void* output); + +/// \brief SpacedExpand for 8 byte values, dispatched on the CPU +ARROW_EXPORT void SpacedExpand64(const void* src, int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset, + void* output); + +/// \brief Spread values out according to a validity bitmap +/// +/// Out of place version of SpacedExpandRightward: the `num_values - null_count` +/// values of `src` go to the valid slots of `output`, in order, and every null +/// slot is zeroed. `src` needs not be aligned, and must not overlap `output`. +/// +/// \param[in] src the dense values +/// \param[in] num_values total size of output including null slots +/// \param[in] null_count number of null slots +/// \param[in] valid_bits bitmap data indicating position of valid slots +/// \param[in] valid_bits_offset offset into valid_bits +/// \param[out] output the spaced values +template +inline void SpacedExpand(const T* src, int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset, + T* output) { + if constexpr (std::is_trivially_copyable_v && sizeof(T) == 4) { + SpacedExpand32(src, num_values, null_count, valid_bits, valid_bits_offset, output); + } else if constexpr (std::is_trivially_copyable_v && sizeof(T) == 8) { + SpacedExpand64(src, num_values, null_count, valid_bits, valid_bits_offset, output); + } else { + SpacedExpandScalar(src, num_values, null_count, valid_bits, valid_bits_offset, + output); + } +} + /// \brief Relocate values according to a validity bitmap, to the left /// /// Non-null values should initially be densely packed at the right of the buffer. diff --git a/cpp/src/arrow/util/spaced_simd_128.cc b/cpp/src/arrow/util/spaced_simd_128.cc index cf6c1262c840..8fd268875301 100644 --- a/cpp/src/arrow/util/spaced_simd_128.cc +++ b/cpp/src/arrow/util/spaced_simd_128.cc @@ -16,25 +16,21 @@ // under the License. #if defined(ARROW_HAVE_NEON) -# define SPACED_EXPAND_PLATFORM SpacedExpandRightwardNeon +# define SPACED_EXPAND_RIGHTWARD_PLATFORM SpacedExpandRightwardNeon +# define SPACED_EXPAND_PLATFORM SpacedExpandNeon #elif defined(ARROW_HAVE_SSE4_2) || defined(ARROW_HAVE_RUNTIME_SSE4_2) -# define SPACED_EXPAND_PLATFORM SpacedExpandRightwardSse4_2 +# define SPACED_EXPAND_RIGHTWARD_PLATFORM SpacedExpandRightwardSse4_2 +# define SPACED_EXPAND_PLATFORM SpacedExpandSse4_2 #endif #if defined(SPACED_EXPAND_PLATFORM) -# include # include -# include -# include # include -# include "arrow/util/bit_util.h" -# include "arrow/util/endian.h" -# include "arrow/util/macros.h" # include "arrow/util/spaced_simd_internal.h" -# include "arrow/util/ubsan.h" +# include "arrow/util/spaced_simd_kernel_internal.h" namespace arrow::util::internal { @@ -42,43 +38,6 @@ namespace { using Arch = xsimd::default_arch; -// --------------------------------------------------------------------------- -// Lookup tables -// -// AArch64 without FEAT_CSSC has no scalar popcount: std::popcount is a round -// trip through a SIMD register (fmov/cnt/addv/fmov), so the hot paths look -// small popcounts up instead. -// --------------------------------------------------------------------------- - -struct PopCountTable { - uint8_t count[256] = {}; - constexpr PopCountTable() { - for (unsigned m = 0; m < 256; ++m) { - count[m] = static_cast(std::popcount(m)); - } - } -}; -inline constexpr PopCountTable kPopCount{}; - -// For an 8 bit chunk mask split in sub-batches of kSubLanes lanes, byte j of -// row[mask] is the number of valid lanes in sub-batches 0..j. -template -struct PrefixTable { - uint32_t row[256] = {}; - constexpr PrefixTable() { - for (unsigned m = 0; m < 256; ++m) { - uint32_t r = 0; - for (int j = 0; j < 4 && (j + 1) * kSubLanes <= 8; ++j) { - const unsigned low = (1u << ((j + 1) * kSubLanes)) - 1u; - r |= static_cast(std::popcount(m & low)) << (8 * j); - } - row[m] = r; - } - } -}; -template -inline constexpr PrefixTable kPrefix{}; - // Byte shuffle expanding the low lanes of a 16 byte register according to a // kLanes bit mask. Null lanes read byte 0 so that every index stays in // [0, 16), the only range where the semantics of a dynamic swizzle are @@ -109,15 +68,9 @@ struct ExpandTable { template inline constexpr ExpandTable kExpand{}; -// --------------------------------------------------------------------------- -// Chunk kernel: expand the low popcount(mask) values of `src` to the 8 slots -// of `dst` according to `mask`, zeroing the null slots. -// -// Reads at most 8 values from `src`: sub-batch j reads kSubLanes values at -// the number of valid lanes below it, which is at most 8 - kSubLanes. -// Everything is loaded before anything is stored, so `dst` may overlap `src`. -// --------------------------------------------------------------------------- - +// Chunk kernel (see spaced::ChunkedWordKernel) on 8 values: sub-batch j reads +// kSubLanes values at the number of valid lanes below it, which is at most +// 8 - kSubLanes. template struct ExpandChunk { using Batch = xsimd::batch; @@ -129,7 +82,7 @@ struct ExpandChunk { static constexpr int kSubs = kLanes / kSubLanes; static void Expand(Uint* dst, const Uint* src, unsigned mask) { - const uint32_t prefixes = kPrefix.row[mask]; + const uint32_t prefixes = spaced::kPrefix.row[mask]; Batch sub[kSubs]; unsigned prefix = 0; for (int j = 0; j < kSubs; ++j) { @@ -148,225 +101,43 @@ struct ExpandChunk { } }; -// --------------------------------------------------------------------------- -// Word driver -// -// The bitmap is walked one 64 bit word at a time from the end of the buffer, -// and every word goes to one of three paths: -// * runs (at most kDenseNulls nulls, or a single valid / null transition): -// memmove whole runs of valid values and memset runs of nulls, keeping a -// run open across words so that a long run is moved only once; -// * sparse (few valid values, far from their destination): zero the word -// and scatter the values one by one; -// * otherwise: 8 chunks of the SIMD kernel. -// --------------------------------------------------------------------------- - -// Bits [pos, pos + 64) of the bitmap. Only reads the bytes holding those bits. -inline uint64_t LoadWord(const uint8_t* bitmap, int64_t pos) { - const uint8_t* p = bitmap + (pos >> 3); - const int shift = static_cast(pos & 7); - const uint64_t lo = bit_util::FromLittleEndian(SafeLoadAs(p)); - if (shift == 0) return lo; - return (lo >> shift) | (static_cast(p[8]) << (64 - shift)); -} - -// Whether the word goes to the run path: at most kDenseNulls nulls, or a -// single transition between a valid run and a null run (including empty -// words). Neither needs a popcount. -template -inline bool IsRunWord(uint64_t bits) { - uint64_t z = ~bits; - for (int k = 0; k < kDenseNulls; ++k) z &= z - 1; - return z == 0 || (bits & (bits + 1)) == 0 || (~bits & (~bits + 1)) == 0; -} +// Words with at most that many valid values are scattered one by one. +constexpr int kSparseValues = 8; -// Handle run words downwards from the one ending at `pos`, whose bits are -// `bits`. `nulls_below` is the number of nulls below `pos` on entry, and -// below the returned position on return. Kept out of line: having the open -// run state in the main loop slows every path down. -template -ARROW_NOINLINE int ExpandRunWords(Uint* out, int pos, int& nulls_below, uint64_t bits, - const uint8_t* valid_bits, int64_t valid_bits_offset) { - // The open run: values [run_hi - x, run_hi) come from run_shift slots lower. - bool pending = false; - int run_hi = 0; - int run_shift = 0; - auto flush = [&](int run_lo) { - if (pending && run_hi > run_lo) { - std::memmove(out + run_lo, out + run_lo - run_shift, - sizeof(Uint) * static_cast(run_hi - run_lo)); - } - pending = false; - }; - while (true) { - const int lo = pos - 64; - if (bits == 0) { - // Empty words close the open run and are zeroed all at once. - flush(pos); - const int hi = pos; - pos = lo; - while (pos >= 64 && LoadWord(valid_bits, valid_bits_offset + pos - 64) == 0) { - pos -= 64; - } - std::memset(out + pos, 0, sizeof(Uint) * static_cast(hi - pos)); - nulls_below -= hi - pos; - } else { - if (!pending && (bits >> 63) != 0) { - pending = true; - run_hi = pos; - run_shift = nulls_below; - } - // One step per run of nulls [null_lo, null_hi], from the top. Zeroing - // them is safe once the run above is moved: everything below still - // comes from lower slots. - uint64_t z = ~bits; - while (z != 0) { - const int null_hi = 63 - std::countl_zero(z); - const uint64_t valid_below = bits & ((uint64_t{1} << null_hi) - 1); - const int null_lo = valid_below == 0 ? 0 : 64 - std::countl_zero(valid_below); - const int len = null_hi - null_lo + 1; - nulls_below -= len; - flush(lo + null_hi + 1); - if (len == 1) { - out[lo + null_lo] = 0; - } else { - std::memset(out + lo + null_lo, 0, sizeof(Uint) * static_cast(len)); - } - z &= (uint64_t{1} << null_lo) - 1; - if (null_lo > 0) { - pending = true; - run_hi = lo + null_lo; - run_shift = nulls_below; - } - } - pos = lo; - } - // No nulls below: the open run (shift 0) and everything below it is - // already in place. - if (nulls_below == 0) return pos; - if (pos < 64) break; - bits = LoadWord(valid_bits, valid_bits_offset + pos - 64); - if (bits == ~uint64_t{0}) { - // Full words only extend the open run: skip them two at a time. - if (!pending) { - pending = true; - run_hi = pos; - run_shift = nulls_below; - } - pos -= 64; - while (pos >= 128) { - const uint64_t a = LoadWord(valid_bits, valid_bits_offset + pos - 64); - const uint64_t b = LoadWord(valid_bits, valid_bits_offset + pos - 128); - if ((a & b) != ~uint64_t{0}) break; - pos -= 128; - } - if (pos < 64) break; - bits = LoadWord(valid_bits, valid_bits_offset + pos - 64); - } - if (!IsRunWord(bits)) break; - } - flush(pos); - return pos; -} - -template -void ExpandRightward(Uint* out, int num_values, int null_count, const uint8_t* valid_bits, - int64_t valid_bits_offset) { - using Kernel = ExpandChunk; - constexpr int kLanes = Kernel::kLanes; - constexpr int kChunks = 64 / kLanes; - // Words with at most that many valid values, far enough from their - // destination, are scattered one by one. - constexpr int kSparseValues = 8; - - if (null_count == 0) return; - int idx_decode = num_values - null_count; - if (idx_decode == 0) { - std::memset(out, 0, sizeof(Uint) * static_cast(num_values)); - return; - } - int pos = num_values; - while (pos >= 64) { - const int lo = pos - 64; - const uint64_t bits = LoadWord(valid_bits, valid_bits_offset + lo); - if (IsRunWord(bits)) { - int nulls_below = pos - idx_decode; - pos = ExpandRunWords(out, pos, nulls_below, bits, valid_bits, - valid_bits_offset); - if (nulls_below == 0) return; - idx_decode = pos - nulls_below; - continue; - } - const int pc = std::popcount(bits); - idx_decode -= pc; - const int src_begin = idx_decode; - const int shift = lo - src_begin; - if (pc <= kSparseValues && shift >= 64) { - // Source and destination do not overlap. - std::memset(out + lo, 0, 64 * sizeof(Uint)); - int q = src_begin + pc; - uint64_t b = bits; - while (b != 0) { - const int i = 63 - std::countl_zero(b); - b ^= uint64_t{1} << i; - out[lo + i] = out[--q]; - } - } else { - // The kernel reads up to kLanes values from its source, which can go - // past the valid values of the top word only. Stage those in a buffer. - Uint staged[64 + kLanes]; - const Uint* base = out; - int rem = src_begin + pc; - if (rem + kLanes > num_values) { - std::memcpy(staged, out + src_begin, sizeof(Uint) * static_cast(pc)); - base = staged; - rem = pc; - } - auto expand_chunk = [&](int c) { - const unsigned mask = static_cast((bits >> (c * kLanes)) & 0xFF); - const int pc_c = kPopCount.count[mask]; - if (pc_c == 0) { - std::memset(out + lo + c * kLanes, 0, sizeof(Uint) * kLanes); - return; - } - const int src = rem - pc_c; - Kernel::Expand(out + lo + c * kLanes, base + src, mask); - rem = src; - }; - // Unrolled by hand: Arrow builds with -O2, which keeps the loop. - [&](std::integer_sequence) { - (expand_chunk(kChunks - 1 - C), ...); - }(std::make_integer_sequence{}); - } - pos = lo; - if (shift == 0) return; - } - for (int i = pos - 1; i >= 0; --i) { - if (bit_util::GetBit(valid_bits, valid_bits_offset + i)) { - out[i] = out[--idx_decode]; - } else { - out[i] = 0; - } - } -} +template +using Kernel = spaced::ChunkedWordKernel>; } // namespace template -void SPACED_EXPAND_PLATFORM(Uint* buffer, int num_values, int null_count, - const uint8_t* valid_bits, int64_t valid_bits_offset) { +void SPACED_EXPAND_RIGHTWARD_PLATFORM(Uint* buffer, int num_values, int null_count, + const uint8_t* valid_bits, + int64_t valid_bits_offset) { // Tuned on Apple M2: words with more nulls are faster through the kernel. constexpr int kDenseNulls = sizeof(Uint) == 4 ? 1 : 2; - ExpandRightward(buffer, num_values, null_count, valid_bits, - valid_bits_offset); + spaced::InPlaceExpander, kDenseNulls, kSparseValues>::Expand( + buffer, num_values, null_count, valid_bits, valid_bits_offset); +} + +template +void SPACED_EXPAND_PLATFORM(const Uint* src, int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset, + Uint* output) { + spaced::OutOfPlaceExpander, kSparseValues>::Expand( + src, num_values, null_count, valid_bits, valid_bits_offset, output); } -template void SPACED_EXPAND_PLATFORM(uint32_t*, int, int, const uint8_t*, - int64_t); -template void SPACED_EXPAND_PLATFORM(uint64_t*, int, int, const uint8_t*, - int64_t); +template void SPACED_EXPAND_RIGHTWARD_PLATFORM(uint32_t*, int, int, + const uint8_t*, int64_t); +template void SPACED_EXPAND_RIGHTWARD_PLATFORM(uint64_t*, int, int, + const uint8_t*, int64_t); +template void SPACED_EXPAND_PLATFORM(const uint32_t*, int, int, const uint8_t*, + int64_t, uint32_t*); +template void SPACED_EXPAND_PLATFORM(const uint64_t*, int, int, const uint8_t*, + int64_t, uint64_t*); } // namespace arrow::util::internal +# undef SPACED_EXPAND_RIGHTWARD_PLATFORM # undef SPACED_EXPAND_PLATFORM #endif // SPACED_EXPAND_PLATFORM diff --git a/cpp/src/arrow/util/spaced_simd_256.cc b/cpp/src/arrow/util/spaced_simd_256.cc new file mode 100644 index 000000000000..f2d35724ba80 --- /dev/null +++ b/cpp/src/arrow/util/spaced_simd_256.cc @@ -0,0 +1,134 @@ +// 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. + +#if !defined(ARROW_HAVE_RUNTIME_AVX2) +# error "This file must be compiled with AVX2 enabled" +#endif + +#include + +#include +#include + +#include "arrow/util/spaced_simd_internal.h" +#include "arrow/util/spaced_simd_kernel_internal.h" + +namespace arrow::util::internal { + +namespace { + +using Arch = xsimd::avx2; + +// For a kSubLanes bit mask, the 8 nibbles of row[mask] give the 32 bit word +// each 32 bit word of the result is taken from, bit 3 meaning that it is kept +// (the value is valid). 1 KiB for 4 byte values, 64 bytes for 8 byte values. +template +struct PermuteTable { + uint32_t row[1 << kSubLanes] = {}; + constexpr PermuteTable() { + constexpr int kWords = 8 / kSubLanes; + for (int m = 0; m < (1 << kSubLanes); ++m) { + uint32_t r = 0; + int prefix = 0; + for (int lane = 0; lane < kSubLanes; ++lane) { + if (((m >> lane) & 1) == 0) continue; + for (int b = 0; b < kWords; ++b) { + const uint32_t nibble = 8u | static_cast(prefix * kWords + b); + r |= nibble << (4 * (lane * kWords + b)); + } + ++prefix; + } + row[m] = r; + } + } +}; +template +inline constexpr PermuteTable kPermute{}; + +// Chunk kernel (see spaced::ChunkedWordKernel) on 8 values: one 32 bit lane +// permutation per batch. +template +struct ExpandChunk { + using Batch = xsimd::batch; + using Words = xsimd::batch; + static_assert(Words::size == 8, "256 bit SIMD only"); + + static constexpr int kLanes = 8; + static constexpr int kSubLanes = static_cast(Batch::size); + static constexpr int kSubs = kLanes / kSubLanes; + + static void Expand(Uint* dst, const Uint* src, unsigned mask) { + const Words shifts{0, 4, 8, 12, 16, 20, 24, 28}; + const uint32_t prefixes = kSubs > 1 ? spaced::kPrefix.row[mask] : 0; + Batch sub[kSubs]; + unsigned prefix = 0; + for (int j = 0; j < kSubs; ++j) { + const unsigned m = (mask >> (j * kSubLanes)) & ((1u << kSubLanes) - 1u); + const Words nibbles = (Words(kPermute.row[m]) >> shifts) & Words(0xF); + const Words words = + xsimd::bitwise_cast(Batch::load_unaligned(src + prefix)); + // TODO: use xsimd::swizzle once xsimd emits vpermd for it + // (https://github.com/xtensor-stack/xsimd/pull/1446). Up to xsimd 14.3 it + // is emulated with AVX instructions and slower than the 128 bit kernel. + // vpermd only looks at the low 3 bits of the indices. + const Words permuted = _mm256_permutevar8x32_epi32(words, nibbles); + // All ones where bit 3 is set + const Words keep = Words(0) - (nibbles >> 3); + sub[j] = xsimd::bitwise_cast(permuted & keep); + prefix = (prefixes >> (8 * j)) & 0xFF; + } + for (int j = 0; j < kSubs; ++j) { + sub[j].store_unaligned(dst + j * kSubLanes); + } + } +}; + +// Words with at most that many valid values are scattered one by one. +constexpr int kSparseValues = 8; + +template +using Kernel = spaced::ChunkedWordKernel>; + +} // namespace + +template +void SpacedExpandRightwardAvx2(Uint* buffer, int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset) { + // Tuned on Zen 4, where other values make little difference. + constexpr int kDenseNulls = sizeof(Uint) == 4 ? 1 : 2; + spaced::InPlaceExpander, kDenseNulls, kSparseValues>::Expand( + buffer, num_values, null_count, valid_bits, valid_bits_offset); +} + +template +void SpacedExpandAvx2(const Uint* src, int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset, + Uint* output) { + spaced::OutOfPlaceExpander, kSparseValues>::Expand( + src, num_values, null_count, valid_bits, valid_bits_offset, output); +} + +template void SpacedExpandRightwardAvx2(uint32_t*, int, int, const uint8_t*, + int64_t); +template void SpacedExpandRightwardAvx2(uint64_t*, int, int, const uint8_t*, + int64_t); +template void SpacedExpandAvx2(const uint32_t*, int, int, const uint8_t*, + int64_t, uint32_t*); +template void SpacedExpandAvx2(const uint64_t*, int, int, const uint8_t*, + int64_t, uint64_t*); + +} // namespace arrow::util::internal diff --git a/cpp/src/arrow/util/spaced_simd_avx512.cc b/cpp/src/arrow/util/spaced_simd_avx512.cc new file mode 100644 index 000000000000..2d84107b849a --- /dev/null +++ b/cpp/src/arrow/util/spaced_simd_avx512.cc @@ -0,0 +1,97 @@ +// 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 "arrow/util/spaced_simd_internal.h" +#include "arrow/util/spaced_simd_kernel_internal.h" + +namespace arrow::util::internal { + +namespace { + +using Arch = xsimd::avx512bw; + +// Word kernel (see spaced::Expander) on full 512 bit batches, with the native +// expand instruction. +template +struct WordKernel { + using Batch = xsimd::batch; + using Mask = xsimd::batch_bool; + + static constexpr int kLanes = static_cast(Batch::size); + static constexpr int kChunks = 64 / kLanes; + static constexpr uint64_t kChunkMask = (uint64_t{1} << kLanes) - 1; + + static void ExpandWord(Uint* dst, const Uint* src, int src_end, uint64_t bits) { + auto chunk_mask = [&](int c) { return (bits >> (c * kLanes)) & kChunkMask; }; + // A chunk only reads below the slots of the chunks above it, so they can + // be stored in any order. Loading every chunk before storing the one above + // keeps the compiler from folding the loads into vpexpand: its memory + // operand form is much slower on Zen 4. + int c = kChunks - 1; + src_end -= std::popcount(chunk_mask(c)); + Batch above = Batch::load_unaligned(src + src_end); + for (c = kChunks - 2; c >= 0; --c) { + src_end -= std::popcount(chunk_mask(c)); + const Batch values = Batch::load_unaligned(src + src_end); + xsimd::expand(above, Mask::from_mask(chunk_mask(c + 1))) + .store_unaligned(dst + (c + 1) * kLanes); + above = values; + } + xsimd::expand(above, Mask::from_mask(chunk_mask(0))).store_unaligned(dst); + } +}; + +// Words with at most that many valid values are scattered one by one. +constexpr int kSparseValues = 8; + +template +using Kernel = WordKernel; + +} // namespace + +template +void SpacedExpandRightwardAvx512(Uint* buffer, int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset) { + // Tuned on Zen 4: the kernel is cheap enough for words with a single null. + constexpr int kDenseNulls = 0; + spaced::InPlaceExpander, kDenseNulls, kSparseValues>::Expand( + buffer, num_values, null_count, valid_bits, valid_bits_offset); +} + +template +void SpacedExpandAvx512(const Uint* src, int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset, + Uint* output) { + spaced::OutOfPlaceExpander, kSparseValues>::Expand( + src, num_values, null_count, valid_bits, valid_bits_offset, output); +} + +template void SpacedExpandRightwardAvx512(uint32_t*, int, int, const uint8_t*, + int64_t); +template void SpacedExpandRightwardAvx512(uint64_t*, int, int, const uint8_t*, + int64_t); +template void SpacedExpandAvx512(const uint32_t*, int, int, const uint8_t*, + int64_t, uint32_t*); +template void SpacedExpandAvx512(const uint64_t*, int, int, const uint8_t*, + int64_t, uint64_t*); + +} // namespace arrow::util::internal diff --git a/cpp/src/arrow/util/spaced_simd_internal.h b/cpp/src/arrow/util/spaced_simd_internal.h index f9ba5a0cc38f..a19cb89a87f8 100644 --- a/cpp/src/arrow/util/spaced_simd_internal.h +++ b/cpp/src/arrow/util/spaced_simd_internal.h @@ -21,32 +21,46 @@ #include "arrow/util/visibility.h" +// SpacedExpandRightward (in place) and SpacedExpand (out of place) on +// unsigned 4 or 8 byte lanes, one bitmap word at a time, for each dispatch +// target. Every null slot is zeroed. + +#define SPACED_EXPAND_DECLARE(TARGET) \ + template \ + ARROW_EXPORT void SpacedExpandRightward##TARGET( \ + Uint* buffer, int num_values, int null_count, const uint8_t* valid_bits, \ + int64_t valid_bits_offset); \ + template \ + ARROW_EXPORT void SpacedExpand##TARGET(const Uint* src, int num_values, \ + int null_count, const uint8_t* valid_bits, \ + int64_t valid_bits_offset, Uint* output); \ + extern template ARROW_TEMPLATE_EXPORT void SpacedExpandRightward##TARGET( \ + uint32_t*, int, int, const uint8_t*, int64_t); \ + extern template ARROW_TEMPLATE_EXPORT void SpacedExpandRightward##TARGET( \ + uint64_t*, int, int, const uint8_t*, int64_t); \ + extern template ARROW_TEMPLATE_EXPORT void SpacedExpand##TARGET( \ + const uint32_t*, int, int, const uint8_t*, int64_t, uint32_t*); \ + extern template ARROW_TEMPLATE_EXPORT void SpacedExpand##TARGET( \ + const uint64_t*, int, int, const uint8_t*, int64_t, uint64_t*) + namespace arrow::util::internal { +SPACED_EXPAND_DECLARE(None); + #if defined(ARROW_HAVE_NEON) -# define SPACED_EXPAND_ARCH128 SpacedExpandRightwardNeon +SPACED_EXPAND_DECLARE(Neon); #elif defined(ARROW_HAVE_SSE4_2) || defined(ARROW_HAVE_RUNTIME_SSE4_2) -# define SPACED_EXPAND_ARCH128 SpacedExpandRightwardSse4_2 +SPACED_EXPAND_DECLARE(Sse4_2); #endif -#if defined(SPACED_EXPAND_ARCH128) - -/// SpacedExpandRightward on unsigned 4 or 8 byte lanes, with 128 bit SIMD. -/// Every null slot is zeroed. -template -ARROW_EXPORT void SPACED_EXPAND_ARCH128(Uint* buffer, int num_values, int null_count, - const uint8_t* valid_bits, - int64_t valid_bits_offset); - -extern template ARROW_TEMPLATE_EXPORT void SPACED_EXPAND_ARCH128( - uint32_t* buffer, int num_values, int null_count, const uint8_t* valid_bits, - int64_t valid_bits_offset); - -extern template ARROW_TEMPLATE_EXPORT void SPACED_EXPAND_ARCH128( - uint64_t* buffer, int num_values, int null_count, const uint8_t* valid_bits, - int64_t valid_bits_offset); +#if defined(ARROW_HAVE_AVX2) || defined(ARROW_HAVE_RUNTIME_AVX2) +SPACED_EXPAND_DECLARE(Avx2); +#endif -#endif // SPACED_EXPAND_ARCH128 -#undef SPACED_EXPAND_ARCH128 +#if defined(ARROW_HAVE_AVX512) || defined(ARROW_HAVE_RUNTIME_AVX512) +SPACED_EXPAND_DECLARE(Avx512); +#endif } // namespace arrow::util::internal + +#undef SPACED_EXPAND_DECLARE diff --git a/cpp/src/arrow/util/spaced_simd_kernel_internal.h b/cpp/src/arrow/util/spaced_simd_kernel_internal.h new file mode 100644 index 000000000000..5ee2430d47cc --- /dev/null +++ b/cpp/src/arrow/util/spaced_simd_kernel_internal.h @@ -0,0 +1,386 @@ +// 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. + +// Word drivers shared by the implementations of SpacedExpandRightward (in +// place) and SpacedExpand (out of place) dispatched on the CPU. +// +// This header is included by translation units compiled with different +// instruction sets. Everything that generates code is therefore a member of +// a class template parameterized on the kernel, which each translation unit +// defines in an anonymous namespace: no instantiation can be shared between +// two of them by the linker. + +#pragma once + +#include +#include +#include +#include +#include + +#include "arrow/util/bit_util.h" +#include "arrow/util/endian.h" +#include "arrow/util/macros.h" +#include "arrow/util/ubsan.h" + +namespace arrow::util::internal::spaced { + +// --------------------------------------------------------------------------- +// Lookup tables (constant data only) +// +// AArch64 without FEAT_CSSC has no scalar popcount: std::popcount is a round +// trip through a SIMD register (fmov/cnt/addv/fmov), so the hot paths look +// small popcounts up instead. +// --------------------------------------------------------------------------- + +struct PopCountTable { + uint8_t count[256] = {}; + constexpr PopCountTable() { + for (unsigned m = 0; m < 256; ++m) { + count[m] = static_cast(std::popcount(m)); + } + } +}; +inline constexpr PopCountTable kPopCount{}; + +// For an 8 bit chunk mask split in sub-batches of kSubLanes lanes, byte j of +// row[mask] is the number of valid lanes in sub-batches 0..j. +template +struct PrefixTable { + uint32_t row[256] = {}; + constexpr PrefixTable() { + for (unsigned m = 0; m < 256; ++m) { + uint32_t r = 0; + for (int j = 0; j < 4 && (j + 1) * kSubLanes <= 8; ++j) { + const unsigned low = (1u << ((j + 1) * kSubLanes)) - 1u; + r |= static_cast(std::popcount(m & low)) << (8 * j); + } + row[m] = r; + } + } +}; +template +inline constexpr PrefixTable kPrefix{}; + +// --------------------------------------------------------------------------- +// Word kernels +// +// A word kernel expands one 64 bit word of the bitmap: +// +// static constexpr int kLanes; +// static void ExpandWord(Uint* dst, const Uint* src, int src_end, uint64_t bits); +// +// writes the 64 slots of `dst`, taking the popcount(bits) values that end at +// src + src_end, and zeroing the null slots. It may read up to kLanes values +// past src + src_end, and `dst` may overlap the values it reads. `src` needs +// not be aligned. +// --------------------------------------------------------------------------- + +/// Word kernel made of 64 / kLanes calls to a chunk kernel, from the top. +/// +/// The chunk kernel provides `kLanes` and +/// `static void Expand(Uint* dst, const Uint* src, unsigned mask)`, which +/// expands the low popcount(mask) values of `src` to the kLanes slots of `dst`, +/// reading at most kLanes values and loading everything before storing. +/// A chunk reads below the slots of the chunks above it, so it never sees +/// what they wrote. Empty and full chunks do not go to the chunk kernel. +template +struct ChunkedWordKernel { + static constexpr int kLanes = Chunk::kLanes; + static constexpr int kChunks = 64 / kLanes; + static constexpr uint64_t kChunkMask = (uint64_t{1} << kLanes) - 1; + + static void ExpandWord(Uint* dst, const Uint* src, int src_end, uint64_t bits) { + auto expand_chunk = [&](int c) { + const unsigned mask = static_cast((bits >> (c * kLanes)) & kChunkMask); + const int pc = kLanes == 8 ? kPopCount.count[mask] : std::popcount(mask); + if (pc == 0) { + std::memset(dst + c * kLanes, 0, sizeof(Uint) * kLanes); + return; + } + src_end -= pc; + if (pc == kLanes) { + // Much cheaper than the chunk kernel on words with few nulls. A memmove + // of a constant size: loads then stores, no call. + std::memmove(dst + c * kLanes, src + src_end, sizeof(Uint) * kLanes); + return; + } + Chunk::Expand(dst + c * kLanes, src + src_end, mask); + }; + // Unrolled by hand: Arrow builds with -O2, which keeps the loop. + [&](std::integer_sequence) { + (expand_chunk(kChunks - 1 - C), ...); + }(std::make_integer_sequence{}); + } +}; + +/// Bits [pos, pos + 64) of the bitmap. Only reads the bytes holding those bits. +/// +/// Templated on the kernel only to get one instantiation per translation unit. +template +uint64_t LoadWord(const uint8_t* bitmap, int64_t pos) { + const uint8_t* p = bitmap + (pos >> 3); + const int shift = static_cast(pos & 7); + const uint64_t lo = bit_util::FromLittleEndian(SafeLoadAs(p)); + if (shift == 0) return lo; + return (lo >> shift) | (static_cast(p[8]) << (64 - shift)); +} + +// --------------------------------------------------------------------------- +// In place driver +// +// The bitmap is walked one 64 bit word at a time from the end of the buffer, +// and every word goes to one of three paths: +// * runs (at most kDenseNulls nulls, or a single valid / null transition): +// memmove whole runs of valid values and memset runs of nulls, keeping a +// run open across words so that a long run is moved only once; +// * sparse (at most kSparseValues valid values, far from their destination): +// zero the word and scatter the values one by one; +// * otherwise: the word kernel. +// --------------------------------------------------------------------------- + +template +struct InPlaceExpander { + // Whether the word goes to the run path: at most kDenseNulls nulls, or a + // single transition between a valid run and a null run (including empty + // words). Neither needs a popcount. + static bool IsRunWord(uint64_t bits) { + uint64_t z = ~bits; + for (int k = 0; k < kDenseNulls; ++k) z &= z - 1; + return z == 0 || (bits & (bits + 1)) == 0 || (~bits & (~bits + 1)) == 0; + } + + // Handle run words downwards from the one ending at `pos`, whose bits are + // `bits`. `nulls_below` is the number of nulls below `pos` on entry, and + // below the returned position on return. Kept out of line: having the open + // run state in the main loop slows every path down. + static ARROW_NOINLINE int ExpandRunWords(Uint* out, int pos, int& nulls_below, + uint64_t bits, const uint8_t* valid_bits, + int64_t valid_bits_offset) { + // The open run: values [run_hi - x, run_hi) come from run_shift slots lower. + bool pending = false; + int run_hi = 0; + int run_shift = 0; + auto flush = [&](int run_lo) { + if (pending && run_hi > run_lo) { + std::memmove(out + run_lo, out + run_lo - run_shift, + sizeof(Uint) * static_cast(run_hi - run_lo)); + } + pending = false; + }; + while (true) { + const int lo = pos - 64; + if (bits == 0) { + // Empty words close the open run and are zeroed all at once. + flush(pos); + const int hi = pos; + pos = lo; + while (pos >= 64 && + LoadWord(valid_bits, valid_bits_offset + pos - 64) == 0) { + pos -= 64; + } + std::memset(out + pos, 0, sizeof(Uint) * static_cast(hi - pos)); + nulls_below -= hi - pos; + } else { + if (!pending && (bits >> 63) != 0) { + pending = true; + run_hi = pos; + run_shift = nulls_below; + } + // One step per run of nulls [null_lo, null_hi], from the top. Zeroing + // them is safe once the run above is moved: everything below still + // comes from lower slots. + uint64_t z = ~bits; + while (z != 0) { + const int null_hi = 63 - std::countl_zero(z); + const uint64_t valid_below = bits & ((uint64_t{1} << null_hi) - 1); + const int null_lo = valid_below == 0 ? 0 : 64 - std::countl_zero(valid_below); + const int len = null_hi - null_lo + 1; + nulls_below -= len; + flush(lo + null_hi + 1); + if (len == 1) { + out[lo + null_lo] = 0; + } else { + std::memset(out + lo + null_lo, 0, sizeof(Uint) * static_cast(len)); + } + z &= (uint64_t{1} << null_lo) - 1; + if (null_lo > 0) { + pending = true; + run_hi = lo + null_lo; + run_shift = nulls_below; + } + } + pos = lo; + } + // No nulls below: the open run (shift 0) and everything below it is + // already in place. + if (nulls_below == 0) return pos; + if (pos < 64) break; + bits = LoadWord(valid_bits, valid_bits_offset + pos - 64); + if (bits == ~uint64_t{0}) { + // Full words only extend the open run: skip them two at a time. + if (!pending) { + pending = true; + run_hi = pos; + run_shift = nulls_below; + } + pos -= 64; + while (pos >= 128) { + const uint64_t a = LoadWord(valid_bits, valid_bits_offset + pos - 64); + const uint64_t b = LoadWord(valid_bits, valid_bits_offset + pos - 128); + if ((a & b) != ~uint64_t{0}) break; + pos -= 128; + } + if (pos < 64) break; + bits = LoadWord(valid_bits, valid_bits_offset + pos - 64); + } + if (!IsRunWord(bits)) break; + } + flush(pos); + return pos; + } + + static void Expand(Uint* out, int num_values, int null_count, const uint8_t* valid_bits, + int64_t valid_bits_offset) { + constexpr int kLanes = Kernel::kLanes; + + if (null_count == 0) return; + int idx_decode = num_values - null_count; + if (idx_decode == 0) { + std::memset(out, 0, sizeof(Uint) * static_cast(num_values)); + return; + } + int pos = num_values; + while (pos >= 64) { + const int lo = pos - 64; + const uint64_t bits = LoadWord(valid_bits, valid_bits_offset + lo); + if (IsRunWord(bits)) { + int nulls_below = pos - idx_decode; + pos = ExpandRunWords(out, pos, nulls_below, bits, valid_bits, valid_bits_offset); + if (nulls_below == 0) return; + idx_decode = pos - nulls_below; + continue; + } + const int pc = std::popcount(bits); + idx_decode -= pc; + const int src_begin = idx_decode; + const int shift = lo - src_begin; + if (pc <= kSparseValues && shift >= 64) { + // Source and destination do not overlap. + std::memset(out + lo, 0, 64 * sizeof(Uint)); + int q = src_begin + pc; + uint64_t b = bits; + while (b != 0) { + const int i = 63 - std::countl_zero(b); + b ^= uint64_t{1} << i; + out[lo + i] = out[--q]; + } + } else if (src_begin + pc + kLanes > num_values) { + // The kernel reads up to kLanes values past the valid values of the + // word, which can go past the end of the buffer for the top words + // only. Stage those in a buffer. + Uint staged[64 + kLanes]; + std::memcpy(staged, out + src_begin, sizeof(Uint) * static_cast(pc)); + Kernel::ExpandWord(out + lo, staged, pc, bits); + } else { + Kernel::ExpandWord(out + lo, out, src_begin + pc, bits); + } + pos = lo; + if (shift == 0) return; + } + for (int i = pos - 1; i >= 0; --i) { + if (bit_util::GetBit(valid_bits, valid_bits_offset + i)) { + out[i] = out[--idx_decode]; + } else { + out[i] = 0; + } + } + } +}; + +// --------------------------------------------------------------------------- +// Out of place driver +// +// Without the in place constraints, runs of full and empty words are a single +// memcpy / memset, words with at most kSparseValues valid values are zeroed +// and scattered, and every other word goes to the kernel. +// --------------------------------------------------------------------------- + +template +struct OutOfPlaceExpander { + static void Expand(const Uint* src, int num_values, int null_count, + const uint8_t* valid_bits, int64_t valid_bits_offset, Uint* out) { + constexpr int kLanes = Kernel::kLanes; + const int num_valid = num_values - null_count; + if (num_valid == 0 || null_count == 0) { + // No need to look at the bitmap + if (num_valid == 0) { + std::memset(out, 0, sizeof(Uint) * static_cast(num_values)); + } else { + std::memcpy(out, src, sizeof(Uint) * static_cast(num_values)); + } + return; + } + // End of the values of src still to be written + int src_end = num_valid; + int pos = num_values; + while (pos >= 64) { + const int lo = pos - 64; + const uint64_t bits = LoadWord(valid_bits, valid_bits_offset + lo); + if (bits == ~uint64_t{0} || bits == 0) { + const int hi = pos; + pos = lo; + while (pos >= 64 && + LoadWord(valid_bits, valid_bits_offset + pos - 64) == bits) { + pos -= 64; + } + const auto size = sizeof(Uint) * static_cast(hi - pos); + if (bits == 0) { + std::memset(out + pos, 0, size); + } else { + src_end -= hi - pos; + std::memcpy(out + pos, src + src_end, size); + } + continue; + } + const int pc = std::popcount(bits); + if (pc <= kSparseValues) { + std::memset(out + lo, 0, 64 * sizeof(Uint)); + const Uint* values = src + src_end - pc; + for (uint64_t b = bits; b != 0; b &= b - 1) { + out[lo + std::countr_zero(b)] = SafeLoad(values++); + } + } else if (src_end + kLanes > num_valid) { + // The kernel would read past the end of src + Uint staged[64 + kLanes]; + std::memcpy(staged, src + src_end - pc, sizeof(Uint) * static_cast(pc)); + Kernel::ExpandWord(out + lo, staged, pc, bits); + } else { + Kernel::ExpandWord(out + lo, src, src_end, bits); + } + src_end -= pc; + pos = lo; + } + for (int i = pos - 1; i >= 0; --i) { + out[i] = bit_util::GetBit(valid_bits, valid_bits_offset + i) + ? SafeLoad(src + --src_end) + : 0; + } + } +}; + +} // namespace arrow::util::internal::spaced diff --git a/cpp/src/arrow/util/spaced_test.cc b/cpp/src/arrow/util/spaced_test.cc index a57bff32fde8..7bfa3925e005 100644 --- a/cpp/src/arrow/util/spaced_test.cc +++ b/cpp/src/arrow/util/spaced_test.cc @@ -27,13 +27,10 @@ #include "arrow/testing/gtest_util.h" #include "arrow/util/bit_util.h" +#include "arrow/util/cpu_info.h" #include "arrow/util/spaced_internal.h" #include "arrow/util/spaced_simd_internal.h" -#if defined(ARROW_HAVE_SSE4_2) || defined(ARROW_HAVE_RUNTIME_SSE4_2) -# include "arrow/util/cpu_info.h" -#endif - namespace arrow::util::internal { namespace { @@ -125,7 +122,7 @@ bool BitwiseEqual(T a, T b) { // Check `expand` against a straightforward implementation. Null slots must // be zero if `zeros_nulls`, and are not checked otherwise. template -void CheckExpand(Expand&& expand, bool zeros_nulls) { +void CheckExpand(Expand&& expand, bool zeros_nulls = true) { constexpr int kGuard = 16; std::mt19937_64 rng(42); for (int n : {0, 1, 7, 63, 64, 65, 127, 128, 129, 900, 2048}) { @@ -166,42 +163,90 @@ void CheckExpand(Expand&& expand, bool zeros_nulls) { } } +// An out of place `expand` through the in place interface of CheckExpand: the +// dense values are first copied to a misaligned buffer of their exact size. +template +auto OutOfPlace(Expand&& expand) { + return + [expand](T* buffer, int n, int null_count, const uint8_t* bitmap, int64_t offset) { + const size_t size = sizeof(T) * static_cast(n - null_count); + std::vector src(size + 1); + std::memcpy(src.data() + 1, buffer, size); + expand(reinterpret_cast(src.data() + 1), n, null_count, bitmap, offset, + buffer); + }; +} + +bool CpuSupports(int64_t flags) { + return ::arrow::internal::CpuInfo::GetInstance()->IsSupported(flags); +} + template -class TestSpacedExpandRightward : public ::testing::Test {}; +class TestSpacedExpandTarget : public ::testing::Test {}; using UintTypes = ::testing::Types; -TYPED_TEST_SUITE(TestSpacedExpandRightward, UintTypes); +TYPED_TEST_SUITE(TestSpacedExpandTarget, UintTypes); -TYPED_TEST(TestSpacedExpandRightward, Scalar) { +TYPED_TEST(TestSpacedExpandTarget, Scalar) { CheckExpand(&SpacedExpandRightwardScalar, /*zeros_nulls=*/false); + CheckExpand(OutOfPlace(&SpacedExpandScalar)); +} + +TYPED_TEST(TestSpacedExpandTarget, None) { + CheckExpand(&SpacedExpandRightwardNone); + CheckExpand(OutOfPlace(&SpacedExpandNone)); } #if defined(ARROW_HAVE_NEON) -TYPED_TEST(TestSpacedExpandRightward, Neon) { - CheckExpand(&SpacedExpandRightwardNeon, /*zeros_nulls=*/true); +TYPED_TEST(TestSpacedExpandTarget, Neon) { + CheckExpand(&SpacedExpandRightwardNeon); + CheckExpand(OutOfPlace(&SpacedExpandNeon)); } #endif #if defined(ARROW_HAVE_SSE4_2) || defined(ARROW_HAVE_RUNTIME_SSE4_2) -TYPED_TEST(TestSpacedExpandRightward, Sse4_2) { - if (!::arrow::internal::CpuInfo::GetInstance()->IsSupported( - ::arrow::internal::CpuInfo::SSE4_2)) { +TYPED_TEST(TestSpacedExpandTarget, Sse4_2) { + if (!CpuSupports(::arrow::internal::CpuInfo::SSE4_2)) { GTEST_SKIP() << "Test requires SSE4.2"; } - CheckExpand(&SpacedExpandRightwardSse4_2, /*zeros_nulls=*/true); + CheckExpand(&SpacedExpandRightwardSse4_2); + CheckExpand(OutOfPlace(&SpacedExpandSse4_2)); +} +#endif + +#if defined(ARROW_HAVE_AVX2) || defined(ARROW_HAVE_RUNTIME_AVX2) +TYPED_TEST(TestSpacedExpandTarget, Avx2) { + if (!CpuSupports(::arrow::internal::CpuInfo::AVX2)) { + GTEST_SKIP() << "Test requires AVX2"; + } + CheckExpand(&SpacedExpandRightwardAvx2); + CheckExpand(OutOfPlace(&SpacedExpandAvx2)); +} +#endif + +#if defined(ARROW_HAVE_AVX512) || defined(ARROW_HAVE_RUNTIME_AVX512) +TYPED_TEST(TestSpacedExpandTarget, Avx512) { + if (!CpuSupports(::arrow::internal::CpuInfo::AVX512)) { + GTEST_SKIP() << "Test requires AVX512"; + } + CheckExpand(&SpacedExpandRightwardAvx512); + CheckExpand(OutOfPlace(&SpacedExpandAvx512)); } #endif -// Through the public template: 4 and 8 byte types are dispatched (possibly -// to the scalar implementation), other widths use the scalar implementation. +// Through the public templates: 4 and 8 byte types are dispatched, other +// widths use the scalar implementations. SpacedExpandRightwardScalar is the +// only one not to zero every null slot. template -class TestSpacedExpandRightwardTyped : public ::testing::Test {}; +class TestSpacedExpandTyped : public ::testing::Test {}; using ValueTypes = ::testing::Types; -TYPED_TEST_SUITE(TestSpacedExpandRightwardTyped, ValueTypes); +TYPED_TEST_SUITE(TestSpacedExpandTyped, ValueTypes); -TYPED_TEST(TestSpacedExpandRightwardTyped, Dispatch) { - CheckExpand(&SpacedExpandRightward, /*zeros_nulls=*/false); +TYPED_TEST(TestSpacedExpandTyped, Dispatch) { + constexpr bool kDispatched = sizeof(TypeParam) == 4 || sizeof(TypeParam) == 8; + CheckExpand(&SpacedExpandRightward, /*zeros_nulls=*/kDispatched); + CheckExpand(OutOfPlace(&SpacedExpand)); } } // namespace diff --git a/cpp/src/parquet/decoder.cc b/cpp/src/parquet/decoder.cc index f90b56a7321c..1c2982011e9a 100644 --- a/cpp/src/parquet/decoder.cc +++ b/cpp/src/parquet/decoder.cc @@ -358,6 +358,9 @@ class PlainDecoder : public TypedDecoderImpl { int Decode(T* buffer, int max_values) override; + int DecodeSpaced(T* buffer, int num_values, int null_count, const uint8_t* valid_bits, + int64_t valid_bits_offset) override; + int DecodeArrow(int num_values, int null_count, const uint8_t* valid_bits, int64_t valid_bits_offset, typename EncodingTraits::Accumulator* builder) override; @@ -532,6 +535,39 @@ int PlainDecoder::Decode(T* buffer, int max_values) { return max_values; } +template +int PlainDecoder::DecodeSpaced(T* buffer, int num_values, int null_count, + const uint8_t* valid_bits, + int64_t valid_bits_offset) { + if constexpr (std::is_same_v || std::is_same_v) { + // Not a plain copy of the page + return TypedDecoderImpl::DecodeSpaced(buffer, num_values, null_count, + valid_bits, valid_bits_offset); + } else { + if (null_count == 0) { + return Decode(buffer, num_values); + } + // Same checks as Decode and TypedDecoderImpl::DecodeSpaced, but the values + // are spread out straight from the page instead of being copied first. + const int values_to_read = num_values - null_count; + const int values_read = std::min(values_to_read, this->num_values_); + const int64_t bytes_to_decode = static_cast(values_read) * sizeof(T); + if (bytes_to_decode > this->len_) { + ParquetException::EofException(); + } + if (values_read != values_to_read) { + throw ParquetException("Number of values / definition_levels read did not match"); + } + ::arrow::util::internal::SpacedExpand(reinterpret_cast(this->data_), + num_values, null_count, valid_bits, + valid_bits_offset, buffer); + this->data_ += bytes_to_decode; + this->len_ -= static_cast(bytes_to_decode); + this->num_values_ -= values_read; + return num_values; + } +} + // PLAIN decoder implementation for BOOLEAN class PlainBooleanDecoder : public TypedDecoderImpl, public BooleanDecoder {