def testMessageCuPy(self):
sbuf = cupy.array([1,2,3], 'i')
rbuf = cupy.array([0,0,0], 'i')
smsg = [sbuf, [3], [0], [MPI.INT]]
rmsg = [rbuf, ([3], [0]), [MPI.INT]]
Alltoallw(smsg, rmsg)
self.assertTrue((sbuf == rbuf).all())
def __init__(self, typecode, initializer):
self._buf = cupy.array(initializer, dtype=typecode)
self.has_dlpack = hasattr(self._buf, '__dlpack_device__')
# TODO(leofang): test CUDA managed memory?
if cupy.cuda.runtime.is_hip:
self.dev_type = dlpack.DLDeviceType.kDLROCM
else:
self.dev_type = dlpack.DLDeviceType.kDLCUDA
def array(self, typecode, initializer):
return cupy.array(initializer, dtype=typecode)
def testMessageCuPy(self):
sbuf = cupy.array([1,2,3], 'i')
rbuf = cupy.array([0,0,0], 'i')
PutGet(sbuf, rbuf)
self.assertTrue((sbuf == rbuf).all())