Optimize RLE decoding using a warp-balanced chunking approach - #23271
Conversation
vyasr
added
Performance
…chunked expand Readability-only. Ptxas SMEM unchanged at 5168 B for preprocess_levels_kernel on SM 8.0. All ctest -R PARQUET pass including PARQUET_RLE_CHUNKED_EQUIVALENCE_TEST. No performance change expected.
Dead scaffolding from earlier iterations of the chunked-expand path. The actual decode_next_chunked function allocates gen_out_off and gen_meta directly as __shared__ arrays; this struct was never instantiated anywhere in the tree. No functional change. SMEM/register footprint of preprocess_levels_kernel unchanged (5168 B / 48 reg on SM 8.0). All PARQUET tests pass.
Both decode_next_ring and decode_next_chunked had the same all-zeros short-circuit for level_bits == 0. Pull it up into decode_next so the helpers stay focused on the general RLE path. Uses cur_values-relative addressing (the ring version's form was only correct because cur_values was always 0 on entry in practice).
…nce test The rle_stream::init signature added a Group parameter as part of the SMEM-staging refactor; update the chunked-equivalence test to pass cg::this_thread_block() to match.
A100 nvbench parquet_read_decode sweep on LIST/STRUCT/STRING shows k=1024 is 80-95% faster than k=512 on the chunked-expand path in preprocess_levels_kernel. k=2048 saturates or regresses ~10% on LIST due to occupancy pressure, so 1024 is the sweet spot. Keep k=512 on sm_70 (V100) where the larger SMEM footprint would exceed the preprocess_levels_kernel budget. H100/Blackwell tiers are left as TODO in the comment; sm_80+ tier applies until they are measured. SMEM cost: (2 * kGenRuns + 1) * 4 bytes = 8196 B at k=1024 (vs 4100 B at k=512). Well within A100/H100 budgets.
level_mask is invariant across the whole decode_next_chunked call (depends only on level_bits, which is class state). Move it out to the top of the function. Verified with cuobjdump --dump-resource-usage that this is a no-op for register counts on sm_80/86/90 - the compiler was already hoisting it - so this is a readability/intent-signalling change only.
The earlier sm_80+ bump to 1024 was justified by an isolated kGenRuns sweep whose k=512 STRUCT baseline was in a pathological regime (245-415 ms with 46-75% noise, hitting nvbench timeouts). A full 36-config parquet_read_decode A/B against upstream/main on A100 shows that k=1024 vs k=512 differences are within noise, and a matching sweep on H100 found the same pattern - k=2048 was numerically best but within noise of both k=512 and k=1024. Since no architecture shows a meaningful preference in real workloads, drop the arch-adaptive split and use a single kGenRuns=512 everywhere. This also reclaims 4 KiB of SMEM per block on sm_80+, easing occupancy pressure on register-constrained architectures like sm_86.
Rename the identifiers we introduced with the chunked-expand path to
match the snake_case convention used throughout libcudf:
kGenRuns -> max_runs_per_chunk (also better describes what the
constant controls)
kWarps -> num_warps (matches existing num_rle_stream_*
naming)
gen_out_off, gen_meta (and their _v span views) -> chunk_out_off,
chunk_meta (the opaque "gen_" prefix came from
"generated in phase 1"; "chunk_"
parallels the existing s_chunk_*
shared vars produced alongside them)
Verified with build-cudf-cpp and PARQUET_RLE_CHUNKED_EQUIVALENCE_TEST (8/8).
warp_fill was only used from the RLE-run arm of the chunked-expand loop. Inline it at the call site and drop the helper. This removes the now-unnecessary __forceinline__ / __restrict__ annotations and the comments that referenced warp_fill by name; the loop is short enough to be self-explanatory at the call site.
Sweeps of {256, 512, 1024, 2048, 4096} on A100, H100, and B200 all show
1024 either as the numerical optimum or within noise of it (H100 leaned
toward 2048 by a small margin, still within noise). The delta over 512
is small - typically a few percent - but 1024 is the consistent winner
across the modern arches, so specialize by __CUDA_ARCH__.
sm_70/sm_75 stay at 512 because 1024 does not fit the
preprocess_levels_kernel SMEM budget on those older architectures.
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe Parquet RLE stream gains an optional chunked-expand decoder with partial-run resumption and configurable staging capacity. Parquet level preprocessing selects this mode and supplies the staging size during repetition- and definition-level decoder initialization. ChangesParquet RLE chunked decoding
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
cpp/tests/io/parquet_rle_chunked_equivalence_test.cu (2)
6-20: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMissing explicit gtest include.
TEST_F/EXPECT_EQare used directly but the file doesn't include<cudf_test/cudf_gtest.hpp>(or rawgtest/gtest.h); it relies on a transitive include viacudf_test/base_fixture.hpp.As per coding guidelines, "Test files must include
#include <cudf_test/cudf_gtest.hpp>instead of rawgtest/gtest.h."🧪 Proposed fix
`#include` "../../src/io/parquet/rle_stream.cuh" +#include <cudf_test/cudf_gtest.hpp> `#include` <cudf_test/base_fixture.hpp> `#include` <cudf_test/testing_main.hpp>🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/io/parquet_rle_chunked_equivalence_test.cu` around lines 6 - 20, Add the explicit <cudf_test/cudf_gtest.hpp> include to parquet_rle_chunked_equivalence_test.cu alongside the other cudf_test headers, so TEST_F and EXPECT_EQ do not rely on the transitive inclusion from base_fixture.hpp.Source: Coding guidelines
179-186: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
ManyShortRepeatedRunsdoesn't reliably cross the chunk boundary on sm_80+.600 short runs is below
max_runs_per_chunk(1024) on sm_80+, so the outer multi-chunk loop indecode_next_chunked(which is the core new logic this PR adds) may not actually be exercised on the modern architectures this feature targets, only on sm_70/75 (max 512). Consider sizing the run count offcudf::io::parquet::detail::max_runs_per_chunkdirectly (e.g.max_runs_per_chunk * 2 + 1) so the test deterministically crosses the boundary regardless of the GPU running CI.As per coding guidelines, "Test suites should cover edge cases such as empty input, null values, sliced columns, boundary sizes, and multi-block sizes."
TEST_F(ParquetRleChunkedEquivalenceTest, ManyShortRepeatedRuns) { using cudf::io::parquet::detail::max_runs_per_chunk; int const num_runs = max_runs_per_chunk * 2 + 1; // guarantees >1 chunk on any arch std::vector<uint8_t> encoded; for (int i = 0; i < num_runs; ++i) { append_repeated(encoded, 1, i & 15, 4); } run_case<uint8_t>(encoded, 4, num_runs, num_runs); }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/io/parquet_rle_chunked_equivalence_test.cu` around lines 179 - 186, Update the ManyShortRepeatedRuns test to derive its run count from cudf::io::parquet::detail::max_runs_per_chunk, using a value greater than two chunk capacities (such as 2 * max_runs_per_chunk + 1). Use that count in the encoding loop and run_case expectations so the test deterministically exercises decode_next_chunked’s multi-chunk path on every supported architecture.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@cpp/tests/io/parquet_rle_chunked_equivalence_test.cu`:
- Around line 6-20: Add the explicit <cudf_test/cudf_gtest.hpp> include to
parquet_rle_chunked_equivalence_test.cu alongside the other cudf_test headers,
so TEST_F and EXPECT_EQ do not rely on the transitive inclusion from
base_fixture.hpp.
- Around line 179-186: Update the ManyShortRepeatedRuns test to derive its run
count from cudf::io::parquet::detail::max_runs_per_chunk, using a value greater
than two chunk capacities (such as 2 * max_runs_per_chunk + 1). Use that count
in the encoding loop and run_case expectations so the test deterministically
exercises decode_next_chunked’s multi-chunk path on every supported
architecture.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0c826ec9-5852-4d01-9114-574ad1224e6b
📒 Files selected for processing (4)
cpp/src/io/parquet/decode_preprocess.cucpp/src/io/parquet/rle_stream.cuhcpp/tests/CMakeLists.txtcpp/tests/io/parquet_rle_chunked_equivalence_test.cu
|
Ah good catch, sorry about that! I fixed it so the signs are consistent. |
mhaseeb123
added
3 - Ready for Review
The chunked-expand decoder had cross-call resume state (partial_run_meta / partial_run_total / partial_run_done plus s_run0_payload_offset) written in anticipation of a chunked-expand dict_stream in decode_page_data_generic that was later benchmarked and formally deferred (see the DEFER verdict recorded on feature branch opt/rle-def-rep-split, geomean B/A 1.0083x inside noise band with a Q18 regression). The current sole caller (preprocess_levels_kernel) invokes decode_next_chunked exactly once per stream with max_output_values = INT_MAX, so a single RLE run cannot exceed the output window and the resume path is structurally unreachable. Untested unreachable state-machine code was flagged in review; deleting it is the cleanest response. Replace the overflow-stash block with a cudf_assert enforcing the single-call invariant, and leave a comment referencing the prior chunked-dict work (commits 4dbde92dd1 / 3e349acb27) as a recovery pointer for whoever revisits multi-call chunked decoding.
The rle_stream::decode_next fast path for level_bits == 0 uses `cur_values + written + t` as the ring index rather than the simpler `written + t`. No current caller enters this fast path with cur_values > 0 (the writer floors dict_rle_bits >= 1 in chunk_dict.cu, and REPETITION/DEFINITION decoders in decode_preprocess.cu are single-call), so the simpler form would also pass all end-to-end tests. Documenting the invariant explicitly so it is preserved defensively for any future caller that iterates decode_next with level_bits == 0.
Introduce class-level constants `run_desc_literal_flag` (1u << 31) and `run_desc_offset_mask` (0x7fffffffu) on rle_stream and use them at all five sites that previously spelled the literal-run flag / offset mask inline. Also folds the standalone block comment about the 2 GiB invariant into the constants' docblock. Pure naming refactor -- no behavior change. Verified PARQUET_TEST still passes 481/481.
| // short runs. | ||
| int p = lo + lane; | ||
| int run_idx = | ||
| static_cast<int>(cuda::std::upper_bound( |
There was a problem hiding this comment.
hmm, if this is slower for long runs, can this first run_idx be something that thread 0 calculates while it's traversing phase 1? it would only need to calculate it for the first thread of each warp. e.g. as it determines it crosses output idx = 2048 (or wherever the warps decide to split their work) it saves that in a small shared memory buf. then each thread in a given warp can use their lane 0 value as an initial guess and the linear search on line 603 will correct it for each thread. so no binary search at all
There was a problem hiding this comment.
I can give this a shot, but I'm a little skeptical since this introduces yet another serialized workstream trading off with what I'd hope is a fairly quick binary search. That new workstream would have to run after the current phase 1 work because you don't know the number of runs you're processing until that initial scan finishes. Maybe binary searching a shmem array will be slow enough that this approach will pay off, though. Let me see.
There was a problem hiding this comment.
I had my agent implement this suggestion and benchmark (A100 80GB, CUDA 12.9) to decide. Here's what it produced. Here's the output:
Baseline (8109e2f1cb): per-lane upper_bound in Phase 2 (current code after the p → out_pos rename).
Variant: thread 0 precomputes warp_seed_run_idx[num_warps] in Phase 1 via a single linear sweep; Phase 2 reads the seed from shared memory instead of searching.
Results — parquet_read_decode + parquet_read_fixed_width_struct, A100, NONE compression
| data_type | cardinality | run_length | Baseline GPU | Variant GPU | Δ | %Δ |
|---|---|---|---|---|---|---|
| INTEGRAL | 0 | 1 | 9.166 ms ±3.6% | 9.683 ms ±4.4% | +517 µs | +5.6% |
| INTEGRAL | 1000 | 1 | 9.917 ms ±0.3% | 10.506 ms ±1.7% | +589 µs | +5.9% |
| INTEGRAL | 0 | 32 | 7.860 ms ±0.6% | 8.314 ms ±2.3% | +454 µs | +5.8% |
| INTEGRAL | 1000 | 32 | 7.816 ms ±0.5% | 8.276 ms ±1.8% | +460 µs | +5.9% |
| LIST | 0 | 1 | 14.808 ms ±0.7% | 15.391 ms ±0.6% | +583 µs | +3.9% |
| LIST | 1000 | 1 | 17.390 ms ±0.2% | 17.852 ms ±0.3% | +463 µs | +2.7% |
| STRUCT | 0 | 1 | 19.853 ms ±6.8% | 19.982 ms ±2.6% | +128 µs | +0.7% (same) |
| STRUCT | 1000 | 1 | 13.734 ms ±0.4% | 14.148 ms ±0.2% | +414 µs | +3.0% |
| STRING | 0 | 1 | 10.706 ms ±0.3% | 10.839 ms ±0.3% | +133 µs | +1.2% |
(20 configs total; no config improved in the variant.)
Verdict: keeping the current per-lane upper_bound. The variant is uniformly slower, with INTEGRAL showing a clear ~6% regression well outside noise. My hypothesis: on sm_80, 32 parallel upper_bound calls over a 1024-entry SMEM array are heavily pipelined L1-cache hits running concurrently across all 32 lanes — effectively free. The serial thread-0 seed loop in Phase 1 adds latency that compounds for short/uniform runs (many chunks per call). The upper_bound is the better choice here.
Address review comment: the loop-carried output position is now named `out_pos` for clarity. Comments updated to match. No behavior change.
decode_next_chunked did not clamp run_len to the remaining output window before appending the run to the chunk table. The cudf_assert that guarded the invariant fires in debug builds but is compiled out in release, leaving Phase 2 free to write past the output buffer whenever a run straddles out_end. This happens legitimately when row-range filtering (skip_rows / num_rows / bounds-page) produces an output_count smaller than INT_MAX, causing silent data corruption: wrong column sizes and null counts on nested (LIST / STRUCT) and nullable columns. Fix: replace the assert with an explicit clamp. cur has already been advanced past the full payload, which is correct -- we emit only the clamped run_len values and the outer loop exits on the next iteration because out_pos_total will equal out_end.
| auto const len = static_cast<int>(cuda::std::distance(_start, _end)); | ||
| if (len > 0 && len <= smem_stage_size) { | ||
| if (len > 0 && len <= stage_capacity) { | ||
| cuda::memcpy_async(group, _smem_stage, _start, static_cast<size_t>(len), *_copy_barrier); |
There was a problem hiding this comment.
perhaps assert if len is too small or large?
There was a problem hiding this comment.
Good idea. I'll toss that into the next PR so we can keep things moving along.
|
/merge |
3 participants
Description
This PR implements an alternative approach for the RLE decoding. The old approach used a producer-consumer model where warp 0 populates a ring buffer of runs for other warps to pick off. That leads to two sources of imbalance:
With the new approach, the full stream is split into chunks of a fixed size (determined at compile time). Within each chunk, thread 0 does a serial pass through the data to find all of the RLE headers and populates the associated splits and metadata (RLE vs bit-packed) into a shared memory array. Then, all warps can cooperatively read through all of that data. Since data is parsed by chunk rather than by run, there is no longer any imbalance between warps. Warps keep track of boundaries via the same shared memory arrays, and therefore warps can start in the middle of any run and completely traverse runs short enough to fit within their chunks.
Checklist