Upgrade to nvCOMP 5.0.0.6

ci/build_wheel_libcudf.sh

@@ -36,8 +36,8 @@ export SKBUILD_CMAKE_ARGS="-DUSE_NVCOMP_FROM_LIBKVIKIO_WHEEL=ON"
 
 # repair wheels and write to the location that artifact-uploading code expects to find them
 python -m auditwheel repair \
-    --exclude libnvcomp.so.4 \
     --exclude libkvikio.so \
+    --exclude libnvcomp.so.5 \
     --exclude librapids_logger.so \
     --exclude librmm.so \
     -w "${RAPIDS_WHEEL_BLD_OUTPUT_DIR}" \

cmake/rapids_config.cmake

@@ -35,6 +35,9 @@ if(NOT _rapids_branch)
   )
 endif()
 
+set(rapids-cmake-repo vuule/rapids-cmake)
+set(rapids-cmake-branch fea-nvcomp-5)
+
 if(NOT rapids-cmake-version)
   set(rapids-cmake-version "${RAPIDS_VERSION_MAJOR_MINOR}")
 endif()

conda/environments/all_cuda-129_arch-aarch64.yaml

@@ -40,7 +40,7 @@ dependencies:
 - jupyter_client
 - libcurand-dev
 - libkvikio==25.10.*,>=0.0.0a0
-- libnvcomp-dev==4.2.0.11
+- libnvcomp-dev==5.0.0.6
 - libnvjitlink-dev
 - librdkafka>=2.8.0,<2.9.0a0
 - librmm==25.10.*,>=0.0.0a0

conda/environments/all_cuda-129_arch-x86_64.yaml

@@ -41,7 +41,7 @@ dependencies:
 - libcufile-dev
 - libcurand-dev
 - libkvikio==25.10.*,>=0.0.0a0
-- libnvcomp-dev==4.2.0.11
+- libnvcomp-dev==5.0.0.6
 - libnvjitlink-dev
 - librdkafka>=2.8.0,<2.9.0a0
 - librmm==25.10.*,>=0.0.0a0

conda/recipes/libcudf/conda_build_config.yaml

@@ -26,7 +26,7 @@ flatbuffers_version:
   - "=24.3.25"
 
 nvcomp_version:
-  - "=4.2.0.11"
+  - "=5.0.0.6"
 
 zlib_version:
   - ">=1.2.13"

dependencies.yaml

@@ -476,7 +476,7 @@ dependencies:
       - output_types: conda
         packages:
           # Align nvcomp version with rapids-cmake
-          - libnvcomp-dev==4.2.0.11
+          - libnvcomp-dev==5.0.0.6
   rapids_build_skbuild:
     common:
       - output_types: [conda, requirements, pyproject]

java/src/main/java/ai/rapids/cudf/nvcomp/BatchedCompressor.java

@@ -91,7 +91,8 @@ public DeviceMemoryBuffer[] compress(BaseDeviceMemoryBuffer[] origInputs, Cuda.S
         buildAddrsAndSizes(inputs, inputChunkAddrs, inputChunkSizes, compressedBuffers,
             outputChunkAddrs);
 
-        final long tempBufferSize = batchedCompressGetTempSize(numChunks, chunkSize);
+        final long tempBufferSize = batchedCompressGetTempSize(numChunks, chunkSize,
+          numChunks * chunkSize);
         try (DeviceMemoryBuffer addrsAndSizes = putAddrsAndSizesOnDevice(inputChunkAddrs,
                 inputChunkSizes, outputChunkAddrs, stream);
              DeviceMemoryBuffer tempBuffer =
@@ -308,9 +309,11 @@ private long[] calcOutputBufferSizes(int[] chunksPerInput, long[] outputChunkSiz
    * Get the temporary workspace size required to perform compression of an entire batch.
    * @param batchSize number of chunks in the batch
    * @param maxChunkSize maximum size of an uncompressed chunk in bytes
+   * @param totalSize Upper bound on the total uncompressed size of all chunks
    * @return The size of required temporary workspace in bytes to compress the batch.
    */
-  protected abstract long batchedCompressGetTempSize(long batchSize, long maxChunkSize);
+  protected abstract long batchedCompressGetTempSize(long batchSize, long maxChunkSize,
+    long totalSize);
 
    /**
    * Asynchronously compress a batch of buffers. Note that compressedSizesOutPtr must

java/src/main/java/ai/rapids/cudf/nvcomp/BatchedDecompressor.java

@@ -71,7 +71,8 @@ public void decompressAsync(BaseDeviceMemoryBuffer[] origInputs,
         totalChunks += numBufferChunks;
       }
 
-      final long tempBufferSize = batchedDecompressGetTempSize(totalChunks, chunkSize);
+      final long tempBufferSize = batchedDecompressGetTempSize(totalChunks, chunkSize,
+        totalChunks * chunkSize);
       try (DeviceMemoryBuffer devAddrsSizes = buildAddrsSizesBuffer(chunkSize, totalChunks,
               inputs.getArray(), chunksPerInput, outputs, stream);
            DeviceMemoryBuffer devTemp = DeviceMemoryBuffer.allocate(tempBufferSize)) {
@@ -198,10 +199,11 @@ private static HostMemoryBuffer fetchMetadata(long totalChunks, BaseDeviceMemory
    * Computes the temporary storage size in bytes needed to decompress a compressed batch.
    * @param numChunks number of chunks in the batch
    * @param maxUncompressedChunkBytes maximum uncompressed size of any chunk in bytes
+   * @param maxTotalSize Upper bound on the total uncompressed size of all chunks
    * @return number of temporary storage bytes needed to decompress the batch
    */
   protected abstract long batchedDecompressGetTempSize(long numChunks,
-      long maxUncompressedChunkBytes);
+      long maxUncompressedChunkBytes, long maxTotalSize);
 
     /**
    * Asynchronously decompress a batch of compressed data buffers.

java/src/main/java/ai/rapids/cudf/nvcomp/BatchedLZ4Compressor.java

@@ -32,8 +32,8 @@ public BatchedLZ4Compressor(long chunkSize, long maxIntermediateBufferSize) {
   }
 
   @Override
-  protected long batchedCompressGetTempSize(long batchSize, long maxChunkSize) {
-    return NvcompJni.batchedLZ4CompressGetTempSize(batchSize, maxChunkSize);
+  protected long batchedCompressGetTempSize(long batchSize, long maxChunkSize, long totalSize) {
+    return NvcompJni.batchedLZ4CompressGetTempSize(batchSize, maxChunkSize, totalSize);
   }
 
   @Override

java/src/main/java/ai/rapids/cudf/nvcomp/BatchedLZ4Decompressor.java

@@ -41,8 +41,10 @@ public static void decompressAsync(long chunkSize, BaseDeviceMemoryBuffer[] orig
   }
 
   @Override
-  protected long batchedDecompressGetTempSize(long numChunks, long maxUncompressedChunkBytes) {
-    return NvcompJni.batchedLZ4DecompressGetTempSize(numChunks, maxUncompressedChunkBytes);
+  protected long batchedDecompressGetTempSize(long numChunks, long maxUncompressedChunkBytes,
+      long maxTotalSize) {
+    return NvcompJni.batchedLZ4DecompressGetTempSize(numChunks, maxUncompressedChunkBytes,
+      maxTotalSize);
   }
 
   @Override

java/src/main/java/ai/rapids/cudf/nvcomp/BatchedZstdCompressor.java

@@ -31,8 +31,8 @@ public BatchedZstdCompressor(long chunkSize, long maxIntermediateBufferSize) {
   }
 
   @Override
-  protected long batchedCompressGetTempSize(long batchSize, long maxChunkSize) {
-    return NvcompJni.batchedZstdCompressGetTempSize(batchSize, maxChunkSize);
+  protected long batchedCompressGetTempSize(long batchSize, long maxChunkSize, long totalSize) {
+    return NvcompJni.batchedZstdCompressGetTempSize(batchSize, maxChunkSize, totalSize);
   }
 
   @Override

java/src/main/java/ai/rapids/cudf/nvcomp/BatchedZstdDecompressor.java

@@ -23,8 +23,10 @@ public BatchedZstdDecompressor(long chunkSize) {
   }
 
   @Override
-  protected long batchedDecompressGetTempSize(long numChunks, long maxUncompressedChunkBytes) {
-    return NvcompJni.batchedZstdDecompressGetTempSize(numChunks, maxUncompressedChunkBytes);
+  protected long batchedDecompressGetTempSize(long numChunks, long maxUncompressedChunkBytes,
+      long maxTotalSize) {
+    return NvcompJni.batchedZstdDecompressGetTempSize(numChunks, maxUncompressedChunkBytes,
+      maxTotalSize);
   }
 
   @Override

java/src/main/java/ai/rapids/cudf/nvcomp/NvcompJni.java

@@ -29,9 +29,10 @@ class NvcompJni {
    * Get the temporary workspace size required to perform compression of entire LZ4 batch.
    * @param batchSize number of chunks in the batch
    * @param maxChunkSize maximum size of an uncompressed chunk in bytes
+   * @param maxTotalSize Upper bound on the total uncompressed size of all chunks
    * @return The size of required temporary workspace in bytes to compress the batch.
    */
-  static native long batchedLZ4CompressGetTempSize(long batchSize, long maxChunkSize);
+  static native long batchedLZ4CompressGetTempSize(long batchSize, long maxChunkSize, long maxTotalSize);
 
   /**
    * Get the maximum size any chunk could compress to in a LZ4 batch. This is the minimum amount of
@@ -74,11 +75,13 @@ static native void batchedLZ4CompressAsync(
    * Computes the temporary storage size in bytes needed to decompress a LZ4-compressed batch.
    * @param numChunks number of chunks in the batch
    * @param maxUncompressedChunkBytes maximum uncompressed size of any chunk in bytes
+   * @param maxTotalSize Upper bound on the total uncompressed size of all chunks
    * @return number of temporary storage bytes needed to decompress the batch
    */
   static native long batchedLZ4DecompressGetTempSize(
       long numChunks,
-      long maxUncompressedChunkBytes);
+      long maxUncompressedChunkBytes,
+      long maxTotalSize);
 
   /**
    * Asynchronously decompress a batch of LZ4-compressed data buffers.
@@ -121,9 +124,10 @@ static native void batchedLZ4GetDecompressSizeAsync(
    * Get the temporary workspace size required to perform compression of entire zstd batch.
    * @param batchSize number of chunks in the batch
    * @param maxChunkSize maximum size of an uncompressed chunk in bytes
+   * @param maxTotalSize Upper bound on the total uncompressed size of all chunks
    * @return The size of required temporary workspace in bytes to compress the batch.
    */
-  static native long batchedZstdCompressGetTempSize(long batchSize, long maxChunkSize);
+  static native long batchedZstdCompressGetTempSize(long batchSize, long maxChunkSize, long maxTotalSize);
 
   /**
    * Get the maximum size any chunk could compress to in a ZSTD batch. This is the minimum
@@ -167,11 +171,13 @@ static native void batchedZstdCompressAsync(
    * ZSTD-compressed batch.
    * @param numChunks number of chunks in the batch
    * @param maxUncompressedChunkBytes maximum uncompressed size of any chunk in bytes
+   * @param maxTotalSize Upper bound on the total uncompressed size of all chunks
    * @return number of temporary storage bytes needed to decompress the batch
    */
   static native long batchedZstdDecompressGetTempSize(
       long numChunks,
-      long maxUncompressedChunkBytes);
+      long maxUncompressedChunkBytes,
+      long maxTotalSize);
 
   /**
    * Asynchronously decompress a batch of ZSTD-compressed data buffers.

java/src/main/native/src/NvcompJni.cpp

@@ -63,15 +63,16 @@ extern "C" {
 
 // methods for lz4
 JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_nvcomp_NvcompJni_batchedLZ4CompressGetTempSize(
-  JNIEnv* env, jclass, jlong j_batch_size, jlong j_max_chunk_size)
+  JNIEnv* env, jclass, jlong j_batch_size, jlong j_max_chunk_size, jlong j_max_total_size)
 {
   try {
     cudf::jni::auto_set_device(env);
     auto batch_size       = static_cast<std::size_t>(j_batch_size);
     auto max_chunk_size   = static_cast<std::size_t>(j_max_chunk_size);
+    auto total_size       = static_cast<std::size_t>(j_max_total_size);
     std::size_t temp_size = 0;
-    auto status           = nvcompBatchedLZ4CompressGetTempSize(
-      batch_size, max_chunk_size, nvcompBatchedLZ4DefaultOpts, &temp_size);
+    auto status           = nvcompBatchedLZ4CompressGetTempSizeAsync(
+      batch_size, max_chunk_size, nvcompBatchedLZ4CompressDefaultOpts, &temp_size, total_size);
     check_nvcomp_status(env, status);
     return static_cast<jlong>(temp_size);
   }
@@ -88,7 +89,7 @@ Java_ai_rapids_cudf_nvcomp_NvcompJni_batchedLZ4CompressGetMaxOutputChunkSize(JNI
     auto max_chunk_size         = static_cast<std::size_t>(j_max_chunk_size);
     std::size_t max_output_size = 0;
     auto status                 = nvcompBatchedLZ4CompressGetMaxOutputChunkSize(
-      max_chunk_size, nvcompBatchedLZ4DefaultOpts, &max_output_size);
+      max_chunk_size, nvcompBatchedLZ4CompressDefaultOpts, &max_output_size);
     check_nvcomp_status(env, status);
     return static_cast<jlong>(max_output_size);
   }
@@ -119,30 +120,36 @@ Java_ai_rapids_cudf_nvcomp_NvcompJni_batchedLZ4CompressAsync(JNIEnv* env,
     auto out_ptrs             = reinterpret_cast<void* const*>(j_out_ptrs);
     auto compressed_out_sizes = reinterpret_cast<std::size_t*>(j_compressed_sizes_out_ptr);
     auto stream               = reinterpret_cast<cudaStream_t>(j_stream);
-    auto status               = nvcompBatchedLZ4CompressAsync(in_ptrs,
+    // FIXME how to use these statuses ? They are not used either in the corresponding
+    // decompressor.
+    auto comp_statuses = rmm::device_uvector<nvcompStatus_t>(batch_size, stream);
+    auto status        = nvcompBatchedLZ4CompressAsync(in_ptrs,
                                                 in_sizes,
                                                 chunk_size,
                                                 batch_size,
                                                 temp_ptr,
                                                 temp_size,
                                                 out_ptrs,
                                                 compressed_out_sizes,
-                                                nvcompBatchedLZ4DefaultOpts,
+                                                nvcompBatchedLZ4CompressDefaultOpts,
+                                                comp_statuses.data(),
                                                 stream);
     check_nvcomp_status(env, status);
   }
   CATCH_STD(env, );
 }
 
 JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_nvcomp_NvcompJni_batchedLZ4DecompressGetTempSize(
-  JNIEnv* env, jclass, jlong j_batch_size, jlong j_chunk_size)
+  JNIEnv* env, jclass, jlong j_batch_size, jlong j_chunk_size, jlong j_max_total_size)
 {
   try {
     cudf::jni::auto_set_device(env);
     auto batch_size       = static_cast<std::size_t>(j_batch_size);
     auto chunk_size       = static_cast<std::size_t>(j_chunk_size);
+    auto total_size       = static_cast<std::size_t>(j_max_total_size);
     std::size_t temp_size = 0;
-    auto status = nvcompBatchedLZ4DecompressGetTempSize(batch_size, chunk_size, &temp_size);
+    auto status           = nvcompBatchedLZ4DecompressGetTempSizeAsync(
+      batch_size, chunk_size, nvcompBatchedLZ4DecompressDefaultOpts, &temp_size, total_size);
     check_nvcomp_status(env, status);
     return static_cast<jlong>(temp_size);
   }
@@ -181,6 +188,7 @@ Java_ai_rapids_cudf_nvcomp_NvcompJni_batchedLZ4DecompressAsync(JNIEnv* env,
                                                   temp_ptr,
                                                   temp_size,
                                                   uncompressed_ptrs,
+                                                  nvcompBatchedLZ4DecompressDefaultOpts,
                                                   uncompressed_statuses.data(),
                                                   stream);
     check_nvcomp_status(env, status);
@@ -218,15 +226,16 @@ Java_ai_rapids_cudf_nvcomp_NvcompJni_batchedLZ4GetDecompressSizeAsync(JNIEnv* en
 
 // methods for zstd
 JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_nvcomp_NvcompJni_batchedZstdCompressGetTempSize(
-  JNIEnv* env, jclass, jlong j_batch_size, jlong j_max_chunk_size)
+  JNIEnv* env, jclass, jlong j_batch_size, jlong j_max_chunk_size, jlong j_max_total_size)
 {
   try {
     cudf::jni::auto_set_device(env);
     auto batch_size       = static_cast<std::size_t>(j_batch_size);
     auto max_chunk_size   = static_cast<std::size_t>(j_max_chunk_size);
+    auto total_size       = static_cast<std::size_t>(j_max_total_size);
     std::size_t temp_size = 0;
-    auto status           = nvcompBatchedZstdCompressGetTempSize(
-      batch_size, max_chunk_size, nvcompBatchedZstdDefaultOpts, &temp_size);
+    auto status           = nvcompBatchedZstdCompressGetTempSizeAsync(
+      batch_size, max_chunk_size, nvcompBatchedZstdCompressDefaultOpts, &temp_size, total_size);
     check_nvcomp_status(env, status);
     return static_cast<jlong>(temp_size);
   }
@@ -242,7 +251,7 @@ Java_ai_rapids_cudf_nvcomp_NvcompJni_batchedZstdCompressGetMaxOutputChunkSize(
     auto max_chunk_size         = static_cast<std::size_t>(j_max_chunk_size);
     std::size_t max_output_size = 0;
     auto status                 = nvcompBatchedZstdCompressGetMaxOutputChunkSize(
-      max_chunk_size, nvcompBatchedZstdDefaultOpts, &max_output_size);
+      max_chunk_size, nvcompBatchedZstdCompressDefaultOpts, &max_output_size);
     check_nvcomp_status(env, status);
     return static_cast<jlong>(max_output_size);
   }
@@ -273,30 +282,36 @@ Java_ai_rapids_cudf_nvcomp_NvcompJni_batchedZstdCompressAsync(JNIEnv* env,
     auto out_ptrs             = reinterpret_cast<void* const*>(j_out_ptrs);
     auto compressed_out_sizes = reinterpret_cast<std::size_t*>(j_compressed_sizes_out_ptr);
     auto stream               = reinterpret_cast<cudaStream_t>(j_stream);
-    auto status               = nvcompBatchedZstdCompressAsync(in_ptrs,
+    // FIXME how to use these statuses ? They are not used either in the corresponding
+    // decompressor.
+    auto comp_statuses = rmm::device_uvector<nvcompStatus_t>(batch_size, stream);
+    auto status        = nvcompBatchedZstdCompressAsync(in_ptrs,
                                                  in_sizes,
                                                  chunk_size,
                                                  batch_size,
                                                  temp_ptr,
                                                  temp_size,
                                                  out_ptrs,
                                                  compressed_out_sizes,
-                                                 nvcompBatchedZstdDefaultOpts,
+                                                 nvcompBatchedZstdCompressDefaultOpts,
+                                                 comp_statuses.data(),
                                                  stream);
     check_nvcomp_status(env, status);
   }
   CATCH_STD(env, );
 }
 
 JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_nvcomp_NvcompJni_batchedZstdDecompressGetTempSize(
-  JNIEnv* env, jclass, jlong j_batch_size, jlong j_chunk_size)
+  JNIEnv* env, jclass, jlong j_batch_size, jlong j_chunk_size, jlong j_max_total_size)
 {
   try {
     cudf::jni::auto_set_device(env);
     auto batch_size       = static_cast<std::size_t>(j_batch_size);
     auto chunk_size       = static_cast<std::size_t>(j_chunk_size);
+    auto total_size       = static_cast<std::size_t>(j_max_total_size);
     std::size_t temp_size = 0;
-    auto status = nvcompBatchedZstdDecompressGetTempSize(batch_size, chunk_size, &temp_size);
+    auto status           = nvcompBatchedZstdDecompressGetTempSizeAsync(
+      batch_size, chunk_size, nvcompBatchedZstdDecompressDefaultOpts, &temp_size, total_size);
     check_nvcomp_status(env, status);
     return static_cast<jlong>(temp_size);
   }
@@ -335,6 +350,7 @@ Java_ai_rapids_cudf_nvcomp_NvcompJni_batchedZstdDecompressAsync(JNIEnv* env,
                                                    temp_ptr,
                                                    temp_size,
                                                    uncompressed_ptrs,
+                                                   nvcompBatchedZstdDecompressDefaultOpts,
                                                    uncompressed_statuses.data(),
                                                    stream);
     check_nvcomp_status(env, status);