def test_reduce(self): cpu, gpu = gen_gpuarray((3, 4, 5), order='c', incr=self.rank, ctx=self.ctx) rescpu = np.empty_like(cpu) resgpu = gpu._empty_like_me() if self.rank != 0: self.gpucomm.reduce(gpu, 'sum', resgpu, root=0) self.mpicomm.Reduce([cpu, MPI.FLOAT], None, op=MPI.SUM, root=0) else: self.gpucomm.reduce(gpu, 'sum', resgpu) self.mpicomm.Reduce([cpu, MPI.FLOAT], [rescpu, MPI.FLOAT], op=MPI.SUM, root=0) if self.rank == 0: assert np.allclose(resgpu, rescpu) resgpu = self.gpucomm.reduce(gpu, 'sum', root=0) if self.rank == 0: assert resgpu.shape == gpu.shape, (resgpu.shape, gpu.shape) assert resgpu.dtype == gpu.dtype, (resgpu.dtype, gpu.dtype) assert resgpu.flags['C'] == gpu.flags['C'] assert resgpu.flags['F'] == gpu.flags['F'] assert np.allclose(resgpu, rescpu) else: assert resgpu is None if self.rank == 0: resgpu = self.gpucomm.reduce(gpu, 'sum') assert resgpu.shape == gpu.shape, (resgpu.shape, gpu.shape) assert resgpu.dtype == gpu.dtype, (resgpu.dtype, gpu.dtype) assert resgpu.flags['C'] == gpu.flags['C'] assert resgpu.flags['F'] == gpu.flags['F'] assert np.allclose(resgpu, rescpu) else: resgpu = self.gpucomm.reduce(gpu, 'sum', root=0) assert resgpu is None
def test_broadcast(self): if self.rank == 0: cpu, gpu = gen_gpuarray((3, 4, 5), order='c', incr=self.rank, ctx=self.ctx) else: cpu = np.zeros((3, 4, 5), dtype='float32') gpu = gpuarray.asarray(cpu, context=self.ctx) if self.rank == 0: self.gpucomm.broadcast(gpu) else: self.gpucomm.broadcast(gpu, root=0) self.mpicomm.Bcast(cpu, root=0) assert np.allclose(gpu, cpu)
def test_all_reduce(self): cpu, gpu = gen_gpuarray((3, 4, 5), order='c', incr=self.rank, ctx=self.ctx) rescpu = np.empty_like(cpu) resgpu = gpu._empty_like_me() self.gpucomm.all_reduce(gpu, 'sum', resgpu) self.mpicomm.Allreduce([cpu, MPI.FLOAT], [rescpu, MPI.FLOAT], op=MPI.SUM) assert np.allclose(resgpu, rescpu) resgpu = self.gpucomm.all_reduce(gpu, 'sum') assert resgpu.shape == gpu.shape, (resgpu.shape, gpu.shape) assert resgpu.dtype == gpu.dtype, (resgpu.dtype, gpu.dtype) assert resgpu.flags['C'] == gpu.flags['C'] assert resgpu.flags['F'] == gpu.flags['F'] assert np.allclose(resgpu, rescpu)