def memalloc(self, size):
"""
Allocate an on-device array from the RMM pool.
"""
buf = librmm.DeviceBuffer(size=size)
ctx = self.context
ptr = ctypes.c_uint64(int(buf.ptr))
finalizer = _make_emm_plugin_finalizer(ptr.value, self.allocations)
# self.allocations is initialized by the parent, HostOnlyCUDAManager,
# and cleared upon context reset, so although we insert into it here
# and delete from it in the finalizer, we need not do any other
# housekeeping elsewhere.
self.allocations[ptr.value] = buf
return MemoryPointer(ctx, ptr, size, finalizer=finalizer)
def device_array(shape, dtype=np.float, strides=None, order="C", stream=0):
"""
device_array(shape, dtype=np.float, strides=None, order='C',
stream=0)
Allocate an empty Numba device array. Clone of Numba `cuda.device_array`,
but uses RMM for device memory management.
"""
shape, strides, dtype = cuda.api._prepare_shape_strides_dtype(
shape, strides, dtype, order)
datasize = cuda.driver.memory_size_from_info(shape, strides,
dtype.itemsize)
buf = librmm.DeviceBuffer(size=datasize, stream=stream)
ctx = cuda.current_context()
ptr = ctypes.c_uint64(int(buf.ptr))
mem = MemoryPointer(ctx, ptr, datasize, owner=buf)
return cuda.cudadrv.devicearray.DeviceNDArray(shape,
strides,
dtype,
gpu_data=mem)
def memalloc(self, size):
"""
Allocate an on-device array from the RMM pool.
"""
buf = librmm.DeviceBuffer(size=size)
ctx = self.context
if config.CUDA_USE_NVIDIA_BINDING:
ptr = CUdeviceptr(int(buf.ptr))
else:
# expect ctypes bindings in numba
ptr = ctypes.c_uint64(int(buf.ptr))
finalizer = _make_emm_plugin_finalizer(int(buf.ptr), self.allocations)
# self.allocations is initialized by the parent, HostOnlyCUDAManager,
# and cleared upon context reset, so although we insert into it here
# and delete from it in the finalizer, we need not do any other
# housekeeping elsewhere.
self.allocations[int(buf.ptr)] = buf
return MemoryPointer(ctx, ptr, size, finalizer=finalizer)