diff --git a/cpp/src/io/parquet/decode_preprocess.cu b/cpp/src/io/parquet/decode_preprocess.cu index f9f886cf9b02..cc4d21796039 100644 --- a/cpp/src/io/parquet/decode_preprocess.cu +++ b/cpp/src/io/parquet/decode_preprocess.cu @@ -409,17 +409,13 @@ 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; - // 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]; + // logic from the decode step. The chunked-expand rle_stream does not need a + // shared-memory ring buffer of run headers; it parses runs directly into + // per-chunk tables, so we default-construct the decoders here. static constexpr int max_output_values = cuda::std::numeric_limits::max(); - rle_stream - decoders[level_type::NUM_LEVEL_TYPES] = {{def_runs}, {rep_runs}}; + using decoder_stream_t = rle_stream_chunked; + decoder_stream_t decoders[level_type::NUM_LEVEL_TYPES] = {}; // Shared-memory staging scratch for the encoded level streams. Level streams // for a page are usually small (definition/repetition levels are dominated by @@ -427,8 +423,7 @@ 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. - using rle_stream_t = rle_stream; - __shared__ __align__(16) uint8_t stage[rle_stream_t::smem_stage_size]; + __shared__ __align__(16) uint8_t stage[decoder_stream_t::smem_stage_size]; __shared__ cuda::barrier copy_barrier; // Get the level decode buffers for this page @@ -452,7 +447,8 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size) rep, num_to_decode, stage, - ©_barrier); + ©_barrier, + decoder_stream_t::smem_stage_size); copy_barrier.arrive_and_wait(); decoders[level_type::REPETITION].decode_next(t, num_to_decode); } @@ -475,7 +471,8 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size) def, num_to_decode, stage, - ©_barrier); + ©_barrier, + decoder_stream_t::smem_stage_size); copy_barrier.arrive_and_wait(); decoders[level_type::DEFINITION].decode_next(t, num_to_decode); } diff --git a/cpp/src/io/parquet/parquet_gpu.hpp b/cpp/src/io/parquet/parquet_gpu.hpp index 46c34c46e3af..b89ecbaae7f6 100644 --- a/cpp/src/io/parquet/parquet_gpu.hpp +++ b/cpp/src/io/parquet/parquet_gpu.hpp @@ -22,6 +22,7 @@ #include #include +#include #include #include @@ -45,9 +46,12 @@ constexpr int LEVEL_DECODE_BUF_SIZE = 2048; template CUDF_HOST_DEVICE constexpr int rolling_index(int index) { - // Cannot divide by 0. But `rolling_size` will be 0 for unused arrays, so this case will never - // actual be executed. - if constexpr (rolling_size == 0) { + // rolling_size == 0 marks unused arrays (never actually indexed). + // rolling_size == INT_MAX is used by paths that treat the output buffer as + // non-rolling (e.g. chunked-expand level decoding), and we short-circuit + // both cases to avoid a modulo (which for non-power-of-two divisors would + // compile to a real integer division on device). + if constexpr (rolling_size == 0 || rolling_size == cuda::std::numeric_limits::max()) { return index; } else { return index % rolling_size; diff --git a/cpp/src/io/parquet/rle_stream.cuh b/cpp/src/io/parquet/rle_stream.cuh index 0f4731003621..f98f1fd00325 100644 --- a/cpp/src/io/parquet/rle_stream.cuh +++ b/cpp/src/io/parquet/rle_stream.cuh @@ -5,6 +5,7 @@ #pragma once +#include "io/utilities/block_utils.cuh" #include "parquet_gpu.hpp" #include @@ -12,11 +13,15 @@ #include #include +#include #include #include +#include namespace cudf::io::parquet::detail { +namespace cg = cooperative_groups; + template __device__ constexpr int rle_stream_required_run_buffer_size() { @@ -154,8 +159,29 @@ struct rle_run { int remaining; // number of output items remaining to be decoded }; +// Controls the number of run headers parsed per chunk in the chunked-expand path. +// SMEM cost is (2 * max_runs_per_chunk + 1) * 4 bytes. Increasing max_runs_per_chunk reduces the +// number of outer-loop iterations in decode_next_chunked (each with a serial +// header-parse phase and a group.sync() at the end), but competes with occupancy in +// preprocess_levels_kernel. +// +// 1024 was chosen empirically for sm_80+. Sweeps of {256, 512, 1024, 2048, +// 4096} on A100, H100, and B200 all showed 1024 either as the numerical +// optimum or within noise of it. The delta over 512 was small (a few percent) +// in every case. sm_70/sm_75 stay at 512, because 1024 does not fit the +// preprocess_levels_kernel SMEM budget on those older arches. +#if __CUDA_ARCH__ >= 800 +static constexpr int max_runs_per_chunk = 1024; +#else +static constexpr int max_runs_per_chunk = 512; +#endif + // a stream of rle_runs -template +template struct rle_stream { static constexpr int num_rle_stream_decode_threads = decode_threads; // the -1 here is for the look-ahead warp that fills in the list of runs to be decoded @@ -167,7 +193,22 @@ struct rle_stream { static constexpr int run_buffer_size = rle_stream_required_run_buffer_size(); + // Bit packing of `run_desc` (chunked-expand path). The 32-bit descriptor + // stores a byte offset into the encoded level stream (`cur - s_start`) in + // its low 31 bits and uses the top bit as a flag: 1 = literal (bit-packed) + // run, 0 = RLE (repeated-value) run. Keeping both in one word lets Phase 1 + // publish a single 32-bit value per run into shared memory and lets Phase 2 + // dispatch on the flag without a second load. + // + // Invariant enforced by `cudf_assert` at parse time: (cur - s_start) fits in + // 31 bits, i.e. the encoded level stream for a single Parquet page is < 2 + // GiB. Parquet page payloads are orders of magnitude smaller than this in + // practice. + static constexpr uint32_t run_desc_literal_flag = 1u << 31; + static constexpr uint32_t run_desc_offset_mask = 0x7fffffffu; + int level_bits; + uint8_t const* s_start; uint8_t const* cur; uint8_t const* end; @@ -192,9 +233,22 @@ 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. - static constexpr int smem_stage_size = 8 * 1024; - - __device__ rle_stream(rle_run* _runs) : runs(_runs) {} + static constexpr int smem_stage_size = smem_stage_size_bytes; + + // Ring-mode streams need a shared-memory ring buffer of run headers to + // coordinate the producer/consumer warps in decode_next_ring. Chunked-expand + // streams parse run headers directly into per-chunk shared tables and never + // touch `runs`, so we forbid supplying one to catch accidental waste. + __device__ rle_stream(rle_run* _runs) + requires(!use_chunked_expand) + : runs(_runs) + { + } + __device__ rle_stream() + requires(use_chunked_expand) + : runs(nullptr) + { + } __device__ inline bool is_last_decode_warp(int warp_id) { @@ -209,11 +263,15 @@ struct rle_stream { level_t* _output, int _total_values, uint8_t* _smem_stage = nullptr, - cuda::barrier* _copy_barrier = nullptr) + cuda::barrier* _copy_barrier = nullptr, + int stage_capacity = smem_stage_size) { level_bits = _level_bits; - cur = _start; - end = _end; + // s_start is set below after any smem-staging rebase, so downstream code + // that stores offsets relative to s_start (chunked-expand meta) works + // uniformly whether cur points into global or shared memory. + cur = _start; + end = _end; output = _output; @@ -224,6 +282,8 @@ struct rle_stream { fill_index = 0; decode_index = -1; // signals the first iteration. Nothing to decode. + cudf_assert(stage_capacity >= 0 and stage_capacity <= smem_stage_size); + // 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 @@ -234,7 +294,7 @@ struct rle_stream { 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) { + if (len > 0 && len <= stage_capacity) { 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 @@ -243,6 +303,9 @@ struct rle_stream { end = smem_stage + len; } } + // Anchor s_start to the (possibly rebased) cur so chunked-expand meta + // offsets index into the same memory space that the parse cursor uses. + s_start = cur; } __device__ inline int get_rle_run_info(rle_run& run) @@ -289,49 +352,42 @@ struct rle_stream { } } - __device__ inline int decode_next(int t, int count) + template + __device__ inline int decode_next_ring(Group const& group, int count) { int const output_count = min(count, total_values - cur_values); - // special case. if level_bits == 0, just return all zeros. this should tremendously speed up - // a very common case: columns with no nulls, especially if they are non-nested - if (level_bits == 0) { - int written = 0; - while (written < output_count) { - int const batch_size = min(num_rle_stream_decode_threads, output_count - written); - if (t < batch_size) { output[rolling_index(written + t)] = 0; } - written += batch_size; - } - cur_values += output_count; - return output_count; - } - // otherwise, full decode. - int const warp_id = t / cudf::detail::warp_size; + auto const warp = cg::tiled_partition(group); + int const warp_id = warp.meta_group_rank(); int const warp_decode_id = warp_id - 1; - int const warp_lane = t % cudf::detail::warp_size; + int const warp_lane = warp.thread_rank(); __shared__ int values_processed_shared; __shared__ int decode_index_shared; __shared__ int fill_index_shared; - if (t == 0) { + // Do not use cg::invoke_one here: rle_stream member state is per-thread, + // so persistent state must be owned by a stable, well-defined thread. + if (group.thread_rank() == 0) { values_processed_shared = 0; decode_index_shared = decode_index; fill_index_shared = fill_index; } - __syncthreads(); + group.sync(); fill_index = fill_index_shared; - do { + while (true) { // protect against threads advancing past the end of this loop // and updating shared variables. - __syncthreads(); + group.sync(); // warp 0 reads ahead and fills `runs` array to be decoded by remaining warps. if (warp_id == 0) { // fill the next set of runs. fill_runs will generally be the bottleneck for any // kernel that uses an rle_stream. + // Do not use cg::invoke_one here: fill_run_batch() advances per-thread + // stream cursors, so the ring producer must always be lane 0. if (warp_lane == 0) { fill_run_batch(); if (decode_index == -1) { @@ -345,7 +401,7 @@ struct rle_stream { // remaining warps decode the runs, starting on the second iteration of this. the pipeline of // runs is also persistent across calls to decode_next, so on the second call to decode_next, // this branch will start doing work immediately. - // do/while loop (decode_index == -1 means "first iteration", so we should skip decoding) + // decode_index == -1 means "first iteration", so we should skip decoding. else if (decode_index >= 0 && decode_index + warp_decode_id < fill_index) { int const run_index = decode_index + warp_decode_id; auto& run = runs[rolling_index(run_index)]; @@ -375,7 +431,7 @@ struct rle_stream { level_bits, warp_lane); - __syncwarp(); + warp.sync(); if (warp_lane == 0) { // after writing this batch, are we at the end of the output buffer? auto const at_end = ((last_run_pos + batch_len - cur_values) == output_count); @@ -393,10 +449,11 @@ struct rle_stream { } } } - __syncthreads(); + group.sync(); decode_index = decode_index_shared; fill_index = fill_index_shared; - } while (values_processed_shared < output_count); + if (values_processed_shared >= output_count) { break; } + } cur_values += values_processed_shared; @@ -404,6 +461,265 @@ struct rle_stream { return values_processed_shared; } + /* Alternate decode path used when `use_chunked_expand` is true. + * + * Instead of the ring-buffer producer/consumer model in decode_next_ring + * (one warp parses run headers, other warps expand one run each), thread 0 + * parses up to `max_runs_per_chunk` headers up-front into shared-memory tables + * (chunk_out_off / chunk_meta), and then *all* warps cooperatively expand a + * slice of the concatenated output range using binary search. This keeps + * every warp busy even when runs are highly non-uniform in size, at the + * cost of an extra intra-block sync per chunk. + * + * The current sole caller passes max_output_values = INT_MAX, so a single + * RLE run can never exceed the requested `count`. A cudf_assert in the + * header-parse loop enforces this invariant; a future caller that needs to + * split runs across calls must restore the partial-run resume machinery + * removed in this commit. + */ + template + __device__ inline int decode_next_chunked(Group const& group, int count) + { + int const output_count = min(count, total_values - cur_values); + + // Per-chunk shared-memory scratch. `chunk_out_off[i]` is the exclusive + // prefix-sum of run lengths within the current chunk, so run `i` + // occupies output positions [chunk_out_off[i], chunk_out_off[i+1]). + // `chunk_meta[i]` encodes both the payload offset (into s_start) and, + // in the top bit, whether the run is literal (1) or RLE (0). + __shared__ int chunk_out_off[max_runs_per_chunk + 1]; + __shared__ uint32_t chunk_meta[max_runs_per_chunk]; + cuda::std::span const chunk_out_off_v{chunk_out_off, max_runs_per_chunk + 1}; + cuda::std::span const chunk_meta_v{chunk_meta, max_runs_per_chunk}; + __shared__ int s_chunk_runs; // number of runs parsed in this chunk (num_runs) + __shared__ int s_chunk_total; // sum of run lengths in this chunk (run_prefix_end) + __shared__ int s_base_out; // absolute output pos where this chunk starts + + auto const warp = cg::tiled_partition(group); + int const lane = warp.thread_rank(); + int const warp_id = warp.meta_group_rank(); + int const num_warps = warp.meta_group_size(); + int const value_width = cudf::util::div_rounding_up_unsafe(level_bits, 8); + // Bit mask used to extract a single level from a bit-packed literal-run + // payload word. Invariant across the whole call; hoisted out of the + // phase-2 expand loop to keep it out of the hot register set. + uint32_t const level_mask = (level_bits == 32) ? 0xffffffffu : ((1u << level_bits) - 1); + int out_pos_total = cur_values; + int const out_end = cur_values + output_count; + + // Outer loop: process the requested output range in chunks of up to + // `max_runs_per_chunk` runs at a time until we have emitted `output_count` + // values or run out of encoded input. + while (out_pos_total < out_end) { + // ----- Phase 1: single-thread run-header parse ------------------ + // Thread 0 walks the encoded stream, decoding VLQ run headers and + // filling chunk_out_off / chunk_meta. Do not use cg::invoke_one here: + // it may choose different threads across calls, but `cur` is per-thread + // rle_stream state that must persist on thread 0. + // The other threads wait at the group.sync() below. This is cheap because + // it is bounded by max_runs_per_chunk headers and header parsing is + // inherently serial. + if (group.thread_rank() == 0) { + int run_prefix_end = 0; + int num_runs = 0; + int out_base = out_pos_total; + chunk_out_off_v[0] = 0; + // Parse up to max_runs_per_chunk headers, stopping early if the output range + // fills up or the encoded stream is exhausted. + while (num_runs < max_runs_per_chunk && (out_base + run_prefix_end) < out_end && + cur < end) { + uint32_t const level_run = get_vlq32(cur, end); + + // Parquet RLE header format: LSB selects the encoding. + // bit 0 = 1 -> literal (bit-packed) run of `groups*8` values + // bit 0 = 0 -> RLE run of `level_run >> 1` copies of one value + // The high bit of `run_desc` distinguishes the two at expand time; + // see `run_desc_literal_flag` / `run_desc_offset_mask` for the + // bit layout and the 31-bit offset invariant enforced below. + cudf_assert(static_cast(cur - s_start) < (uint64_t{run_desc_offset_mask} + 1)); + int run_len; + uint32_t run_desc; + if (level_run & 1u) { + int const groups = level_run >> 1; + run_len = groups * 8; + run_desc = static_cast(cur - s_start) | run_desc_literal_flag; + cur += groups * level_bits; + } else { + run_len = level_run >> 1; + run_desc = static_cast(cur - s_start); + cur += value_width; + } + // Clamp the run to the remaining output window. A run can legally + // straddle out_end when row-range filtering makes output_count + // smaller than INT_MAX (e.g. skip_rows / num_rows / bounds-page + // filtering in preprocess_levels_kernel). Without clamping, Phase 2 + // would write past the output buffer causing silent data corruption. + // `cur` has already been advanced past the full payload above, which + // is correct: we emit only `run_len` values but do not need to + // re-parse the header on any subsequent call (the outer loop exits + // after this chunk because out_pos_total will equal out_end). + int const room = out_end - (out_base + run_prefix_end); + run_len = min(run_len, room); + run_prefix_end += run_len; + chunk_meta_v[num_runs] = run_desc; + chunk_out_off_v[++num_runs] = run_prefix_end; + } + s_chunk_runs = num_runs; + s_chunk_total = run_prefix_end; + s_base_out = out_base; + } + group.sync(); + + // ----- Phase 2: cooperative expand ------------------------------ + // All warps see the same chunk_out_off / chunk_meta tables. We split + // the flat output range [0, chunk_total) into `num_warps` equal-ish + // slices and each warp writes its slice. + int const chunk_runs = s_chunk_runs; + int const chunk_total = s_chunk_total; + int const base_out = s_base_out; + + if (chunk_runs == 0) { break; } + + int const per = cudf::util::div_rounding_up_safe(chunk_total, num_warps); + int const lo = warp_id * per; + int const hi = min(lo + per, chunk_total); + if (lo < hi) { + // Per-lane expand: each lane owns output positions + // out_pos = lo + lane, lo + lane + 32, lo + lane + 64, ... + // and finds its own run. run_idx starts by binary-search on the + // lane's first out_pos, then advances forward by linear walk (usually + // 0 steps when still in the same run). Across all 32 lanes the + // linear walks amortize to <= (num_runs_in_slice / 32) warp cycles + // total. + // + // This keeps all 32 lanes writing on every iteration, instead of the + // per-run loop where only lanes 0..(run_len-1) do useful work on + // short runs. + int out_pos = lo + lane; + int run_idx = + static_cast(cuda::std::upper_bound(chunk_out_off_v.begin(), + chunk_out_off_v.begin() + chunk_runs + 1, + out_pos) - + chunk_out_off_v.begin()) - + 1; + while (out_pos < hi) { + // Linear walk forward: no iterations if we're still in the same + // run (long run case), 1+ iterations only when out_pos crosses one + // or more short-run boundaries. + while (run_idx < chunk_runs && chunk_out_off_v[run_idx + 1] <= out_pos) { + ++run_idx; + } + int const run_start_out = chunk_out_off_v[run_idx]; + uint32_t const run_desc = chunk_meta_v[run_idx]; + + if (run_desc & run_desc_literal_flag) { + // Literal (bit-packed) run: bit-field extract for this lane's + // out_pos. + uint32_t const payload_off = run_desc & run_desc_offset_mask; + uint8_t const* payload = s_start + payload_off; + int const local = out_pos - run_start_out; + int bitpos = local * level_bits; + uint8_t const* source = payload + (bitpos >> 3); + bitpos &= 7; + uint32_t level_val; + if (source + sizeof(uint32_t) <= end) { + // Fast path: whole 32-bit field is in-bounds, so one unaligned + // load replaces up to four dependent byte reads. + level_val = cudf::io::unaligned_load(source); + } else { + // Tail path: within the last 4 bytes of the encoded stream, so + // fall back to per-byte reads and guard each against `end`. + level_val = 0; + if (source < end) { level_val = source[0]; } + if (level_bits > 8 - bitpos && (source + 1) < end) { + level_val |= static_cast(source[1]) << 8; + if (level_bits > 16 - bitpos && (source + 2) < end) { + level_val |= static_cast(source[2]) << 16; + if (level_bits > 24 - bitpos && (source + 3) < end) { + level_val |= static_cast(source[3]) << 24; + } + } + } + } + level_val = (level_val >> bitpos) & level_mask; + output[rolling_index(base_out + out_pos)] = + static_cast(level_val); + } else { + // RLE run: read the single repeated value from s_start. + // Guard each byte against `end` to match the literal path, since a + // truncated Parquet page can leave `run_desc`'s payload offset + // pointing within [s_start, end) while `level_bits > 8` implies + // vptr[1..3] may lie past `end`. + uint8_t const* vptr = s_start + (run_desc & run_desc_offset_mask); + uint32_t level_val = 0; + if (vptr < end) { level_val = vptr[0]; } + if constexpr (sizeof(level_t) > 1) { + if (level_bits > 8 && (vptr + 1) < end) { + level_val |= static_cast(vptr[1]) << 8; + if constexpr (sizeof(level_t) > 2) { + if (level_bits > 16 && (vptr + 2) < end) { + level_val |= static_cast(vptr[2]) << 16; + if (level_bits > 24 && (vptr + 3) < end) { + level_val |= static_cast(vptr[3]) << 24; + } + } + } + } + } + output[rolling_index(base_out + out_pos)] = + static_cast(level_val); + } + out_pos += warp.size(); + } + } + // Barrier before rewriting the shared tables on the next iteration. + group.sync(); + + out_pos_total = base_out + chunk_total; + } + + int const decoded = out_pos_total - cur_values; + cur_values = out_pos_total; + return decoded; + } + + __device__ inline int decode_next(int t, int count) + { + // Fast path: level_bits == 0 means every level is implicitly 0, so no + // headers or payloads need parsing. This is a very common case: columns + // with no nulls (especially non-nested ones) have all-zero definition + // levels. Handled here so both decode_next_ring and decode_next_chunked + // stay focused on the general RLE path. + // + // The write uses `cur_values + written + t` rather than `written + t` so it + // targets the correct ring slots regardless of how many times decode_next + // has already been called. No current caller enters this fast path with + // cur_values > 0 -- the writer floors dict_rle_bits >= 1 in chunk_dict.cu, + // and the REPETITION/DEFINITION decoders in decode_preprocess.cu are + // single-call -- so all reachable end-to-end tests still pass with the + // simpler `written + t` form. This invariant is preserved defensively so + // any future caller that iterates decode_next with level_bits == 0 stays + // correct without a silent off-by-one in the ring buffer. + int const output_count = min(count, total_values - cur_values); + if (level_bits == 0) { + int written = 0; + while (written < output_count) { + int const batch_size = min(num_rle_stream_decode_threads, output_count - written); + if (t < batch_size) { + output[rolling_index(cur_values + written + t)] = 0; + } + written += batch_size; + } + cur_values += output_count; + return output_count; + } + if constexpr (use_chunked_expand) { + return decode_next_chunked(cg::this_thread_block(), count); + } else { + return decode_next_ring(cg::this_thread_block(), count); + } + } + __device__ inline int skip_runs(int target_count) { // we want to process all runs UP TO BUT NOT INCLUDING the run that overlaps with the skip @@ -429,6 +745,7 @@ struct rle_stream { __device__ inline int skip_decode(int t, int count) { + static_assert(not use_chunked_expand, "skip_decode is not supported by chunked-expand"); int const output_count = min(count, total_values - cur_values); // if level_bits == 0, there's nothing to do @@ -440,4 +757,7 @@ struct rle_stream { __device__ inline int decode_next(int t) { return decode_next(t, max_output_values); } }; +template +using rle_stream_chunked = rle_stream; + } // namespace cudf::io::parquet::detail