From 8bd096460c1eb65b27d52edaa625288ce52ea690 Mon Sep 17 00:00:00 2001 From: ykiran Date: Wed, 16 Sep 2026 20:21:55 -0700 Subject: [PATCH] Interleaved word copies --- cpp/src/io/parquet/page_string_decode.cu | 58 ++++++++++++++++++------ 1 file changed, 43 insertions(+), 15 deletions(-) diff --git a/cpp/src/io/parquet/page_string_decode.cu b/cpp/src/io/parquet/page_string_decode.cu index ae66893382a8..d7439a455213 100644 --- a/cpp/src/io/parquet/page_string_decode.cu +++ b/cpp/src/io/parquet/page_string_decode.cu @@ -1088,22 +1088,50 @@ inline __device__ bool prefetch_string_data(int t, if (total_bytes_to_copy <= 0) { return false; } // No data left to copy buffer_end = buffer_base + total_bytes_to_copy; - // Nominally, each thread will copy an equal number of bytes; this rounds up. - auto const nominal_thread_bytes_to_copy = - cudf::util::div_rounding_up_unsafe(total_bytes_to_copy, block_size); - int32_t const thread_offset = nominal_thread_bytes_to_copy * t; - - if (thread_offset < total_bytes_to_copy) { - // Guard against the end of the data stream - int32_t const thread_bytes_to_copy = - cuda::std::min(nominal_thread_bytes_to_copy, total_bytes_to_copy - thread_offset); - - if (thread_bytes_to_copy > 0) { - int32_t const thread_copy_from_index = buffer_base + thread_offset; - cuda::std::memcpy(reinterpret_cast(&prefetch_buffer[thread_offset]), - reinterpret_cast(&cur[thread_copy_from_index]), - thread_bytes_to_copy); + // Copy in word-sized units, interleaved across the threads: thread `t` writes the shared + // words t, t + block_size, ... A word index maps to the bank of the same number, so every + // lane in a warp writes a distinct bank and the stores are conflict free. The reads feeding + // them are contiguous across the warp, so they coalesce as well. + // + // Handing each thread a contiguous run instead -- thread `t` copying the bytes + // [t * run, (t + 1) * run) -- is the obvious split, but it places consecutive lanes `run` + // bytes apart. For a 1024-byte buffer and a warp-sized block that is 32 bytes, which is only + // four distinct banks for 32 lanes: an eight-way conflict on every byte stored, on top of + // reads scattered over the whole window rather than coalescing. + constexpr int32_t bytes_per_word = sizeof(uint32_t); + static_assert(prefetch_size % bytes_per_word == 0, + "the prefetch buffer must hold a whole number of words"); + + auto* const dst_words = reinterpret_cast(prefetch_buffer); + uint8_t const* const src = cur + buffer_base; + + // Rounded up: the final word may reach past the copied data. That is harmless, as + // prefetch_size is a whole number of words, so the store still lands inside the buffer. + int32_t const num_words = + cudf::util::div_rounding_up_unsafe(total_bytes_to_copy, bytes_per_word); + + // Warp uniform: every thread derives the same source pointer, so this never diverges. + bool const src_is_word_aligned = (reinterpret_cast(src) % bytes_per_word) == 0; + + for (int32_t w = t; w < num_words; w += block_size) { + int32_t const byte_offset = w * bytes_per_word; + uint32_t packed = 0; + + if (src_is_word_aligned and (byte_offset + bytes_per_word <= total_bytes_to_copy)) { + packed = reinterpret_cast(src)[w]; + } else { + // Either the source is unaligned -- the offset into the stream is data dependent, so it + // only lands on a word boundary by chance -- or this is a partial word at the end of the + // data. Assemble the bytes that are actually there; the rest stay zero and are never read, + // because the scan stops at buffer_end. + int32_t const valid_bytes = + cuda::std::min(bytes_per_word, total_bytes_to_copy - byte_offset); + for (int32_t b = 0; b < valid_bytes; ++b) { + packed |= static_cast(src[byte_offset + b]) << (CHAR_BIT * b); + } } + + dst_words[w] = packed; } return true; }