From 9e67823138583b2b99206913e3e37532bbac44f7 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Thu, 16 Jul 2026 21:32:53 +0000 Subject: [PATCH] perf(parquet): decode def and rep levels concurrently via 2D grid --- cpp/src/io/parquet/decode_preprocess.cu | 74 ++++++++++--------------- 1 file changed, 28 insertions(+), 46 deletions(-) diff --git a/cpp/src/io/parquet/decode_preprocess.cu b/cpp/src/io/parquet/decode_preprocess.cu index 85e4a6f0d989..fc8a4884986a 100644 --- a/cpp/src/io/parquet/decode_preprocess.cu +++ b/cpp/src/io/parquet/decode_preprocess.cu @@ -393,7 +393,7 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size) page_state_s* const s = &state_g; auto const block = cg::this_thread_block(); - int const page_idx = cg::this_grid().block_rank(); + int const page_idx = blockIdx.x; int const t = block.thread_rank(); PageInfo* pp = &pages[page_idx]; @@ -409,17 +409,20 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size) // whether or not we have repetition levels (lists) bool const has_repetition = chunks[pp->chunk_idx].max_level[level_type::REPETITION] > 0; + // Grid-expansion dispatch: each page is decoded by two blocks. + // blockIdx.y maps to the level_type enum (parquet_gpu.hpp:184-189). + int const stream_id = blockIdx.y; + static_assert(level_type::DEFINITION == 0 && level_type::REPETITION == 1, + "grid dispatch assumes DEFINITION=0, REPETITION=1"); + // the required number of runs in shared memory we will need to provide the // rle_stream object constexpr int rle_run_buffer_size = rle_stream_required_run_buffer_size(); - // the level stream decoders. max_output_values is max to remove rolling buffer - __shared__ rle_run def_runs[rle_run_buffer_size]; - __shared__ rle_run rep_runs[rle_run_buffer_size]; + // One rle_run buffer per block (previously two: def_runs + rep_runs) + __shared__ rle_run runs[rle_run_buffer_size]; static constexpr int max_output_values = cuda::std::numeric_limits::max(); - 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 @@ -431,54 +434,33 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size) __shared__ __align__(16) uint8_t stage[rle_stream_t::smem_stage_size]; __shared__ cuda::barrier copy_barrier; - // 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]); - // Determine how many values need to be decoded size_t const num_to_decode = precompute_page_num_values_in_range(*pp, chunks[pp->chunk_idx], min_row, num_rows); if (num_to_decode == 0) { return; } - // Initialize the stream decoders + // Skip if this block's stream is absent for this page. bool const process_nulls = should_process_nulls(s); - 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, - stage, - ©_barrier); - copy_barrier.arrive_and_wait(); - decoders[level_type::REPETITION].decode_next(t, num_to_decode); - } + if (stream_id == level_type::REPETITION && !has_repetition) { return; } + if (stream_id == level_type::DEFINITION && !process_nulls) { return; } - // 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. - // Must sync as shared variables in decode_next() are shared between decoders!! + // Dispatch to the level stream this block owns. + auto* const out = reinterpret_cast(pp->lvl_decode_buf[stream_id]); + + cg::invoke_one(block, [&]() { init(©_barrier, block.size()); }); block.sync(); - 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); - } + rle_stream decoder{runs}; + decoder.init(block, + s->col.level_bits[stream_id], + s->abs_lvl_start[stream_id], + s->abs_lvl_end[stream_id], + out, + num_to_decode, + stage, + ©_barrier); + copy_barrier.arrive_and_wait(); + decoder.decode_next(t, num_to_decode); } } // anonymous namespace @@ -534,7 +516,7 @@ void preprocess_levels(cudf::detail::hostdevice_span pages, if (pages.size() == 0) { return; } dim3 dim_block(level_decode_block_size, 1); - dim3 dim_grid(pages.size(), 1); // 1 threadblock per page + dim3 dim_grid(pages.size(), 2); // 2 threadblocks per page: one per level stream if (level_type_size == 1) { preprocess_levels_kernel