From c181628367f4164f482853c3ce87150c0c31835f Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Fri, 21 Aug 2026 18:35:37 +0000 Subject: [PATCH 1/2] Pass a CUDA stream to the nvCOMP batched temp size queries Newer nvCOMP releases rename the host-only batched temp size queries and give them a trailing stream parameter, which the library uses to identify the device whose properties determine the temp size. Wrap the call in a macro so both spellings build, and thread a stream down from the existing call sites, all of which already had one available. --- cpp/src/io/comp/decompression.cpp | 10 +- cpp/src/io/comp/decompression.hpp | 3 +- cpp/src/io/comp/nvcomp_adapter.cpp | 193 ++++++++++-------- cpp/src/io/comp/nvcomp_adapter.hpp | 4 +- .../io/parquet/reader_impl_chunking_utils.cu | 9 +- 5 files changed, 127 insertions(+), 92 deletions(-) diff --git a/cpp/src/io/comp/decompression.cpp b/cpp/src/io/comp/decompression.cpp index 4a7c518c5eb8..b5cc1d6fcb2f 100644 --- a/cpp/src/io/comp/decompression.cpp +++ b/cpp/src/io/comp/decompression.cpp @@ -619,7 +619,8 @@ size_t get_uncompressed_size(compression_type compression, host_span= MAKE_SEMANTIC_VERSION(5, 3, 0)) +// nvCOMP 6.0 drops the `Async` suffix from the host-only +// `nvcompBatchedXXX(De)CompressGetTempSizeAsync` functions and gave them a stream parameter. +#if NVCOMP_VER >= MAKE_SEMANTIC_VERSION(6, 0, 0) +// Call sites spell `fn` without the `Async` suffix +#define NVCOMP_BATCHED_GET_TEMP_SIZE(fn, stream, ...) fn(__VA_ARGS__, stream) +#else +// ingore the `stream` parameter +#define NVCOMP_BATCHED_GET_TEMP_SIZE(fn, stream, ...) fn##Async(__VA_ARGS__) +#endif + namespace cudf::io::detail::nvcomp { namespace { @@ -89,45 +99,55 @@ namespace { CUDF_FAIL("Unsupported compression type: " + compression_type_name(compression)); \ } while (0) -// Dispatcher for nvcompBatchedDecompressGetTempSizeAsync -auto batched_decompress_get_temp_size_async(compression_type compression, - size_t num_chunks, - size_t max_uncompressed_chunk_bytes, - size_t* temp_bytes, - size_t max_total_uncompressed_bytes) +// Dispatcher for nvcompBatchedDecompressGetTempSize +auto batched_decompress_get_temp_size(compression_type compression, + size_t num_chunks, + size_t max_uncompressed_chunk_bytes, + size_t* temp_bytes, + size_t max_total_uncompressed_bytes, + [[maybe_unused]] cuda::stream_ref stream) { switch (compression) { case compression_type::SNAPPY: - return nvcompBatchedSnappyDecompressGetTempSizeAsync(num_chunks, - max_uncompressed_chunk_bytes, - nvcompBatchedSnappyDecompressDefaultOpts, - temp_bytes, - max_total_uncompressed_bytes); + return NVCOMP_BATCHED_GET_TEMP_SIZE(nvcompBatchedSnappyDecompressGetTempSize, + stream.get(), + num_chunks, + max_uncompressed_chunk_bytes, + nvcompBatchedSnappyDecompressDefaultOpts, + temp_bytes, + max_total_uncompressed_bytes); case compression_type::ZSTD: - return nvcompBatchedZstdDecompressGetTempSizeAsync(num_chunks, - max_uncompressed_chunk_bytes, - nvcompBatchedZstdDecompressDefaultOpts, - temp_bytes, - max_total_uncompressed_bytes); + return NVCOMP_BATCHED_GET_TEMP_SIZE(nvcompBatchedZstdDecompressGetTempSize, + stream.get(), + num_chunks, + max_uncompressed_chunk_bytes, + nvcompBatchedZstdDecompressDefaultOpts, + temp_bytes, + max_total_uncompressed_bytes); case compression_type::LZ4: - return nvcompBatchedLZ4DecompressGetTempSizeAsync(num_chunks, - max_uncompressed_chunk_bytes, - nvcompBatchedLZ4DecompressDefaultOpts, - temp_bytes, - max_total_uncompressed_bytes); + return NVCOMP_BATCHED_GET_TEMP_SIZE(nvcompBatchedLZ4DecompressGetTempSize, + stream.get(), + num_chunks, + max_uncompressed_chunk_bytes, + nvcompBatchedLZ4DecompressDefaultOpts, + temp_bytes, + max_total_uncompressed_bytes); case compression_type::DEFLATE: - return nvcompBatchedDeflateDecompressGetTempSizeAsync( - num_chunks, - max_uncompressed_chunk_bytes, - nvcompBatchedDeflateDecompressDefaultOpts, - temp_bytes, - max_total_uncompressed_bytes); + return NVCOMP_BATCHED_GET_TEMP_SIZE(nvcompBatchedDeflateDecompressGetTempSize, + stream.get(), + num_chunks, + max_uncompressed_chunk_bytes, + nvcompBatchedDeflateDecompressDefaultOpts, + temp_bytes, + max_total_uncompressed_bytes); case compression_type::GZIP: - return nvcompBatchedGzipDecompressGetTempSizeAsync(num_chunks, - max_uncompressed_chunk_bytes, - nvcompBatchedGzipDecompressDefaultOpts, - temp_bytes, - max_total_uncompressed_bytes); + return NVCOMP_BATCHED_GET_TEMP_SIZE(nvcompBatchedGzipDecompressGetTempSize, + stream.get(), + num_chunks, + max_uncompressed_chunk_bytes, + nvcompBatchedGzipDecompressDefaultOpts, + temp_bytes, + max_total_uncompressed_bytes); default: UNSUPPORTED_COMPRESSION(compression); } } @@ -217,50 +237,56 @@ auto batched_decompress_async(compression_type compression, } } -// Wrapper for nvcompBatchedCompressGetTempSizeAsync -nvcompStatus_t batched_compress_get_temp_size_async(compression_type compression, - size_t batch_size, - size_t max_uncompressed_chunk_bytes, - size_t* temp_size, - size_t max_total_uncompressed_bytes) +// Wrapper for nvcompBatchedCompressGetTempSize +nvcompStatus_t batched_compress_get_temp_size(compression_type compression, + size_t batch_size, + size_t max_uncompressed_chunk_bytes, + size_t* temp_size, + size_t max_total_uncompressed_bytes, + [[maybe_unused]] cuda::stream_ref stream) { switch (compression) { case compression_type::SNAPPY: - return nvcompBatchedSnappyCompressGetTempSizeAsync(batch_size, - max_uncompressed_chunk_bytes, - nvcompBatchedSnappyCompressDefaultOpts, - temp_size, - max_total_uncompressed_bytes); - break; + return NVCOMP_BATCHED_GET_TEMP_SIZE(nvcompBatchedSnappyCompressGetTempSize, + stream.get(), + batch_size, + max_uncompressed_chunk_bytes, + nvcompBatchedSnappyCompressDefaultOpts, + temp_size, + max_total_uncompressed_bytes); case compression_type::DEFLATE: - return nvcompBatchedDeflateCompressGetTempSizeAsync(batch_size, - max_uncompressed_chunk_bytes, - nvcompBatchedDeflateCompressDefaultOpts, - temp_size, - max_total_uncompressed_bytes); - break; + return NVCOMP_BATCHED_GET_TEMP_SIZE(nvcompBatchedDeflateCompressGetTempSize, + stream.get(), + batch_size, + max_uncompressed_chunk_bytes, + nvcompBatchedDeflateCompressDefaultOpts, + temp_size, + max_total_uncompressed_bytes); case compression_type::ZSTD: - return nvcompBatchedZstdCompressGetTempSizeAsync(batch_size, - max_uncompressed_chunk_bytes, - nvcompBatchedZstdCompressDefaultOpts, - temp_size, - max_total_uncompressed_bytes); - break; + return NVCOMP_BATCHED_GET_TEMP_SIZE(nvcompBatchedZstdCompressGetTempSize, + stream.get(), + batch_size, + max_uncompressed_chunk_bytes, + nvcompBatchedZstdCompressDefaultOpts, + temp_size, + max_total_uncompressed_bytes); case compression_type::LZ4: - return nvcompBatchedLZ4CompressGetTempSizeAsync(batch_size, - max_uncompressed_chunk_bytes, - nvcompBatchedLZ4CompressDefaultOpts, - temp_size, - max_total_uncompressed_bytes); - break; + return NVCOMP_BATCHED_GET_TEMP_SIZE(nvcompBatchedLZ4CompressGetTempSize, + stream.get(), + batch_size, + max_uncompressed_chunk_bytes, + nvcompBatchedLZ4CompressDefaultOpts, + temp_size, + max_total_uncompressed_bytes); #if CUDF_NVCOMP_HAS_GZIP_COMPRESSION case compression_type::GZIP: - return nvcompBatchedGzipCompressGetTempSizeAsync(batch_size, - max_uncompressed_chunk_bytes, - nvcompBatchedGzipCompressDefaultOpts, - temp_size, - max_total_uncompressed_bytes); - break; + return NVCOMP_BATCHED_GET_TEMP_SIZE(nvcompBatchedGzipCompressGetTempSize, + stream.get(), + batch_size, + max_uncompressed_chunk_bytes, + nvcompBatchedGzipCompressDefaultOpts, + temp_size, + max_total_uncompressed_bytes); #endif default: UNSUPPORTED_COMPRESSION(compression); } @@ -269,16 +295,16 @@ nvcompStatus_t batched_compress_get_temp_size_async(compression_type compression size_t batched_compress_temp_size(compression_type compression, size_t batch_size, size_t max_uncompressed_chunk_bytes, - size_t max_total_uncompressed_bytes) + size_t max_total_uncompressed_bytes, + cuda::stream_ref stream) { - size_t temp_size = 0; - nvcompStatus_t nvcomp_status = nvcompStatus_t::nvcompSuccess; - - nvcomp_status = batched_compress_get_temp_size_async(compression, - batch_size, - max_uncompressed_chunk_bytes, - &temp_size, - max_total_uncompressed_bytes); + size_t temp_size = 0; + nvcompStatus_t const nvcomp_status = batched_compress_get_temp_size(compression, + batch_size, + max_uncompressed_chunk_bytes, + &temp_size, + max_total_uncompressed_bytes, + stream); CHECK_NVCOMP_STATUS(nvcomp_status); return temp_size; @@ -526,7 +552,7 @@ size_t batched_decompress_temp_size_ex(compression_type compression, } // Fallback to the original batched decompress temp size calculation return batched_decompress_temp_size( - compression, input_data_ptrs.size(), max_uncomp_chunk_size, max_total_uncomp_size); + compression, input_data_ptrs.size(), max_uncomp_chunk_size, max_total_uncomp_size, stream); } } // namespace @@ -534,11 +560,12 @@ size_t batched_decompress_temp_size_ex(compression_type compression, size_t batched_decompress_temp_size(compression_type compression, size_t num_chunks, size_t max_uncomp_chunk_size, - size_t max_total_uncomp_size) + size_t max_total_uncomp_size, + cuda::stream_ref stream) { size_t temp_size = 0; - nvcompStatus_t const nvcomp_status = batched_decompress_get_temp_size_async( - compression, num_chunks, max_uncomp_chunk_size, &temp_size, max_total_uncomp_size); + nvcompStatus_t const nvcomp_status = batched_decompress_get_temp_size( + compression, num_chunks, max_uncomp_chunk_size, &temp_size, max_total_uncomp_size, stream); CHECK_NVCOMP_STATUS(nvcomp_status); return temp_size; } @@ -673,8 +700,8 @@ void batched_compress(compression_type compression, auto const [max_uncomp_chunk_size, total_uncomp_size] = max_chunk_and_total_input_size(nvcomp_args.input_data_sizes, stream); - auto const temp_size = - batched_compress_temp_size(compression, num_chunks, max_uncomp_chunk_size, total_uncomp_size); + auto const temp_size = batched_compress_temp_size( + compression, num_chunks, max_uncomp_chunk_size, total_uncomp_size, stream); rmm::device_buffer scratch(temp_size, stream); CUDF_EXPECTS(is_aligned(scratch.data(), 8), "Compression failed, misaligned scratch buffer"); diff --git a/cpp/src/io/comp/nvcomp_adapter.hpp b/cpp/src/io/comp/nvcomp_adapter.hpp index 5b7d91ab10c9..c22ab2ec8747 100644 --- a/cpp/src/io/comp/nvcomp_adapter.hpp +++ b/cpp/src/io/comp/nvcomp_adapter.hpp @@ -46,12 +46,14 @@ void batched_decompress(compression_type compression, * @param[in] num_chunks The number of decompression chunks to be processed * @param[in] max_uncomp_chunk_size Maximum size of any single uncompressed chunk * @param[in] max_total_uncomp_size Maximum total size of uncompressed data + * @param[in] stream CUDA stream that the decompression will run on * @returns The total required size in bytes */ size_t batched_decompress_temp_size(compression_type compression, size_t num_chunks, size_t max_uncomp_chunk_size, - size_t max_total_uncomp_size); + size_t max_total_uncomp_size, + cuda::stream_ref stream); /** * @brief Return the amount of temporary space required in bytes for a given decompression diff --git a/cpp/src/io/parquet/reader_impl_chunking_utils.cu b/cpp/src/io/parquet/reader_impl_chunking_utils.cu index 68c9ff5f6581..164554a26402 100644 --- a/cpp/src/io/parquet/reader_impl_chunking_utils.cu +++ b/cpp/src/io/parquet/reader_impl_chunking_utils.cu @@ -731,9 +731,10 @@ rmm::device_uvector compute_decompression_scratch_sizes( // retrieve to host so we can get compression scratch sizes auto temp_cost = cudf::detail::make_pinned_vector_async(pages.size(), stream); auto h_decomp_info = cudf::detail::make_pinned_vector(decomp_info, stream); - std::transform(h_decomp_info.begin(), h_decomp_info.end(), temp_cost.begin(), [](auto const& d) { - return cudf::io::detail::get_decompression_scratch_size(d); - }); + std::transform( + h_decomp_info.begin(), h_decomp_info.end(), temp_cost.begin(), [stream](auto const& d) { + return cudf::io::detail::get_decompression_scratch_size(d, stream); + }); rmm::device_uvector d_temp_cost = cudf::detail::make_device_uvector(temp_cost, stream, cudf::get_current_device_resource_ref()); @@ -790,7 +791,7 @@ rmm::device_uvector compute_decompression_scratch_sizes( } page_spans.resize(end_iter - page_spans.begin(), stream); - auto const total_temp_size = get_decompression_scratch_size(total_decomp_info); + auto const total_temp_size = get_decompression_scratch_size(total_decomp_info, stream); auto const total_temp_size_ex = cudf::io::detail::get_decompression_scratch_size_ex( total_decomp_info.type, page_spans, From 0245d39085b7a3ca03fe7316f24b9ab83d2711e7 Mon Sep 17 00:00:00 2001 From: Vukasin Milovanovic Date: Tue, 8 Sep 2026 17:51:15 +0000 Subject: [PATCH 2/2] Fix trailing whitespace and typo in the nvCOMP macro comments --- cpp/src/io/comp/nvcomp_adapter.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/src/io/comp/nvcomp_adapter.cpp b/cpp/src/io/comp/nvcomp_adapter.cpp index d76ca790a0d4..c01435a3703d 100644 --- a/cpp/src/io/comp/nvcomp_adapter.cpp +++ b/cpp/src/io/comp/nvcomp_adapter.cpp @@ -23,13 +23,13 @@ #define CUDF_NVCOMP_HAS_GZIP_COMPRESSION (NVCOMP_VER >= MAKE_SEMANTIC_VERSION(5, 3, 0)) -// nvCOMP 6.0 drops the `Async` suffix from the host-only +// nvCOMP 6.0 drops the `Async` suffix from the host-only // `nvcompBatchedXXX(De)CompressGetTempSizeAsync` functions and gave them a stream parameter. #if NVCOMP_VER >= MAKE_SEMANTIC_VERSION(6, 0, 0) // Call sites spell `fn` without the `Async` suffix #define NVCOMP_BATCHED_GET_TEMP_SIZE(fn, stream, ...) fn(__VA_ARGS__, stream) #else -// ingore the `stream` parameter +// ignore the `stream` parameter #define NVCOMP_BATCHED_GET_TEMP_SIZE(fn, stream, ...) fn##Async(__VA_ARGS__) #endif