def get_ipc_handle(ary, stream=0):
"""
Get an IPC handle from the DeviceArray ary with offset modified by
the RMM memory pool.
"""
ipch = cuda.devices.get_context().get_ipc_handle(ary.gpu_data)
ptr = ary.device_ctypes_pointer.value
offset = librmm.rmm_getallocationoffset(ptr, stream)
# replace offset with RMM's offset
ipch.offset = offset
desc = dict(shape=ary.shape, strides=ary.strides, dtype=ary.dtype)
return cuda.cudadrv.devicearray.IpcArrayHandle(ipc_handle=ipch,
array_desc=desc)
def get_ipc_handle(self, memory):
"""
Get an IPC handle for the MemoryPointer memory with offset modified by
the RMM memory pool.
"""
ipchandle = (ctypes.c_byte * 64)() # IPC handle is 64 bytes
cuda.cudadrv.driver.driver.cuIpcGetMemHandle(
ctypes.byref(ipchandle),
memory.owner.handle,
)
source_info = cuda.current_context().device.get_device_identity()
ptr = memory.device_ctypes_pointer.value
offset = librmm.rmm_getallocationoffset(ptr, 0)
return IpcHandle(memory,
ipchandle,
memory.size,
source_info,
offset=offset)