Coverage for gpaw/test/gpu/test_mpi.py: 100%

14 statements  

« prev     ^ index     » next       coverage.py v7.7.1, created at 2025-07-08 00:17 +0000

1import pytest 

2 

3import gpaw.cgpaw as cgpaw 

4from gpaw.gpu import cupy as cp 

5from gpaw.gpu.mpi import CuPyMPI 

6import gpaw.mpi as mpi 

7 

8 

9@pytest.mark.gpu 

10def test_mpi(): 

11 a = cp.ones(1) 

12 world = mpi.world 

13 print(a, world, a.shape, a.nbytes, a.size, a.dtype, a.data.ptr) 

14 if not getattr(cgpaw, 'gpu_aware_mpi', False): 

15 world = CuPyMPI(world) 

16 world.sum(a) 

17 assert a[0].get() == world.size