Optimize RLE decoding by staging in shared memory - #23090
Conversation
|
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:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThis PR adds shared-memory staging for Parquet RLE level decoding, wires a block barrier through preprocessing, and updates fixed decode call sites to use the new cooperative-group ChangesParquet RLE level staging
Estimated code review effort: 3 (Moderate) | ~25 minutes 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.
Actionable comments posted: 2
🤖 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.
Inline comments:
In `@cpp/src/io/parquet/decode_preprocess.cu`:
- Around line 428-429: The staging buffers in decode_preprocess.cu are declared
as plain __shared__ uint8_t arrays, but rle_stream::init may write to them with
16-byte vectorized stores via uint4*, so they must be 16-byte aligned. Update
the def_stage and rep_stage shared-buffer declarations to use explicit 16-byte
alignment, matching the alignment pattern already used by state_g, and keep the
rest of the staging logic unchanged.
In `@cpp/src/io/parquet/rle_stream.cuh`:
- Around line 236-250: The vectorized copy path in the rle_stream copy logic
only checks `_start` alignment and then casts `s_dst` to `uint4*`, which can
misalign the destination store. Update the alignment guard in the copy routine
so the `uint4` path is used only when both `_start` and `s_dst` are 16-byte
aligned, and otherwise fall back to the byte-wise loop. Use the existing copy
block in `rle_stream` as the fix point, and ensure the caller-side staging
buffer in `decode_preprocess` remains compatible with the fast path.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: b36eaded-4324-4d2e-8aa4-12c6ec20f49d
📒 Files selected for processing (2)
cpp/src/io/parquet/decode_preprocess.cucpp/src/io/parquet/rle_stream.cuh
| // 16-byte vectorized copy of the aligned body; byte copy for the tail. | ||
| // Falls back to a plain byte copy when the source is not 16B-aligned. |
There was a problem hiding this comment.
Does CCCL offer any similar vectorized copying utilities? I'd like to keep this code a bit higher-level if there are utilities we can leverage.
There was a problem hiding this comment.
Or at least in a reusable device utility somewhere in src/io/utilities
There was a problem hiding this comment.
Ideally we would be using ublckcp on hardware that supports it
There was a problem hiding this comment.
Probably cuda::memcpy_async https://nvidia.github.io/cccl/unstable/libcudacxx/extended_api/asynchronous_operations/memcpy_async.html
There was a problem hiding this comment.
Yes, I'm aware of memcpy_async, it's what I'm using in my local copy of the branch. That will pick the right algorithm for each hardware, including executing TMA loads if available. I'm trying to get benchmark numbers on both pre- and post-hopper GPUs to see if there are observable performance differences.
There was a problem hiding this comment.
I switched to 26bf960.
Ideally we would be using ublckcp on hardware that supports it
There is some more subtlety to this part. cuda::memcpy_async should use the best PTX instructions available, but both cp.async and cp.async.bulk impose stricter alignment requirements that the current code doesn't satisfy by default. I tried to implement some basic peeling off of enough initial/final elements to ensure the requisite alignment, but I had mixed results. With 4-byte alignment I failed to produce cp.async instructions. My guess is that I missed one of the other requirements here, I'll inspect the code again tomorrow to confirm. With 16-byte alignment I did observe cp.async.bulk instructions, but the resulting code was about 5% slower in performance. My guess is that the head/tail peeling logic adds enough extra branching, register usage (from the local loop variables etc), and possible divergence to wash out any potential benefits from TMAs.
I do think this question is worth exploring further, but it will require some more careful benchmarking that is probably best done in a follow-up PR.
# parquet_read_decode
## [0] NVIDIA H100 80GB HBM3
| data_type | io_type | compression_type | cardinality | run_length | data_size | row_group_size_bytes | row_group_size_rows | Ref Time | Ref Noise | Cmp Time | Cmp Noise | Diff | %Diff | Status |
|-------------|---------------|--------------------|---------------|--------------|-------------|------------------------|-----------------------|------------|-------------|------------|-------------|----------------|---------|----------|
| INTEGRAL | DEVICE_BUFFER | NONE | 0 | 1 | 536870912 | 0 | 0 | 63.176 ms | 20.58% | 18.578 ms | 128.37% | -44598.708 us | -70.59% | ESC[32mFASTESC[39m |
| INTEGRAL | DEVICE_BUFFER | NONE | 1000 | 1 | 536870912 | 0 | 0 | 64.432 ms | 16.26% | 20.610 ms | 123.00% | -43821.974 us | -68.01% | ESC[32mFASTESC[39m |
| INTEGRAL | DEVICE_BUFFER | NONE | 0 | 32 | 536870912 | 0 | 0 | 63.687 ms | 17.89% | 16.202 ms | 139.35% | -47485.116 us | -74.56% | ESC[32mFASTESC[39m |
| INTEGRAL | DEVICE_BUFFER | NONE | 1000 | 32 | 536870912 | 0 | 0 | 60.351 ms | 25.01% | 16.850 ms | 136.81% | -43501.693 us | -72.08% | ESC[32mFASTESC[39m |
| FLOAT | DEVICE_BUFFER | NONE | 0 | 1 | 536870912 | 0 | 0 | 32.465 ms | 16.52% | 6.933 ms | 118.82% | -25531.934 us | -78.64% | ESC[32mFASTESC[39m |
| FLOAT | DEVICE_BUFFER | NONE | 1000 | 1 | 536870912 | 0 | 0 | 32.179 ms | 25.84% | 13.352 ms | 101.52% | -18826.981 us | -58.51% | ESC[32mFASTESC[39m |
| FLOAT | DEVICE_BUFFER | NONE | 0 | 32 | 536870912 | 0 | 0 | 34.147 ms | 23.71% | 11.013 ms | 114.63% | -23133.590 us | -67.75% | ESC[32mFASTESC[39m |
| FLOAT | DEVICE_BUFFER | NONE | 1000 | 32 | 536870912 | 0 | 0 | 32.584 ms | 18.22% | 11.382 ms | 106.87% | -21202.834 us | -65.07% | ESC[32mFASTESC[39m |
| BOOL8 | DEVICE_BUFFER | NONE | 0 | 1 | 536870912 | 0 | 0 | 214.633 ms | 19.68% | 41.164 ms | 173.74% | -173469.378 us | -80.82% | ESC[32mFASTESC[39m |
| BOOL8 | DEVICE_BUFFER | NONE | 1000 | 1 | 536870912 | 0 | 0 | 204.488 ms | 29.38% | 38.173 ms | 181.19% | -166315.640 us | -81.33% | ESC[32mFASTESC[39m |
| BOOL8 | DEVICE_BUFFER | NONE | 0 | 32 | 536870912 | 0 | 0 | 205.120 ms | 30.56% | 35.119 ms | 186.37% | -170001.083 us | -82.88% | ESC[32mFASTESC[39m |
| BOOL8 | DEVICE_BUFFER | NONE | 1000 | 32 | 536870912 | 0 | 0 | 209.800 ms | 26.85% | 33.666 ms | 198.17% | -176134.032 us | -83.95% | ESC[32mFASTESC[39m |
| DECIMAL | DEVICE_BUFFER | NONE | 0 | 1 | 536870912 | 0 | 0 | 19.005 ms | 25.44% | 7.817 ms | 93.29% | -11188.078 us | -58.87% | ESC[32mFASTESC[39m |
| DECIMAL | DEVICE_BUFFER | NONE | 1000 | 1 | 536870912 | 0 | 0 | 21.292 ms | 20.74% | 9.174 ms | 87.75% | -12117.822 us | -56.91% | ESC[32mFASTESC[39m |
| DECIMAL | DEVICE_BUFFER | NONE | 0 | 32 | 536870912 | 0 | 0 | 19.836 ms | 21.48% | 8.443 ms | 93.90% | -11392.889 us | -57.43% | ESC[32mFASTESC[39m |
| DECIMAL | DEVICE_BUFFER | NONE | 1000 | 32 | 536870912 | 0 | 0 | 20.423 ms | 25.74% | 8.057 ms | 94.28% | -12365.332 us | -60.55% | ESC[32mFASTESC[39m |
| TIMESTAMP | DEVICE_BUFFER | NONE | 0 | 1 | 536870912 | 0 | 0 | 28.283 ms | 21.98% | 4.009 ms | 1.58% | -24274.046 us | -85.82% | ESC[32mFASTESC[39m |
| TIMESTAMP | DEVICE_BUFFER | NONE | 1000 | 1 | 536870912 | 0 | 0 | 29.028 ms | 14.80% | 11.921 ms | 96.07% | -17107.517 us | -58.93% | ESC[32mFASTESC[39m |
| TIMESTAMP | DEVICE_BUFFER | NONE | 0 | 32 | 536870912 | 0 | 0 | 29.196 ms | 24.97% | 9.650 ms | 105.65% | -19545.256 us | -66.95% | ESC[32mFASTESC[39m |
| TIMESTAMP | DEVICE_BUFFER | NONE | 1000 | 32 | 536870912 | 0 | 0 | 27.585 ms | 19.37% | 9.823 ms | 99.64% | -17761.693 us | -64.39% | ESC[32mFASTESC[39m |
| DURATION | DEVICE_BUFFER | NONE | 0 | 1 | 536870912 | 0 | 0 | 27.329 ms | 21.69% | 3.791 ms | 1.21% | -23538.050 us | -86.13% | ESC[32mFASTESC[39m |
| DURATION | DEVICE_BUFFER | NONE | 1000 | 1 | 536870912 | 0 | 0 | 9.350 ms | 105.72% | 12.979 ms | 73.82% | 3.629 ms | 38.81% | ESC[34mSAMEESC[39m |
| DURATION | DEVICE_BUFFER | NONE | 0 | 32 | 536870912 | 0 | 0 | 3.663 ms | 1.61% | 7.275 ms | 76.54% | 3.612 ms | 98.60% | ESC[31mSLOWESC[39m |
| DURATION | DEVICE_BUFFER | NONE | 1000 | 32 | 536870912 | 0 | 0 | 10.068 ms | 114.43% | 6.682 ms | 73.62% | -3386.333 us | -33.63% | ESC[34mSAMEESC[39m |
| STRING | DEVICE_BUFFER | NONE | 0 | 1 | 536870912 | 0 | 0 | 12.565 ms | 41.96% | 8.438 ms | 0.33% | -4126.815 us | -32.84% | ESC[32mFASTESC[39m |
| STRING | DEVICE_BUFFER | NONE | 1000 | 1 | 536870912 | 0 | 0 | 8.246 ms | 60.28% | 4.883 ms | 0.74% | -3363.360 us | -40.79% | ESC[32mFASTESC[39m |
| LIST | DEVICE_BUFFER | NONE | 0 | 1 | 536870912 | 0 | 0 | 20.140 ms | 45.60% | 12.448 ms | 0.48% | -7692.423 us | -38.19% | ESC[32mFASTESC[39m |
| LIST | DEVICE_BUFFER | NONE | 1000 | 32 | 536870912 | 0 | 0 | 21.315 ms | 36.34% | 14.130 ms | 0.23% | -7185.127 us | -33.71% | ESC[32mFASTESC[39m |
| STRUCT | DEVICE_BUFFER | NONE | 0 | 1 | 536870912 | 0 | 0 | 26.064 ms | 53.16% | 49.879 ms | 38.93% | 23.815 ms | 91.37% | ESC[31mSLOWESC[39m |
| STRUCT | DEVICE_BUFFER | NONE | 1000 | 1 | 536870912 | 0 | 0 | 20.409 ms | 70.69% | 25.206 ms | 97.45% | 4.796 ms | 23.50% | ESC[34mSAMEESC[39m |
| STRUCT | DEVICE_BUFFER | NONE | 0 | 32 | 536870912 | 0 | 0 | 27.551 ms | 62.71% | 33.091 ms | 82.29% | 5.540 ms | 20.11% | ESC[34mSAMEESC[39m |
| STRUCT | DEVICE_BUFFER | NONE | 1000 | 32 | 536870912 | 0 | 0 | 23.070 ms | 76.04% | 26.259 ms | 101.00% | 3.189 ms | 13.83% | ESC[34mSAMEESC[39m |
# Summary
- Total Matches: 36
- Pass (diff <= min_noise): 7
- Unknown (infinite noise): 0
- Failure (diff > min_noise): 29 |
|
Confirm that this does not slow down the decode of dictionary-encoded string columns; i.e. kernels that use a lot of shared memory and benefit from having more cache lines available. |
|
See #23090 (comment) for a bit of discussion on some of the benchmarking I did of testing out different alignments to see if we could take advantage of some of the newer instructions available to use on newer hardware. |
@pmattione-nvidia do you have an idea of which benchmarks you'd want to run? I'm not that familiar with our Parquet reader code, so correct me if I'm wrong, but from a brief inspection I'm having trouble seeing where there could be concerning interactions. |
| constexpr int level_stage_bytes = 8 * 1024; | ||
| __shared__ __align__(16) uint8_t def_stage[level_stage_bytes]; | ||
| __shared__ __align__(16) uint8_t rep_stage[level_stage_bytes]; | ||
| using barrier_t = cuda::barrier<cuda::thread_scope_block>; |
There was a problem hiding this comment.
nit: Does decltype(block) allow us to tie the barrier scope to the block scope?
There was a problem hiding this comment.
Unfortunately no, CCCL doesn't currently provide a trait to derive thread_scope_block from the cg::thread_block type. Happy to revisit if CCCL adds such a trait.
- Remove smem_stage data member from rle_stream; use local variable in init() - Replace char storage + reinterpret_cast with direct __shared__ barrier declaration - Suppress nvcc static_var_with_dynamic_init for direct __shared__ barrier use
pmattione-nvidia
left a comment
pmattione-nvidia
left a comment
There was a problem hiding this comment.
Approved, contingent on addressing all of the other feedback of course.
Both decoders use the same 8 KiB staging buffer by serializing the async copies: rep is staged and decoded first, then the barrier is reinitialised and def reuses the same buffer. This saves 8 KiB of shared memory per block with no performance cost: benchmarking on an H100 shows the serialised path is never slower and is marginally faster on several types (BOOL8, STRUCT) where the reduced register/smem pressure improves occupancy. Addresses review comment from pmattione-nvidia.
kingcrimsontianyu
left a comment
kingcrimsontianyu
left a comment
There was a problem hiding this comment.
Looks good. Great work!
|
/merge |
7 participants
Description
The streams to be decoded are optionally prefetched into a shared memory buffer based on a size constraint. I did a benchmark sweep of smem sizes in powers of 2 from 2 to 32 and I found that 8KB was the sweet spot for the smem size based on our microbenchmarks. This gives ~6x speedups on the Parquet read decode microbenchmarks on H100s.
Checklist