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