Fix

cpp/src/io/orc/stripe_data.cu

@@ -145,53 +145,59 @@ class run_cache {
  public:
   __forceinline__ __device__ void initialize(orcdec_state_s* s)
   {
-    _status = (s->top.data.index.run_pos[CI_DATA2] > 0 and s->chunk.type_kind == TIMESTAMP)
-                ? status::CAN_WRITE_TO_CACHE
-                : status::DISABLED;
-    _num_rows_to_reuse = 0;
-    _run_length        = 0;
+    _status          = (s->top.data.index.run_pos[CI_DATA2] > 0 and s->chunk.type_kind == TIMESTAMP)
+                         ? status::CAN_WRITE_TO_CACHE
+                         : status::DISABLED;
+    _reusable_length = 0;
+    _run_length      = 0;
   }
 
-  __forceinline__ __device__ void set_num_rows_to_reuse(uint32_t run_length, uint32_t max_length)
+  __forceinline__ __device__ void set_reusable_length(uint32_t run_length, uint32_t max_length)
   {
     if (_status == status::CAN_WRITE_TO_CACHE) {
-      _run_length        = run_length;
-      _num_rows_to_reuse = (_run_length > max_length) ? (_run_length - max_length) : 0;
+      _run_length      = run_length;
+      _reusable_length = (_run_length > max_length) ? (_run_length - max_length) : 0;
     }
   }
 
   __forceinline__ __device__ uint32_t adjust_max_length(uint32_t max_length)
   {
     auto new_max_length{max_length};
-    if (_status == status::CAN_READ_FROM_CACHE) {
-      if (_num_rows_to_reuse > 0) { new_max_length -= _num_rows_to_reuse; }
+    if (_status == status::CAN_READ_FROM_CACHE and _reusable_length > 0) {
+      new_max_length -= _reusable_length;
     }
     return new_max_length;
   }
 
-  __forceinline__ __device__ void write_to_cache_if_needed(int64_t* src)
+  __forceinline__ __device__ void write_to_cache(int64_t* src)
   {
-    const auto warp_idx = cub::WarpId();
-    if (warp_idx != 0) { return; }
-
-    const auto lane_idx = cub::LaneId();
-    if (_status == status::CAN_WRITE_TO_CACHE and _num_rows_to_reuse > 0) {
-      for (uint32_t idx = lane_idx; idx < _run_length; idx += 32) {
-        if (idx < _run_length) { _buf[idx] = src[idx]; }
+    // Block until the src data, generated by the 1st warp, for the thread block are ready.
+    __syncthreads();
+    const auto tid = threadIdx.x;
+
+    // All threads in the block take a uniform code path.
+    // _reusable_length ranges between [0, 512]
+    if (_status == status::CAN_WRITE_TO_CACHE and _reusable_length > 0) {
+      const auto length_to_skip = _run_length - _reusable_length;
+      if (tid < _reusable_length) {
+        const auto src_idx = tid + length_to_skip;
+        _buf[tid]          = src[src_idx];
       }
-      __syncwarp();
-      if (lane_idx == 0) { _status = status::CAN_READ_FROM_CACHE; }
+      // Block until all writes are done to safely change _status.
+      __syncthreads();
+      if (tid == 0) { _status = status::CAN_READ_FROM_CACHE; }
     } else {
-      __syncwarp();
-      if (lane_idx == 0) { _status = status::DISABLED; }
+      __syncthreads();
+      if (tid == 0) { _status = status::DISABLED; }
     }
+    __syncthreads();
   }
 
-  __forceinline__ __device__ void read_from_cache_if_needed([[maybe_unused]] uint64_t* dst) {}
+  __forceinline__ __device__ void read_from_cache([[maybe_unused]] uint64_t* dst) {}
 
  private:
   status _status;
-  uint32_t _num_rows_to_reuse;
+  uint32_t _reusable_length;
   uint32_t _run_length;
   int64_t _buf[bytestream_buffer_size >> 4];
 };
@@ -752,7 +758,7 @@ static __device__ uint32_t Integer_RLEv2(orc_bytestream_s* bs,
         }
       }
 
-      if (run_cache_bs != nullptr) { run_cache_bs->set_num_rows_to_reuse(n, maxvals); }
+      if (run_cache_bs != nullptr) { run_cache_bs->set_reusable_length(n, maxvals); }
 
       if ((numvals != 0) and (numvals + n > maxvals)) break;
       // case where there are buffered values and can't consume a whole chunk
@@ -935,9 +941,8 @@ static __device__ uint32_t Integer_RLEv2(orc_bytestream_s* bs,
     __syncwarp();
   }
   if constexpr (cuda::std::is_same_v<T, int64_t>) {
-    if (run_cache_bs != nullptr) { run_cache_bs->write_to_cache_if_needed(vals); }
+    if (run_cache_bs != nullptr) { run_cache_bs->write_to_cache(vals); }
   }
-  __syncthreads();
   return rle->num_vals;
 }