Skip to content

Commit 0337847

Browse files
madsbkmadsbk
committed
removed stream and device_id
1 parent 2a8f203 commit 0337847

File tree

3 files changed

+21
-78
lines changed

3 files changed

+21
-78
lines changed

‎python/kvikio/kvikio/_lib/libnvcomp.pyx

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -150,8 +150,6 @@ cdef class _ANSManager(_nvcompManager):
150150
def __cinit__(
151151
self,
152152
size_t uncomp_chunk_size,
153-
user_stream,
154-
const int device_id,
155153
):
156154
self._impl = <nvcompManagerBase*>new ANSManager(
157155
uncomp_chunk_size,
@@ -165,8 +163,6 @@ cdef class _BitcompManager(_nvcompManager):
165163
size_t uncomp_chunk_size,
166164
nvcompType_t data_type,
167165
int bitcomp_algo,
168-
user_stream,
169-
const int device_id
170166
):
171167
cdef opts = nvcompBatchedBitcompFormatOpts(bitcomp_algo, data_type)
172168
self._impl = <nvcompManagerBase*>new BitcompManager(
@@ -179,8 +175,6 @@ cdef class _CascadedManager(_nvcompManager):
179175
def __cinit__(
180176
self,
181177
_options,
182-
user_stream,
183-
const int device_id,
184178
):
185179
self._impl = <nvcompManagerBase*>new CascadedManager(
186180
_options["chunk_size"],
@@ -193,8 +187,6 @@ cdef class _GdeflateManager(_nvcompManager):
193187
self,
194188
size_t uncomp_chunk_size,
195189
int algo,
196-
user_stream,
197-
const int device_id
198190
):
199191
cdef opts = nvcompBatchedGdeflateOpts_t(algo)
200192
self._impl = <nvcompManagerBase*>new GdeflateManager(
@@ -208,8 +200,6 @@ cdef class _LZ4Manager(_nvcompManager):
208200
self,
209201
size_t uncomp_chunk_size,
210202
nvcompType_t data_type,
211-
user_stream,
212-
const int device_id,
213203
):
214204
# TODO: Doesn't work with user specified streams passed down
215205
# from anywhere up. I'm not going to rabbit hole on it until
@@ -226,8 +216,6 @@ cdef class _SnappyManager(_nvcompManager):
226216
def __cinit__(
227217
self,
228218
size_t uncomp_chunk_size,
229-
user_stream,
230-
const int device_id,
231219
):
232220
# TODO: Doesn't work with user specified streams passed down
233221
# from anywhere up. I'm not going to rabbit hole on it until

‎python/kvikio/kvikio/nvcomp.py

Lines changed: 6 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved.
1+
# Copyright (c) 2021-2025, NVIDIA CORPORATION. All rights reserved.
22
# See file LICENSE for terms.
33

44
from enum import Enum
@@ -62,12 +62,10 @@ class nvCompManager:
6262
# Default options exist for every option type for every class that inherits
6363
# from nvCompManager, which takes advantage of the below property-setting
6464
# code.
65-
stream: cp.cuda.Stream = cp.cuda.Stream()
6665
chunk_size: int = 1 << 16
6766
data_type: _lib.pyNvcompType_t = _lib.pyNvcompType_t.pyNVCOMP_TYPE_UCHAR
6867
# Some classes have this defined as type, some as data_type.
6968
type: _lib.pyNvcompType_t = _lib.pyNvcompType_t.pyNVCOMP_TYPE_UCHAR
70-
device_id: int = 0
7169

7270
# Bitcomp Defaults
7371
bitcomp_algo: int = 0
@@ -84,12 +82,6 @@ def __init__(self, kwargs):
8482
8583
Special case: Convert data_type to a _lib.pyNvcompType_t
8684
"""
87-
# Special case: Throw error if stream or device_id are specified
88-
if kwargs.get("stream") is not None:
89-
raise NotImplementedError(
90-
"stream argument not yet supported: " "Use the default argument"
91-
)
92-
9385
# data_type will be passed in as a python object. Convert it to
9486
# a C++ nvcompType_t here.
9587
if kwargs.get("data_type"):
@@ -221,13 +213,10 @@ def __init__(self, **kwargs):
221213
----------
222214
chunk_size: int (optional)
223215
Defaults to 4096.
224-
device_id: int (optional)
225-
Specify which device_id on the node to use for allocation and compression.
226-
Defaults to 0.
227216
"""
228217
super().__init__(kwargs)
229218

230-
self._manager = _lib._ANSManager(self.chunk_size, self.stream, self.device_id)
219+
self._manager = _lib._ANSManager(self.chunk_size)
231220

232221

233222
class BitcompManager(nvCompManager):
@@ -241,18 +230,13 @@ def __init__(self, **kwargs):
241230
----------
242231
chunk_size: int (optional)
243232
Defaults to 4096.
244-
device_id: int (optional)
245-
Specify which device_id on the node to use
246-
Defaults to 0.
247233
"""
248234
super().__init__(kwargs)
249235

250236
self._manager = _lib._BitcompManager(
251237
self.chunk_size,
252238
self.data_type.value,
253239
self.bitcomp_algo,
254-
self.stream,
255-
self.device_id,
256240
)
257241

258242

@@ -278,9 +262,6 @@ def __init__(self, **kwargs):
278262
use_bp: bool (optional)
279263
Enable Bitpacking, see [algorithms overview.md](
280264
https://github.com/NVIDIA/nvcomp/blob/main/doc/algorithms_overview.md#bitpacking) # noqa: E501
281-
device_id: int (optional)
282-
Specify which device_id on the node to use
283-
Defaults to 0.
284265
"""
285266
super().__init__(kwargs)
286267
default_options = {
@@ -304,9 +285,7 @@ def __init__(self, **kwargs):
304285
"num_deltas": self.num_deltas,
305286
"use_bp": self.use_bp,
306287
}
307-
self._manager = _lib._CascadedManager(
308-
default_options, self.stream, self.device_id
309-
)
288+
self._manager = _lib._CascadedManager(default_options)
310289

311290

312291
class GdeflateManager(nvCompManager):
@@ -322,18 +301,10 @@ def __init__(self, **kwargs):
322301
algo: int (optional)
323302
Integer in the range [0, 1, 2]. Only algorithm #0 is currently
324303
supported.
325-
stream: cudaStream_t (optional)
326-
Which CUDA stream to perform the operation on. Not currently
327-
supported.
328-
device_id: int (optional)
329-
Specify which device_id on the node to use
330-
Defaults to 0.
331304
"""
332305
super().__init__(kwargs)
333306

334-
self._manager = _lib._GdeflateManager(
335-
self.chunk_size, self.algo, self.stream, self.device_id
336-
)
307+
self._manager = _lib._GdeflateManager(self.chunk_size, self.algo)
337308

338309

339310
class LZ4Manager(nvCompManager):
@@ -354,17 +325,9 @@ def __init__(self, **kwargs):
354325
data_type: pyNVCOMP_TYPE (optional)
355326
The data type returned for decompression.
356327
Defaults to pyNVCOMP_TYPE.UCHAR
357-
stream: cudaStream_t (optional)
358-
Which CUDA stream to perform the operation on. Not currently
359-
supported.
360-
device_id: int (optional)
361-
Specify which device_id on the node to use
362-
Defaults to 0.
363328
"""
364329
super().__init__(kwargs)
365-
self._manager = _lib._LZ4Manager(
366-
self.chunk_size, self.data_type.value, self.stream, self.device_id
367-
)
330+
self._manager = _lib._LZ4Manager(self.chunk_size, self.data_type.value)
368331

369332

370333
class SnappyManager(nvCompManager):
@@ -377,17 +340,9 @@ def __init__(self, **kwargs):
377340
Parameters
378341
----------
379342
chunk_size: int (optional)
380-
stream: cudaStream_t (optional)
381-
Which CUDA stream to perform the operation on. Not currently
382-
supported.
383-
device_id: int (optional)
384-
Specify which device_id on the node to use
385-
Defaults to 0.
386343
"""
387344
super().__init__(kwargs)
388-
self._manager = _lib._SnappyManager(
389-
self.chunk_size, self.stream, self.device_id
390-
)
345+
self._manager = _lib._SnappyManager(self.chunk_size)
391346

392347

393348
class ManagedDecompressionManager(nvCompManager):

‎python/kvikio/tests/test_nvcomp.py

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved.
1+
# Copyright (c) 2021-2025, NVIDIA CORPORATION. All rights reserved.
22
# See file LICENSE for terms.
33

44
import pytest
@@ -77,12 +77,11 @@ def test_round_trip_dtypes(manager, dtype):
7777
"inputs",
7878
[
7979
{},
80-
{"chunk_size": 1 << 16, "device_id": 0},
8180
{
8281
"chunk_size": 1 << 16,
8382
},
8483
{
85-
"device_id": 0,
84+
"chunk_size": 1 << 16,
8685
},
8786
],
8887
)
@@ -99,13 +98,13 @@ def test_ans_inputs(inputs):
9998
"inputs",
10099
[
101100
{},
102-
{"data_type": np.uint8, "algo": 0, "device_id": 0},
103-
{"data_type": np.uint8},
104101
{
102+
"data_type": np.uint8,
105103
"algo": 0,
106104
},
105+
{"data_type": np.uint8},
107106
{
108-
"device_id": 0,
107+
"algo": 0,
109108
},
110109
],
111110
)
@@ -179,7 +178,6 @@ def test_bitcomp_algorithms(inputs, expected):
179178
"num_deltas": 1,
180179
"use_bp": True,
181180
},
182-
"device_id": 0,
183181
},
184182
],
185183
)
@@ -196,15 +194,15 @@ def test_cascaded_inputs(inputs):
196194
"inputs",
197195
[
198196
{},
199-
{"chunk_size": 1 << 16, "algo": 0, "device_id": 0},
200197
{
201198
"chunk_size": 1 << 16,
199+
"algo": 0,
202200
},
203201
{
204-
"algo": 0,
202+
"chunk_size": 1 << 16,
205203
},
206204
{
207-
"device_id": 0,
205+
"algo": 0,
208206
},
209207
],
210208
)
@@ -253,15 +251,15 @@ def test_gdeflate_algorithms_not_implemented(inputs, expected):
253251
"inputs",
254252
[
255253
{},
256-
{"chunk_size": 1 << 16, "data_type": np.uint8, "device_id": 0},
257254
{
258255
"chunk_size": 1 << 16,
256+
"data_type": np.uint8,
259257
},
260258
{
261-
"data_type": np.uint8,
259+
"chunk_size": 1 << 16,
262260
},
263261
{
264-
"device_id": 0,
262+
"data_type": np.uint8,
265263
},
266264
],
267265
)
@@ -278,11 +276,13 @@ def test_lz4_inputs(inputs):
278276
"inputs",
279277
[
280278
{},
281-
{"chunk_size": 1 << 16, "device_id": 0},
282279
{
283280
"chunk_size": 1 << 16,
284281
},
285-
{"device_id": 0},
282+
{
283+
"chunk_size": 1 << 16,
284+
},
285+
{},
286286
],
287287
)
288288
def test_snappy_inputs(inputs):

0 commit comments

Comments
 (0)