bitcomp fix

rapidsai · Nov 7, 2023 · 6264a64 · 6264a64
1 parent 4ad2f06
commit 6264a64
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 3 deletions.
diff --git a/python/kvikio/_lib/libnvcomp.pyx b/python/kvikio/_lib/libnvcomp.pyx
@@ -168,14 +168,16 @@ cdef class _ANSManager(_nvcompManager):
 cdef class _BitcompManager(_nvcompManager):
     def __cinit__(
         self,
+        size_t uncomp_chunk_size,
         nvcompType_t data_type,
         int bitcomp_algo,
         user_stream,
         const int device_id
     ):
+        cdef opts = nvcompBatchedBitcompFormatOpts(bitcomp_algo, data_type)
         self._impl = <nvcompManagerBase*>new BitcompManager(
-            0,  # TODO
-            <nvcompBatchedBitcompFormatOpts>nvcompBatchedBitcompDefaultOpts,  # TODO
+            uncomp_chunk_size,
+            opts,
             <cudaStream_t><void*>0,  # TODO
             device_id
         )

diff --git a/python/kvikio/_lib/nvcomp_cxx_api.pxd b/python/kvikio/_lib/nvcomp_cxx_api.pxd
@@ -146,6 +146,7 @@ cdef extern from "nvcomp/ans.hpp":
 # C++ Concrete Bitcomp Manager
 cdef extern from "nvcomp/bitcomp.h" nogil:
     ctypedef struct nvcompBatchedBitcompFormatOpts:
+        int algorithm_type
         nvcompType_t data_type
     cdef nvcompBatchedBitcompFormatOpts nvcompBatchedBitcompDefaultOpts
 

diff --git a/python/kvikio/nvcomp.py b/python/kvikio/nvcomp.py
@@ -248,7 +248,11 @@ def __init__(self, **kwargs):
         super().__init__(kwargs)
 
         self._manager = _lib._BitcompManager(
-            self.data_type.value, self.bitcomp_algo, self.stream, self.device_id
+            self.chunk_size,
+            self.data_type.value,
+            self.bitcomp_algo,
+            self.stream,
+            self.device_id,
         )