def test_rmm_pool_cupy_allocator_with_stream(stream):
cupy = pytest.importorskip("cupy")
rmm.reinitialize(pool_allocator=True)
cupy.cuda.set_allocator(rmm.rmm_cupy_allocator)
if stream == "null":
stream = cupy.cuda.stream.Stream.null
else:
stream = cupy.cuda.stream.Stream()
with stream:
m = rmm.rmm_cupy_allocator(42)
assert m.mem.size == 42
assert m.mem.ptr != 0
assert isinstance(m.mem._owner, rmm.DeviceBuffer)
m = rmm.rmm_cupy_allocator(0)
assert m.mem.size == 0
assert m.mem.ptr == 0
assert isinstance(m.mem._owner, rmm.DeviceBuffer)
a = cupy.arange(10)
assert isinstance(a.data.mem._owner, rmm.DeviceBuffer)
# Deleting all allocations known by the RMM pool is required
# before rmm.reinitialize(), otherwise it may segfault.
del a
rmm.reinitialize()
def test_rmm_cupy_allocator():
cupy = pytest.importorskip("cupy")
m = rmm.rmm_cupy_allocator(42)
assert m.mem.size == 42
assert m.mem.ptr != 0
assert m.mem.rmm_array is not None
m = rmm.rmm_cupy_allocator(0)
assert m.mem.size == 0
assert m.mem.ptr == 0
assert m.mem.rmm_array is None
cupy.cuda.set_allocator(rmm.rmm_cupy_allocator)
a = cupy.arange(10)
assert hasattr(a.data.mem, "rmm_array")
def test_rmm_cupy_allocator():
cupy = pytest.importorskip("cupy")
m = rmm.rmm_cupy_allocator(42)
assert m.mem.size == 42
assert m.mem.ptr != 0
assert isinstance(m.mem._owner, rmm.DeviceBuffer)
m = rmm.rmm_cupy_allocator(0)
assert m.mem.size == 0
assert m.mem.ptr == 0
assert isinstance(m.mem._owner, rmm.DeviceBuffer)
cupy.cuda.set_allocator(rmm.rmm_cupy_allocator)
a = cupy.arange(10)
assert isinstance(a.data.mem._owner, rmm.DeviceBuffer)
def _rmm_cupy6_array_like(ary, order):
nbytes = np.ndarray(ary.shape,
dtype=ary.dtype,
strides=ary.strides,
order=order).nbytes
memptr = rmm.rmm_cupy_allocator(nbytes)
arr = cp.ndarray(ary.shape,
dtype=ary.dtype,
memptr=memptr,
strides=ary.strides,
order=order)
return arr
