From cf03e8a4a543d206e56ba7d03ce260ecb783d942 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Thu, 2 Jul 2026 01:51:35 +0000 Subject: [PATCH 1/8] perf(parquet): stage RLE level streams into shared memory for faster header parsing --- cpp/src/io/parquet/decode_preprocess.cu | 20 +++++++-- cpp/src/io/parquet/rle_stream.cuh | 58 +++++++++++++++++++++++-- 2 files changed, 72 insertions(+), 6 deletions(-) diff --git a/cpp/src/io/parquet/decode_preprocess.cu b/cpp/src/io/parquet/decode_preprocess.cu index 4aa40f3cd1ab..c9c1484dc8e7 100644 --- a/cpp/src/io/parquet/decode_preprocess.cu +++ b/cpp/src/io/parquet/decode_preprocess.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -418,6 +418,16 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size) rle_stream decoders[level_type::NUM_LEVEL_TYPES] = {{def_runs}, {rep_runs}}; + // Shared-memory staging scratch for the encoded level streams. Level streams + // for a page are usually small (definition/repetition levels are dominated by + // short RLE runs), and their serial run-header parse is latency-bound on + // dependent global loads. Staging the bytes into shared memory once removes + // that latency from fill_run_batch(). Streams larger than the per-stream + // budget fall back to parsing from global with no behavior change. + constexpr int level_stage_bytes = 8 * 1024; + __shared__ uint8_t def_stage[level_stage_bytes]; + __shared__ uint8_t rep_stage[level_stage_bytes]; + // Get the level decode buffers for this page auto* const def = reinterpret_cast(pp->lvl_decode_buf[level_type::DEFINITION]); auto* const rep = reinterpret_cast(pp->lvl_decode_buf[level_type::REPETITION]); @@ -434,14 +444,18 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size) s->abs_lvl_start[level_type::DEFINITION], s->abs_lvl_end[level_type::DEFINITION], def, - num_to_decode); + num_to_decode, + def_stage, + level_stage_bytes); } if (has_repetition) { decoders[level_type::REPETITION].init(s->col.level_bits[level_type::REPETITION], s->abs_lvl_start[level_type::REPETITION], s->abs_lvl_end[level_type::REPETITION], rep, - num_to_decode); + num_to_decode, + rep_stage, + level_stage_bytes); } block.sync(); diff --git a/cpp/src/io/parquet/rle_stream.cuh b/cpp/src/io/parquet/rle_stream.cuh index 4c826e1d5400..9200c651c681 100644 --- a/cpp/src/io/parquet/rle_stream.cuh +++ b/cpp/src/io/parquet/rle_stream.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -178,7 +178,19 @@ struct rle_stream { int fill_index; int decode_index; - __device__ rle_stream(rle_run* _runs) : runs(_runs) {} + // Optional shared-memory staging of the encoded byte stream. When init() is + // given a scratch buffer large enough to hold [start, end), the stream is + // copied into it once (block-cooperatively) and cur/end are rebased into + // shared memory. This turns the serial run-header parse that dominates + // fill_run_batch() from a chain of dependent L2 loads (~200 cyc each) into + // shared-memory loads (~30 cyc). It stages *raw encoded bytes*, so it is + // level_t- and level_bits-agnostic: definition/repetition levels, dictionary + // indices, and boolean streams all benefit with identical code. Streams that + // do not fit the budget transparently fall back to parsing from global. + uint8_t const* smem_stage; + int smem_stage_size; + + __device__ rle_stream(rle_run* _runs) : runs(_runs), smem_stage(nullptr), smem_stage_size(0) {} __device__ inline bool is_last_decode_warp(int warp_id) { @@ -189,7 +201,9 @@ struct rle_stream { uint8_t const* _start, uint8_t const* _end, level_t* _output, - int _total_values) + int _total_values, + uint8_t* _smem_stage = nullptr, + int _smem_stage_size = 0) { level_bits = _level_bits; cur = _start; @@ -203,6 +217,44 @@ struct rle_stream { cur_values = 0; fill_index = 0; decode_index = -1; // signals the first iteration. Nothing to decode. + + // Optionally stage the encoded stream into shared memory. init() is called + // by every thread in the block, so the copy below is block-cooperative. It + // deliberately performs no barrier of its own: every existing caller issues + // a block-wide sync immediately after init() (before the first parse in + // fill_run_batch/skip_runs), which is what publishes the staged bytes. + smem_stage = _smem_stage; + smem_stage_size = _smem_stage_size; + if (smem_stage != nullptr) { + int const len = static_cast(_end - _start); + if (len > 0 && len <= smem_stage_size) { + auto* const s_dst = _smem_stage; + int const t = threadIdx.x; + int const nthreads = blockDim.x; + // 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. + if ((reinterpret_cast(_start) & 15u) == 0) { + int const nvec = len >> 4; + auto const* g4 = reinterpret_cast(_start); + auto* const s4 = reinterpret_cast(s_dst); + for (int i = t; i < nvec; i += nthreads) { + s4[i] = g4[i]; + } + for (int i = (nvec << 4) + t; i < len; i += nthreads) { + s_dst[i] = _start[i]; + } + } else { + for (int i = t; i < len; i += nthreads) { + s_dst[i] = _start[i]; + } + } + // Rebase the parse cursor and end onto the shared copy. All downstream + // reads (get_rle_run_info, decode, skip_runs) follow cur/end and now hit + // shared memory with no other changes required. + cur = smem_stage; + end = smem_stage + len; + } + } } __device__ inline int get_rle_run_info(rle_run& run) From 72a3fa5101566b495d9ba526c5e79236632f237e Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Thu, 2 Jul 2026 19:38:39 +0000 Subject: [PATCH 2/8] Minor tweaks --- cpp/src/io/parquet/rle_stream.cuh | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/cpp/src/io/parquet/rle_stream.cuh b/cpp/src/io/parquet/rle_stream.cuh index 9200c651c681..bd56e08908f7 100644 --- a/cpp/src/io/parquet/rle_stream.cuh +++ b/cpp/src/io/parquet/rle_stream.cuh @@ -182,11 +182,11 @@ struct rle_stream { // given a scratch buffer large enough to hold [start, end), the stream is // copied into it once (block-cooperatively) and cur/end are rebased into // shared memory. This turns the serial run-header parse that dominates - // fill_run_batch() from a chain of dependent L2 loads (~200 cyc each) into - // shared-memory loads (~30 cyc). It stages *raw encoded bytes*, so it is - // level_t- and level_bits-agnostic: definition/repetition levels, dictionary - // indices, and boolean streams all benefit with identical code. Streams that - // do not fit the budget transparently fall back to parsing from global. + // fill_run_batch() from a chain of dependent L2 loads into shared-memory + // loads. It stages *raw encoded bytes*, so it is level_t- and + // level_bits-agnostic: definition/repetition levels, dictionary indices, and + // boolean streams all benefit with identical code. Streams that do not fit + // the budget transparently fall back to parsing from global. uint8_t const* smem_stage; int smem_stage_size; @@ -218,15 +218,15 @@ struct rle_stream { fill_index = 0; decode_index = -1; // signals the first iteration. Nothing to decode. - // Optionally stage the encoded stream into shared memory. init() is called - // by every thread in the block, so the copy below is block-cooperative. It - // deliberately performs no barrier of its own: every existing caller issues - // a block-wide sync immediately after init() (before the first parse in + // If smem staging is active, since init() is called by every thread in the + // block the copy below is block-cooperative. It deliberately performs no + // barrier of its own: every existing caller issues a block-wide sync + // immediately after init() (before the first parse in // fill_run_batch/skip_runs), which is what publishes the staged bytes. smem_stage = _smem_stage; smem_stage_size = _smem_stage_size; if (smem_stage != nullptr) { - int const len = static_cast(_end - _start); + auto const len = static_cast(_end - _start); if (len > 0 && len <= smem_stage_size) { auto* const s_dst = _smem_stage; int const t = threadIdx.x; From 6942aea1e351f63cf2a0702f421dc52fd1650d67 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Mon, 6 Jul 2026 23:30:35 +0000 Subject: [PATCH 3/8] Switch to using cuda::memcpy_async with a suitable barrier --- cpp/src/io/parquet/decode_preprocess.cu | 14 +++++++-- cpp/src/io/parquet/rle_stream.cuh | 41 +++++++++---------------- 2 files changed, 25 insertions(+), 30 deletions(-) diff --git a/cpp/src/io/parquet/decode_preprocess.cu b/cpp/src/io/parquet/decode_preprocess.cu index c9c1484dc8e7..73d0cb473bf4 100644 --- a/cpp/src/io/parquet/decode_preprocess.cu +++ b/cpp/src/io/parquet/decode_preprocess.cu @@ -14,6 +14,7 @@ #include #include +#include #include #include @@ -427,6 +428,11 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size) constexpr int level_stage_bytes = 8 * 1024; __shared__ uint8_t def_stage[level_stage_bytes]; __shared__ uint8_t rep_stage[level_stage_bytes]; + using barrier_t = cuda::barrier; + __shared__ alignas(barrier_t) char copy_barrier_storage[sizeof(barrier_t)]; + auto* copy_barrier = reinterpret_cast(copy_barrier_storage); + if (t == 0) { init(copy_barrier, block.size()); } + block.sync(); // Get the level decode buffers for this page auto* const def = reinterpret_cast(pp->lvl_decode_buf[level_type::DEFINITION]); @@ -446,7 +452,8 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size) def, num_to_decode, def_stage, - level_stage_bytes); + level_stage_bytes, + copy_barrier); } if (has_repetition) { decoders[level_type::REPETITION].init(s->col.level_bits[level_type::REPETITION], @@ -455,9 +462,10 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size) rep, num_to_decode, rep_stage, - level_stage_bytes); + level_stage_bytes, + copy_barrier); } - block.sync(); + copy_barrier->arrive_and_wait(); // Decode levels for this page up to the last row needed. // If skipping the first rows, we still need to decode their levels. diff --git a/cpp/src/io/parquet/rle_stream.cuh b/cpp/src/io/parquet/rle_stream.cuh index bd56e08908f7..14d8ab8b5e5e 100644 --- a/cpp/src/io/parquet/rle_stream.cuh +++ b/cpp/src/io/parquet/rle_stream.cuh @@ -10,6 +10,9 @@ #include #include +#include +#include + namespace cudf::io::parquet::detail { template @@ -202,8 +205,9 @@ struct rle_stream { uint8_t const* _end, level_t* _output, int _total_values, - uint8_t* _smem_stage = nullptr, - int _smem_stage_size = 0) + uint8_t* _smem_stage = nullptr, + int _smem_stage_size = 0, + cuda::barrier* _copy_barrier = nullptr) { level_bits = _level_bits; cur = _start; @@ -218,36 +222,19 @@ struct rle_stream { fill_index = 0; decode_index = -1; // signals the first iteration. Nothing to decode. - // If smem staging is active, since init() is called by every thread in the - // block the copy below is block-cooperative. It deliberately performs no - // barrier of its own: every existing caller issues a block-wide sync - // immediately after init() (before the first parse in - // fill_run_batch/skip_runs), which is what publishes the staged bytes. + // If smem staging is active, use cuda::memcpy_async for a + // block-cooperative global-to-shared copy that automatically dispatches to + // the best copy path (cp.async, cp.async.bulk, or TMA) depending on the + // hardware. Callers must provide a copy_barrier when using smem staging, + // and must issue copy_barrier->arrive_and_wait() after init() to complete + // the async copy. smem_stage = _smem_stage; smem_stage_size = _smem_stage_size; if (smem_stage != nullptr) { auto const len = static_cast(_end - _start); if (len > 0 && len <= smem_stage_size) { - auto* const s_dst = _smem_stage; - int const t = threadIdx.x; - int const nthreads = blockDim.x; - // 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. - if ((reinterpret_cast(_start) & 15u) == 0) { - int const nvec = len >> 4; - auto const* g4 = reinterpret_cast(_start); - auto* const s4 = reinterpret_cast(s_dst); - for (int i = t; i < nvec; i += nthreads) { - s4[i] = g4[i]; - } - for (int i = (nvec << 4) + t; i < len; i += nthreads) { - s_dst[i] = _start[i]; - } - } else { - for (int i = t; i < len; i += nthreads) { - s_dst[i] = _start[i]; - } - } + auto group = cooperative_groups::this_thread_block(); + cuda::memcpy_async(group, _smem_stage, _start, static_cast(len), *_copy_barrier); // Rebase the parse cursor and end onto the shared copy. All downstream // reads (get_rle_run_info, decode, skip_runs) follow cur/end and now hit // shared memory with no other changes required. From 26bf9600365cfa08d561f7bc82f5d075b2420677 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Tue, 7 Jul 2026 05:50:47 +0000 Subject: [PATCH 4/8] Align shared staging buffers to 16 bytes for vectorized async copy --- cpp/src/io/parquet/decode_preprocess.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/src/io/parquet/decode_preprocess.cu b/cpp/src/io/parquet/decode_preprocess.cu index 73d0cb473bf4..45fcfd5b211e 100644 --- a/cpp/src/io/parquet/decode_preprocess.cu +++ b/cpp/src/io/parquet/decode_preprocess.cu @@ -426,8 +426,8 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size) // that latency from fill_run_batch(). Streams larger than the per-stream // budget fall back to parsing from global with no behavior change. constexpr int level_stage_bytes = 8 * 1024; - __shared__ uint8_t def_stage[level_stage_bytes]; - __shared__ uint8_t rep_stage[level_stage_bytes]; + __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; __shared__ alignas(barrier_t) char copy_barrier_storage[sizeof(barrier_t)]; auto* copy_barrier = reinterpret_cast(copy_barrier_storage); From be363266243d9da5224a9fb47dbd7570be5b41f4 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Tue, 7 Jul 2026 05:50:54 +0000 Subject: [PATCH 5/8] Use cuda::std::distance for length computation in rle_stream --- cpp/src/io/parquet/rle_stream.cuh | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cpp/src/io/parquet/rle_stream.cuh b/cpp/src/io/parquet/rle_stream.cuh index 14d8ab8b5e5e..328f649030f8 100644 --- a/cpp/src/io/parquet/rle_stream.cuh +++ b/cpp/src/io/parquet/rle_stream.cuh @@ -12,6 +12,7 @@ #include #include +#include namespace cudf::io::parquet::detail { @@ -231,7 +232,7 @@ struct rle_stream { smem_stage = _smem_stage; smem_stage_size = _smem_stage_size; if (smem_stage != nullptr) { - auto const len = static_cast(_end - _start); + auto const len = static_cast(cuda::std::distance(_start, _end)); if (len > 0 && len <= smem_stage_size) { auto group = cooperative_groups::this_thread_block(); cuda::memcpy_async(group, _smem_stage, _start, static_cast(len), *_copy_barrier); From 499252f6eb34b5408d9299f5b958a77820d0fc34 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Tue, 7 Jul 2026 19:33:40 +0000 Subject: [PATCH 6/8] fix(io): address wence- review comments on RLE stream staging --- cpp/src/io/parquet/decode_fixed.cu | 5 +++-- cpp/src/io/parquet/decode_preprocess.cu | 16 ++++++++-------- cpp/src/io/parquet/rle_stream.cuh | 16 +++++++++------- 3 files changed, 20 insertions(+), 17 deletions(-) diff --git a/cpp/src/io/parquet/decode_fixed.cu b/cpp/src/io/parquet/decode_fixed.cu index c2adc419434b..45a7aeb4ec6d 100644 --- a/cpp/src/io/parquet/decode_fixed.cu +++ b/cpp/src/io/parquet/decode_fixed.cu @@ -1081,7 +1081,7 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size_t, 8) rle_stream dict_stream{dict_runs}; if constexpr (has_dict_t) { dict_stream.init( - s->dict_bits, s->data_start, s->data_end, sb->dict_idx, s->page.num_input_values); + block, s->dict_bits, s->data_start, s->data_end, sb->dict_idx, s->page.num_input_values); } // Use dictionary stream memory for bools @@ -1089,7 +1089,8 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size_t, 8) bool bools_are_rle_stream = (s->dict_run == 0); if constexpr (has_bools_t) { if (bools_are_rle_stream) { - bool_stream.init(1, s->data_start, s->data_end, sb->dict_idx, s->page.num_input_values); + bool_stream.init( + block, 1, s->data_start, s->data_end, sb->dict_idx, s->page.num_input_values); } } block.sync(); diff --git a/cpp/src/io/parquet/decode_preprocess.cu b/cpp/src/io/parquet/decode_preprocess.cu index 45fcfd5b211e..43e1241b727b 100644 --- a/cpp/src/io/parquet/decode_preprocess.cu +++ b/cpp/src/io/parquet/decode_preprocess.cu @@ -425,13 +425,13 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size) // dependent global loads. Staging the bytes into shared memory once removes // that latency from fill_run_batch(). Streams larger than the per-stream // budget fall back to parsing from global with no behavior change. - 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 rle_stream_t = rle_stream; + __shared__ __align__(16) uint8_t def_stage[rle_stream_t::smem_stage_size]; + __shared__ __align__(16) uint8_t rep_stage[rle_stream_t::smem_stage_size]; using barrier_t = cuda::barrier; __shared__ alignas(barrier_t) char copy_barrier_storage[sizeof(barrier_t)]; auto* copy_barrier = reinterpret_cast(copy_barrier_storage); - if (t == 0) { init(copy_barrier, block.size()); } + cg::invoke_one(block, [&]() { init(copy_barrier, block.size()); }); block.sync(); // Get the level decode buffers for this page @@ -446,23 +446,23 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size) // Initialize the stream decoders bool const process_nulls = should_process_nulls(s); if (process_nulls) { - decoders[level_type::DEFINITION].init(s->col.level_bits[level_type::DEFINITION], + decoders[level_type::DEFINITION].init(block, + s->col.level_bits[level_type::DEFINITION], s->abs_lvl_start[level_type::DEFINITION], s->abs_lvl_end[level_type::DEFINITION], def, num_to_decode, def_stage, - level_stage_bytes, copy_barrier); } if (has_repetition) { - decoders[level_type::REPETITION].init(s->col.level_bits[level_type::REPETITION], + decoders[level_type::REPETITION].init(block, + s->col.level_bits[level_type::REPETITION], s->abs_lvl_start[level_type::REPETITION], s->abs_lvl_end[level_type::REPETITION], rep, num_to_decode, rep_stage, - level_stage_bytes, copy_barrier); } copy_barrier->arrive_and_wait(); diff --git a/cpp/src/io/parquet/rle_stream.cuh b/cpp/src/io/parquet/rle_stream.cuh index 328f649030f8..dffc8e008e17 100644 --- a/cpp/src/io/parquet/rle_stream.cuh +++ b/cpp/src/io/parquet/rle_stream.cuh @@ -13,6 +13,7 @@ #include #include #include +#include namespace cudf::io::parquet::detail { @@ -192,22 +193,23 @@ struct rle_stream { // boolean streams all benefit with identical code. Streams that do not fit // the budget transparently fall back to parsing from global. uint8_t const* smem_stage; - int smem_stage_size; + static constexpr int smem_stage_size = 8 * 1024; - __device__ rle_stream(rle_run* _runs) : runs(_runs), smem_stage(nullptr), smem_stage_size(0) {} + __device__ rle_stream(rle_run* _runs) : runs(_runs), smem_stage(nullptr) {} __device__ inline bool is_last_decode_warp(int warp_id) { return warp_id == num_rle_stream_decode_warps; } - __device__ void init(int _level_bits, + template + __device__ void init(Group const& group, + int _level_bits, uint8_t const* _start, uint8_t const* _end, level_t* _output, int _total_values, uint8_t* _smem_stage = nullptr, - int _smem_stage_size = 0, cuda::barrier* _copy_barrier = nullptr) { level_bits = _level_bits; @@ -229,12 +231,12 @@ struct rle_stream { // hardware. Callers must provide a copy_barrier when using smem staging, // and must issue copy_barrier->arrive_and_wait() after init() to complete // the async copy. - smem_stage = _smem_stage; - smem_stage_size = _smem_stage_size; + smem_stage = _smem_stage != nullptr + ? static_cast(cuda::std::assume_aligned<16>(_smem_stage)) + : nullptr; if (smem_stage != nullptr) { auto const len = static_cast(cuda::std::distance(_start, _end)); if (len > 0 && len <= smem_stage_size) { - auto group = cooperative_groups::this_thread_block(); cuda::memcpy_async(group, _smem_stage, _start, static_cast(len), *_copy_barrier); // Rebase the parse cursor and end onto the shared copy. All downstream // reads (get_rle_run_info, decode, skip_runs) follow cur/end and now hit From b6e114950a7aa49be01e9b59b12d3bf18a061641 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Wed, 8 Jul 2026 15:38:32 +0000 Subject: [PATCH 7/8] fix(io): address kingcrimsontianyu review comments on RLE stream staging - 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 --- cpp/src/io/parquet/decode_preprocess.cu | 13 ++++++------- cpp/src/io/parquet/rle_stream.cuh | 10 ++++------ 2 files changed, 10 insertions(+), 13 deletions(-) diff --git a/cpp/src/io/parquet/decode_preprocess.cu b/cpp/src/io/parquet/decode_preprocess.cu index 43e1241b727b..7233b8a673a4 100644 --- a/cpp/src/io/parquet/decode_preprocess.cu +++ b/cpp/src/io/parquet/decode_preprocess.cu @@ -379,6 +379,7 @@ CUDF_KERNEL void __launch_bounds__(preprocess_block_size) * @param min_row Minimum row index to read * @param num_rows Number of rows to read starting from min_row */ +#pragma nv_diag_suppress static_var_with_dynamic_init template CUDF_KERNEL void __launch_bounds__(level_decode_block_size) preprocess_levels_kernel(PageInfo* pages, @@ -428,10 +429,8 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size) using rle_stream_t = rle_stream; __shared__ __align__(16) uint8_t def_stage[rle_stream_t::smem_stage_size]; __shared__ __align__(16) uint8_t rep_stage[rle_stream_t::smem_stage_size]; - using barrier_t = cuda::barrier; - __shared__ alignas(barrier_t) char copy_barrier_storage[sizeof(barrier_t)]; - auto* copy_barrier = reinterpret_cast(copy_barrier_storage); - cg::invoke_one(block, [&]() { init(copy_barrier, block.size()); }); + __shared__ cuda::barrier copy_barrier; + cg::invoke_one(block, [&]() { init(©_barrier, block.size()); }); block.sync(); // Get the level decode buffers for this page @@ -453,7 +452,7 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size) def, num_to_decode, def_stage, - copy_barrier); + ©_barrier); } if (has_repetition) { decoders[level_type::REPETITION].init(block, @@ -463,9 +462,9 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size) rep, num_to_decode, rep_stage, - copy_barrier); + ©_barrier); } - copy_barrier->arrive_and_wait(); + copy_barrier.arrive_and_wait(); // Decode levels for this page up to the last row needed. // If skipping the first rows, we still need to decode their levels. diff --git a/cpp/src/io/parquet/rle_stream.cuh b/cpp/src/io/parquet/rle_stream.cuh index dffc8e008e17..0f4731003621 100644 --- a/cpp/src/io/parquet/rle_stream.cuh +++ b/cpp/src/io/parquet/rle_stream.cuh @@ -192,10 +192,9 @@ struct rle_stream { // level_bits-agnostic: definition/repetition levels, dictionary indices, and // boolean streams all benefit with identical code. Streams that do not fit // the budget transparently fall back to parsing from global. - uint8_t const* smem_stage; static constexpr int smem_stage_size = 8 * 1024; - __device__ rle_stream(rle_run* _runs) : runs(_runs), smem_stage(nullptr) {} + __device__ rle_stream(rle_run* _runs) : runs(_runs) {} __device__ inline bool is_last_decode_warp(int warp_id) { @@ -231,10 +230,9 @@ struct rle_stream { // hardware. Callers must provide a copy_barrier when using smem staging, // and must issue copy_barrier->arrive_and_wait() after init() to complete // the async copy. - smem_stage = _smem_stage != nullptr - ? static_cast(cuda::std::assume_aligned<16>(_smem_stage)) - : nullptr; - if (smem_stage != nullptr) { + if (_smem_stage != nullptr) { + auto* const smem_stage = + static_cast(cuda::std::assume_aligned<16>(_smem_stage)); auto const len = static_cast(cuda::std::distance(_start, _end)); if (len > 0 && len <= smem_stage_size) { cuda::memcpy_async(group, _smem_stage, _start, static_cast(len), *_copy_barrier); From 633468486db796d845c5fe3ed64d21549b094952 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Wed, 8 Jul 2026 21:43:05 +0000 Subject: [PATCH 8/8] fix(io): use single shared staging buffer for def/rep level decode 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. --- cpp/src/io/parquet/decode_preprocess.cu | 39 +++++++++++++------------ 1 file changed, 20 insertions(+), 19 deletions(-) diff --git a/cpp/src/io/parquet/decode_preprocess.cu b/cpp/src/io/parquet/decode_preprocess.cu index 7233b8a673a4..13dad751c9ca 100644 --- a/cpp/src/io/parquet/decode_preprocess.cu +++ b/cpp/src/io/parquet/decode_preprocess.cu @@ -427,11 +427,8 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size) // that latency from fill_run_batch(). Streams larger than the per-stream // budget fall back to parsing from global with no behavior change. using rle_stream_t = rle_stream; - __shared__ __align__(16) uint8_t def_stage[rle_stream_t::smem_stage_size]; - __shared__ __align__(16) uint8_t rep_stage[rle_stream_t::smem_stage_size]; + __shared__ __align__(16) uint8_t stage[rle_stream_t::smem_stage_size]; __shared__ cuda::barrier copy_barrier; - cg::invoke_one(block, [&]() { init(©_barrier, block.size()); }); - block.sync(); // Get the level decode buffers for this page auto* const def = reinterpret_cast(pp->lvl_decode_buf[level_type::DEFINITION]); @@ -444,39 +441,43 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size) // Initialize the stream decoders bool const process_nulls = should_process_nulls(s); - if (process_nulls) { - decoders[level_type::DEFINITION].init(block, - s->col.level_bits[level_type::DEFINITION], - s->abs_lvl_start[level_type::DEFINITION], - s->abs_lvl_end[level_type::DEFINITION], - def, - num_to_decode, - def_stage, - ©_barrier); - } if (has_repetition) { + cg::invoke_one(block, [&]() { init(©_barrier, block.size()); }); + block.sync(); decoders[level_type::REPETITION].init(block, s->col.level_bits[level_type::REPETITION], s->abs_lvl_start[level_type::REPETITION], s->abs_lvl_end[level_type::REPETITION], rep, num_to_decode, - rep_stage, + stage, ©_barrier); + copy_barrier.arrive_and_wait(); + decoders[level_type::REPETITION].decode_next(t, num_to_decode); } - copy_barrier.arrive_and_wait(); // Decode levels for this page up to the last row needed. // If skipping the first rows, we still need to decode their levels. // This is because we need to determine the number of non-null values we skipped. // Note that for lists we haven't computed skipped_leaf_values yet; this is used as input for // that. - if (has_repetition) { decoders[level_type::REPETITION].decode_next(t, num_to_decode); } - // Must sync as shared variables in decode_next() are shared between decoders!! block.sync(); - if (process_nulls) { decoders[level_type::DEFINITION].decode_next(t, num_to_decode); } + if (process_nulls) { + cg::invoke_one(block, [&]() { init(©_barrier, block.size()); }); + block.sync(); + decoders[level_type::DEFINITION].init(block, + s->col.level_bits[level_type::DEFINITION], + s->abs_lvl_start[level_type::DEFINITION], + s->abs_lvl_end[level_type::DEFINITION], + def, + num_to_decode, + stage, + ©_barrier); + copy_barrier.arrive_and_wait(); + decoders[level_type::DEFINITION].decode_next(t, num_to_decode); + } } } // anonymous namespace