diff --git a/python/kvikio/kvikio/_lib/libnvcomp.pyx b/python/kvikio/kvikio/_lib/libnvcomp.pyx index ddb890a943..dc5359e9b3 100644 --- a/python/kvikio/kvikio/_lib/libnvcomp.pyx +++ b/python/kvikio/kvikio/_lib/libnvcomp.pyx @@ -150,8 +150,6 @@ cdef class _ANSManager(_nvcompManager): def __cinit__( self, size_t uncomp_chunk_size, - user_stream, - const int device_id, ): self._impl = new ANSManager( uncomp_chunk_size, @@ -165,8 +163,6 @@ cdef class _BitcompManager(_nvcompManager): 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 = new BitcompManager( @@ -179,8 +175,6 @@ cdef class _CascadedManager(_nvcompManager): def __cinit__( self, _options, - user_stream, - const int device_id, ): self._impl = new CascadedManager( _options["chunk_size"], @@ -193,8 +187,6 @@ cdef class _GdeflateManager(_nvcompManager): self, size_t uncomp_chunk_size, int algo, - user_stream, - const int device_id ): cdef opts = nvcompBatchedGdeflateOpts_t(algo) self._impl = new GdeflateManager( @@ -208,8 +200,6 @@ cdef class _LZ4Manager(_nvcompManager): self, size_t uncomp_chunk_size, nvcompType_t data_type, - user_stream, - const int device_id, ): # TODO: Doesn't work with user specified streams passed down # from anywhere up. I'm not going to rabbit hole on it until @@ -226,8 +216,6 @@ cdef class _SnappyManager(_nvcompManager): def __cinit__( self, size_t uncomp_chunk_size, - user_stream, - const int device_id, ): # TODO: Doesn't work with user specified streams passed down # from anywhere up. I'm not going to rabbit hole on it until diff --git a/python/kvikio/kvikio/nvcomp.py b/python/kvikio/kvikio/nvcomp.py index 482d7fbeb7..5606ad5ce5 100644 --- a/python/kvikio/kvikio/nvcomp.py +++ b/python/kvikio/kvikio/nvcomp.py @@ -1,4 +1,4 @@ -# Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2021-2025, NVIDIA CORPORATION. All rights reserved. # See file LICENSE for terms. from enum import Enum @@ -62,12 +62,10 @@ class nvCompManager: # Default options exist for every option type for every class that inherits # from nvCompManager, which takes advantage of the below property-setting # code. - stream: cp.cuda.Stream = cp.cuda.Stream() chunk_size: int = 1 << 16 data_type: _lib.pyNvcompType_t = _lib.pyNvcompType_t.pyNVCOMP_TYPE_UCHAR # Some classes have this defined as type, some as data_type. type: _lib.pyNvcompType_t = _lib.pyNvcompType_t.pyNVCOMP_TYPE_UCHAR - device_id: int = 0 # Bitcomp Defaults bitcomp_algo: int = 0 @@ -84,12 +82,6 @@ def __init__(self, kwargs): Special case: Convert data_type to a _lib.pyNvcompType_t """ - # Special case: Throw error if stream or device_id are specified - if kwargs.get("stream") is not None: - raise NotImplementedError( - "stream argument not yet supported: " "Use the default argument" - ) - # data_type will be passed in as a python object. Convert it to # a C++ nvcompType_t here. if kwargs.get("data_type"): @@ -221,13 +213,10 @@ def __init__(self, **kwargs): ---------- chunk_size: int (optional) Defaults to 4096. - device_id: int (optional) - Specify which device_id on the node to use for allocation and compression. - Defaults to 0. """ super().__init__(kwargs) - self._manager = _lib._ANSManager(self.chunk_size, self.stream, self.device_id) + self._manager = _lib._ANSManager(self.chunk_size) class BitcompManager(nvCompManager): @@ -241,9 +230,6 @@ def __init__(self, **kwargs): ---------- chunk_size: int (optional) Defaults to 4096. - device_id: int (optional) - Specify which device_id on the node to use - Defaults to 0. """ super().__init__(kwargs) @@ -251,8 +237,6 @@ def __init__(self, **kwargs): self.chunk_size, self.data_type.value, self.bitcomp_algo, - self.stream, - self.device_id, ) @@ -278,9 +262,6 @@ def __init__(self, **kwargs): use_bp: bool (optional) Enable Bitpacking, see [algorithms overview.md]( https://github.com/NVIDIA/nvcomp/blob/main/doc/algorithms_overview.md#bitpacking) # noqa: E501 - device_id: int (optional) - Specify which device_id on the node to use - Defaults to 0. """ super().__init__(kwargs) default_options = { @@ -304,9 +285,7 @@ def __init__(self, **kwargs): "num_deltas": self.num_deltas, "use_bp": self.use_bp, } - self._manager = _lib._CascadedManager( - default_options, self.stream, self.device_id - ) + self._manager = _lib._CascadedManager(default_options) class GdeflateManager(nvCompManager): @@ -322,18 +301,10 @@ def __init__(self, **kwargs): algo: int (optional) Integer in the range [0, 1, 2]. Only algorithm #0 is currently supported. - stream: cudaStream_t (optional) - Which CUDA stream to perform the operation on. Not currently - supported. - device_id: int (optional) - Specify which device_id on the node to use - Defaults to 0. """ super().__init__(kwargs) - self._manager = _lib._GdeflateManager( - self.chunk_size, self.algo, self.stream, self.device_id - ) + self._manager = _lib._GdeflateManager(self.chunk_size, self.algo) class LZ4Manager(nvCompManager): @@ -354,17 +325,9 @@ def __init__(self, **kwargs): data_type: pyNVCOMP_TYPE (optional) The data type returned for decompression. Defaults to pyNVCOMP_TYPE.UCHAR - stream: cudaStream_t (optional) - Which CUDA stream to perform the operation on. Not currently - supported. - device_id: int (optional) - Specify which device_id on the node to use - Defaults to 0. """ super().__init__(kwargs) - self._manager = _lib._LZ4Manager( - self.chunk_size, self.data_type.value, self.stream, self.device_id - ) + self._manager = _lib._LZ4Manager(self.chunk_size, self.data_type.value) class SnappyManager(nvCompManager): @@ -377,17 +340,9 @@ def __init__(self, **kwargs): Parameters ---------- chunk_size: int (optional) - stream: cudaStream_t (optional) - Which CUDA stream to perform the operation on. Not currently - supported. - device_id: int (optional) - Specify which device_id on the node to use - Defaults to 0. """ super().__init__(kwargs) - self._manager = _lib._SnappyManager( - self.chunk_size, self.stream, self.device_id - ) + self._manager = _lib._SnappyManager(self.chunk_size) class ManagedDecompressionManager(nvCompManager): diff --git a/python/kvikio/tests/test_nvcomp.py b/python/kvikio/tests/test_nvcomp.py index a2ea895dd4..356c5e77cd 100644 --- a/python/kvikio/tests/test_nvcomp.py +++ b/python/kvikio/tests/test_nvcomp.py @@ -1,4 +1,4 @@ -# Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2021-2025, NVIDIA CORPORATION. All rights reserved. # See file LICENSE for terms. import pytest @@ -77,12 +77,11 @@ def test_round_trip_dtypes(manager, dtype): "inputs", [ {}, - {"chunk_size": 1 << 16, "device_id": 0}, { "chunk_size": 1 << 16, }, { - "device_id": 0, + "chunk_size": 1 << 16, }, ], ) @@ -99,13 +98,13 @@ def test_ans_inputs(inputs): "inputs", [ {}, - {"data_type": np.uint8, "algo": 0, "device_id": 0}, - {"data_type": np.uint8}, { + "data_type": np.uint8, "algo": 0, }, + {"data_type": np.uint8}, { - "device_id": 0, + "algo": 0, }, ], ) @@ -179,7 +178,6 @@ def test_bitcomp_algorithms(inputs, expected): "num_deltas": 1, "use_bp": True, }, - "device_id": 0, }, ], ) @@ -196,15 +194,15 @@ def test_cascaded_inputs(inputs): "inputs", [ {}, - {"chunk_size": 1 << 16, "algo": 0, "device_id": 0}, { "chunk_size": 1 << 16, + "algo": 0, }, { - "algo": 0, + "chunk_size": 1 << 16, }, { - "device_id": 0, + "algo": 0, }, ], ) @@ -253,15 +251,15 @@ def test_gdeflate_algorithms_not_implemented(inputs, expected): "inputs", [ {}, - {"chunk_size": 1 << 16, "data_type": np.uint8, "device_id": 0}, { "chunk_size": 1 << 16, + "data_type": np.uint8, }, { - "data_type": np.uint8, + "chunk_size": 1 << 16, }, { - "device_id": 0, + "data_type": np.uint8, }, ], ) @@ -278,11 +276,13 @@ def test_lz4_inputs(inputs): "inputs", [ {}, - {"chunk_size": 1 << 16, "device_id": 0}, { "chunk_size": 1 << 16, }, - {"device_id": 0}, + { + "chunk_size": 1 << 16, + }, + {}, ], ) def test_snappy_inputs(inputs):