Optimize RLE decoding by staging in shared memory

cpp/src/io/parquet/decode_fixed.cu

@@ -1081,15 +1081,16 @@ CUDF_KERNEL void __launch_bounds__(decode_block_size_t, 8)
   rle_stream<uint32_t, decode_block_size_t, rolling_buf_size> 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
   rle_stream<uint32_t, decode_block_size_t, rolling_buf_size> bool_stream{bool_runs};
   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();

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
  */
 
@@ -14,6 +14,7 @@
 #include <rmm/exec_policy.hpp>
 
 #include <cooperative_groups.h>
+#include <cuda/barrier>
 #include <cuda/std/iterator>
 #include <cuda/std/limits>
 
@@ -378,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 <typename level_t, int level_decode_block_size>
 CUDF_KERNEL void __launch_bounds__(level_decode_block_size)
   preprocess_levels_kernel(PageInfo* pages,
@@ -418,6 +420,16 @@ CUDF_KERNEL void __launch_bounds__(level_decode_block_size)
   rle_stream<level_t, level_decode_block_size, max_output_values>
     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.
+  using rle_stream_t = rle_stream<level_t, level_decode_block_size, max_output_values>;
+  __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]);
@@ -429,33 +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(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);
-  }
   if (has_repetition) {
-    decoders[level_type::REPETITION].init(s->col.level_bits[level_type::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);
+                                          num_to_decode,
+                                          stage,
+                                          &copy_barrier);
+    copy_barrier.arrive_and_wait();
+    decoders[level_type::REPETITION].decode_next(t, num_to_decode);
   }
-  block.sync();
 
   // 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(&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);
+  }
 }
 
 }  // anonymous namespace

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
  */
 
@@ -10,6 +10,11 @@
 #include <cudf/detail/utilities/cuda.cuh>
 #include <cudf/detail/utilities/integer_utils.hpp>
 
+#include <cooperative_groups.h>
+#include <cuda/barrier>
+#include <cuda/std/iterator>
+#include <cuda/std/memory>
+
 namespace cudf::io::parquet::detail {
 
 template <int num_threads>
@@ -178,18 +183,33 @@ struct rle_stream {
   int fill_index;
   int decode_index;
 
+  // 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 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.
+  static constexpr int smem_stage_size = 8 * 1024;
+
   __device__ rle_stream(rle_run* _runs) : runs(_runs) {}
 
   __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 <typename Group>
+  __device__ void init(Group const& group,
+                       int _level_bits,
                        uint8_t const* _start,
                        uint8_t const* _end,
                        level_t* _output,
-                       int _total_values)
+                       int _total_values,
+                       uint8_t* _smem_stage                                   = nullptr,
+                       cuda::barrier<cuda::thread_scope_block>* _copy_barrier = nullptr)
   {
     level_bits = _level_bits;
     cur        = _start;
@@ -203,6 +223,26 @@ struct rle_stream {
     cur_values   = 0;
     fill_index   = 0;
     decode_index = -1;  // signals the first iteration. Nothing to decode.
+
+    // 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.
+    if (_smem_stage != nullptr) {
+      auto* const smem_stage =
+        static_cast<uint8_t const*>(cuda::std::assume_aligned<16>(_smem_stage));
+      auto const len = static_cast<int>(cuda::std::distance(_start, _end));
+      if (len > 0 && len <= smem_stage_size) {
+        cuda::memcpy_async(group, _smem_stage, _start, static_cast<size_t>(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.
+        cur = smem_stage;
+        end = smem_stage + len;
+      }
+    }
   }
 
   __device__ inline int get_rle_run_info(rle_run& run)