Coverage for gpaw/test/parallel/test_blacsdist.py: 100%

1import numpy as np 

2 

3from gpaw.mpi import world 

4from gpaw.blacs import BlacsGrid, Redistributor 

5 

6 

7def check(comm, M, N, mcpus, ncpus, mb, nb): 

8 grid0 = BlacsGrid(comm, 1, 1) 

9 desc0 = grid0.new_descriptor(M, N, M, N, 0, 0) 

10 A_mn = desc0.zeros(dtype=float) 

11 A_mn[:] = comm.size + 1 

12 

13 grid1 = BlacsGrid(comm, mcpus, ncpus) 

14 desc1 = grid1.new_descriptor(M, N, mb, nb, 0, 0) # ??? 

15 B_mn = desc1.zeros(dtype=float) 

16 B_mn[:] = comm.rank 

17 

18 if comm.rank == 0: 

19 msg = 'Slices of global matrix indices by rank' 

20 print(msg) 

21 print('-' * len(msg)) 

22 

23 for rank in range(comm.size): 

24 comm.barrier() 

25 if rank == comm.rank: 

26 print('Rank %d:' % rank) 

27 last_Mstart = -1 

28 for Mstart, Mstop, Nstart, Nstop, block in desc1.my_blocks(B_mn): 

29 if Mstart > last_Mstart and last_Mstart >= 0: 

30 print() 

31 print('[%3d:%3d, %3d:%3d]' % (Mstart, Mstop, Nstart, Nstop), 

32 end=' ') 

33 last_Mstart = Mstart 

34 assert (block == comm.rank).all() 

35 print() 

36 print() 

37 comm.barrier() 

38 

39 redistributor = Redistributor(comm, desc1, desc0) 

40 redistributor.redistribute(B_mn, A_mn) 

41 

42 if comm.rank == 0: 

43 msg = 'Rank where each element of the global matrix is stored' 

44 print(msg) 

45 print('-' * len(msg)) 

46 print(A_mn) 

47 

48 

49def test_parallel_blacsdist(scalapack): 

50 M, N = 10, 10 

51 mb, nb = 2, 2 

52 mcpus = int(np.ceil(world.size**0.5)) 

53 ncpus = world.size // mcpus 

54 

55 if world.rank == 0: 

56 print('world size: ', world.size) 

57 print('M x N: ', M, 'x', N) 

58 print('mcpus x ncpus:', mcpus, 'x', ncpus) 

59 print('mb x nb: ', mb, 'x', nb) 

60 print() 

61 

62 check(world, M, N, mcpus, ncpus, mb, nb)