From 33e9e495447851ac603ad04e372624cd354dfbb3 Mon Sep 17 00:00:00 2001
From: Vukasin Milovanovic <vmilovanovic@nvidia.com>
Date: Fri, 3 Nov 2023 13:46:09 -0700
Subject: [PATCH] remove

---
 python/kvikio/_lib/libnvcomp.pyx      |  8 ------
 python/kvikio/_lib/nvcomp_cxx_api.pxd |  4 ---
 python/kvikio/nvcomp.py               | 28 ---------------------
 python/tests/test_nvcomp.py           | 36 ++-------------------------
 4 files changed, 2 insertions(+), 74 deletions(-)

diff --git a/python/kvikio/_lib/libnvcomp.pyx b/python/kvikio/_lib/libnvcomp.pyx
index 19237d24ca..7e87ff37a0 100644
--- a/python/kvikio/_lib/libnvcomp.pyx
+++ b/python/kvikio/_lib/libnvcomp.pyx
@@ -134,14 +134,6 @@ cdef class _nvcompManager:
             <DecompressionConfig&>self._decompression_config.get()[0]
         )
 
-    def set_scratch_buffer(self, Array new_scratch_buffer):
-        return self._impl.set_scratch_buffer(
-            <uint8_t*>new_scratch_buffer.ptr
-        )
-
-    def get_required_scratch_buffer_size(self):
-        return self._impl.get_required_scratch_buffer_size()
-
     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
index e5b464d5c2..a50967134e 100644
--- a/python/kvikio/_lib/nvcomp_cxx_api.pxd
+++ b/python/kvikio/_lib/nvcomp_cxx_api.pxd
@@ -106,8 +106,6 @@ cdef extern from "nvcomp/nvcompManager.hpp" namespace 'nvcomp':
             uint8_t* decomp_buffer,
             const uint8_t* comp_buffer,
             const DecompressionConfig& decomp_config)
-        void set_scratch_buffer(uint8_t* new_scratch_buffer) except +
-        size_t get_required_scratch_buffer_size() except +
         size_t get_compressed_output_size(uint8_t* comp_buffer) except +
 
     cdef cppclass PimplManager "nvcomp::PimplManager":
@@ -125,8 +123,6 @@ cdef extern from "nvcomp/nvcompManager.hpp" namespace 'nvcomp':
             uint8_t* decomp_buffer,
             const uint8_t* comp_buffer,
             const DecompressionConfig& decomp_config) except +
-        void set_scratch_buffer(uint8_t* new_scratch_buffer) except +
-        size_t get_required_scratch_buffer_size() except +
         size_t get_compressed_output_size(uint8_t* comp_buffer) except +
 
 # C++ Concrete ANS Manager
diff --git a/python/kvikio/nvcomp.py b/python/kvikio/nvcomp.py
index 38a326d410..190786d12c 100644
--- a/python/kvikio/nvcomp.py
+++ b/python/kvikio/nvcomp.py
@@ -192,34 +192,6 @@ def configure_decompression_with_compressed_buffer(
             asarray(data)
         )
 
-    def get_required_scratch_buffer_size(self) -> int:
-        """Return the size of the optional scratch buffer.
-
-        Returns
-        -------
-        int
-        """
-        return self._manager.get_required_scratch_buffer_size()
-
-    def set_scratch_buffer(self, new_scratch_buffer: cp.ndarray) -> None:
-        """Use a pre-allocated buffer for compression.
-
-        Use a GPU-allocated buffer that will be used for compression
-        temporary storage instead of allowing the library to create the
-        scratch buffer.
-        Can reduce memory usage.
-
-        Parameters
-        ----------
-        new_scratch_buffer : cp.ndarray
-            The buffer that you allocated on the GPU for compressor temporary
-            storage.
-
-        Returns
-        -------
-        cp.ndarray
-        """
-        return self._manager.set_scratch_buffer(asarray(new_scratch_buffer))
 
     def get_compressed_output_size(self, comp_buffer: cp.ndarray) -> int:
         """Return the actual size of compression result.
diff --git a/python/tests/test_nvcomp.py b/python/tests/test_nvcomp.py
index 59ba24869f..951f15a1e0 100644
--- a/python/tests/test_nvcomp.py
+++ b/python/tests/test_nvcomp.py
@@ -407,7 +407,7 @@ def test_get_decompression_config_with_default_options(manager, expected):
 
 
 @pytest.mark.parametrize("manager", managers())
-def test_set_scratch_buffer(manager):
+def test_set_scratch_alloc(manager):
     length = 10000
     dtype = cupy.uint8
     data = cupy.array(
@@ -419,9 +419,7 @@ def test_set_scratch_buffer(manager):
     )
     compressor_instance = manager()
     compressor_instance.configure_compression(len(data))
-    buffer_size = compressor_instance.get_required_scratch_buffer_size()
-    buffer = cupy.zeros(buffer_size, dtype="int8")
-    compressor_instance.set_scratch_buffer(buffer)
+    # TODO set_scratch_allocators
     compressor_instance.compress(data)
     if isinstance(compressor_instance, libnvcomp.BitcompManager):
         # Bitcomp does not use the scratch buffer
@@ -430,36 +428,6 @@ def test_set_scratch_buffer(manager):
         assert (buffer[0:5] != cupy.array([0, 0, 0, 0, 0])).any()
 
 
-@pytest.mark.parametrize(
-    "manager,expected",
-    zip(
-        managers(),
-        [
-            378355712,  # ANS
-            8,  # Bitcomp
-            1641608,  # Cascaded
-            393222400,  # Gdeflate
-            252334080,  # LZ4
-            67311208,  # Snappy
-        ],
-    ),
-)
-def test_get_required_scratch_buffer_size(manager, expected):
-    length = 10000
-    dtype = cupy.uint8
-    data = cupy.array(
-        np.arange(
-            0,
-            length // cupy.dtype(dtype).type(0).itemsize,
-            dtype=dtype,
-        )
-    )
-    compressor_instance = manager()
-    compressor_instance.configure_compression(len(data))
-    buffer_size = compressor_instance.get_required_scratch_buffer_size()
-    assert_compression_size(buffer_size, expected)
-
-
 @pytest.mark.parametrize(
     "manager, expected",
     zip(managers(), list(LEN.values())),