Skip to content

perf(parquet): decode def and rep levels concurrently via 2D grid #7

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
vyasr wants to merge 1 commit into
base: opt/rle-chunked-expand
Choose a base branch
from
Open

perf(parquet): decode def and rep levels concurrently via 2D grid #7

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 28 additions & 46 deletions cpp/src/io/parquet/decode_preprocess.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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];

Expand All @@ -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<level_decode_block_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<int>::max();
rle_stream<level_t, level_decode_block_size, max_output_values, true>
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
Expand All @@ -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<cuda::thread_scope_block> copy_barrier;

// Get the level decode buffers for this page
auto* const def = reinterpret_cast<level_t*>(pp->lvl_decode_buf[level_type::DEFINITION]);
auto* const rep = reinterpret_cast<level_t*>(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(&copy_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,
&copy_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<level_t*>(pp->lvl_decode_buf[stream_id]);

cg::invoke_one(block, [&]() { init(&copy_barrier, block.size()); });
block.sync();

if (process_nulls) {
cg::invoke_one(block, [&]() { init(&copy_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,
&copy_barrier);
copy_barrier.arrive_and_wait();
decoders[level_type::DEFINITION].decode_next(t, num_to_decode);
}
rle_stream<level_t, level_decode_block_size, max_output_values, true> 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,
&copy_barrier);
copy_barrier.arrive_and_wait();
decoder.decode_next(t, num_to_decode);
}

} // anonymous namespace
Expand Down Expand Up @@ -534,7 +516,7 @@ void preprocess_levels(cudf::detail::hostdevice_span<PageInfo> 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<uint8_t, level_decode_block_size>
Expand Down
Loading