add most of set_scratch_allocators

rapidsai · Nov 3, 2023 · 8866101 · 8866101
1 parent 33e9e49
commit 8866101
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 0 deletions.
diff --git a/python/kvikio/_lib/libnvcomp.pyx b/python/kvikio/_lib/libnvcomp.pyx
@@ -29,10 +29,12 @@ from libcpp.utility cimport move
 
 from kvikio._lib.arr cimport Array
 from kvikio._lib.nvcomp_cxx_api cimport (
+    AllocFn_t,
     ANSManager,
     BitcompManager,
     CascadedManager,
     CompressionConfig,
+    DeAllocFn_t,
     DecompressionConfig,
     GdeflateManager,
     LZ4Manager,
@@ -134,6 +136,16 @@ cdef class _nvcompManager:
             <DecompressionConfig&>self._decompression_config.get()[0]
         )
 
+    def set_scratch_allocators(
+        self,
+        alloc_fn,
+        dealloc_fn
+    ):
+        self._impl.set_scratch_allocators(
+            <AllocFn_t>(alloc_fn),
+            <DeAllocFn_t>(dealloc_fn)
+        )
+
     def get_compressed_output_size(self, Array comp_buffer):
         return self._impl.get_compressed_output_size(
             <uint8_t*>comp_buffer.ptr

diff --git a/python/kvikio/_lib/nvcomp_cxx_api.pxd b/python/kvikio/_lib/nvcomp_cxx_api.pxd
@@ -61,6 +61,9 @@ cdef extern from "nvcomp/nvcompManagerFactory.hpp" namespace 'nvcomp':
 
 # Compression Manager
 cdef extern from "nvcomp/nvcompManager.hpp" namespace 'nvcomp':
+    ctypedef void* (*AllocFn_t)(size_t)
+    ctypedef void (*DeAllocFn_t)(void*, size_t)
+
     cdef cppclass PinnedPtrPool[T]:
         pass
 

diff --git a/python/kvikio/nvcomp.py b/python/kvikio/nvcomp.py
@@ -192,6 +192,7 @@ def configure_decompression_with_compressed_buffer(
             asarray(data)
         )
 
+    # TODO def set_scratch_allocators
 
     def get_compressed_output_size(self, comp_buffer: cp.ndarray) -> int:
         """Return the actual size of compression result.