Coverage for pySDC/helpers/blocks.py: 87%

1import numpy as np 

2 

3 

4class BlockDecomposition(object): 

5 """ 

6 Class decomposing a cartesian space domain (1D to 3D) into a given number of processors. 

7 

8 Parameters 

9 ---------- 

10 nProcs : int 

11 Total number of processors for space block decomposition. 

12 gridSizes : list[int] 

13 Number of grid points in each dimension 

14 algo : str, optional 

15 Algorithm used for the block decomposition : 

16 

17 - Hybrid : approach minimizing interface communication, inspired from 

18 the `[Hybrid CFD solver] <https://web.stanford.edu/group/ctr/ResBriefs07/5_larsson1_pp47_58.pdf>`_. 

19 - ChatGPT : quickly generated using `[ChatGPT] <https://chatgpt.com>`_. 

20 

21 The default is "Hybrid". 

22 gRank : int, optional 

23 If provided, the global rank that will determine the local block distribution. Default is None. 

24 order : str, optional 

25 The order used when computing the rank block distribution. Default is `C`. 

26 """ 

27 

28 def __init__(self, nProcs, gridSizes, algo="Hybrid", gRank=None, order="C"): 

29 """Compute the number of blocks in each dimension, see the class docstring.""" 

30 dim = len(gridSizes) 

31 assert dim in [1, 2, 3], "block decomposition only works for 1D, 2D or 3D domains" 

32 

33 if algo == "ChatGPT": 

34 

35 nBlocks = [1] * dim 

36 for i in range(2, int(nProcs**0.5) + 1): 

37 while nProcs % i == 0: 

38 nBlocks[0] *= i 

39 nProcs //= i 

40 nBlocks.sort() 

41 

42 if nProcs > 1: 

43 nBlocks[0] *= nProcs 

44 

45 nBlocks.sort() 

46 while len(nBlocks) < dim: 

47 smallest = nBlocks.pop(0) 

48 nBlocks += [1, smallest] 

49 nBlocks.sort() 

50 

51 while len(nBlocks) > dim: 

52 smallest = nBlocks.pop(0) 

53 next_smallest = nBlocks.pop(0) 

54 nBlocks.append(smallest * next_smallest) 

55 nBlocks.sort() 

56 

57 elif algo == "Hybrid": 

58 rest = nProcs 

59 facs = { 

60 1: [1], 

61 2: [2, 1], 

62 3: [2, 3, 1], 

63 }[dim] 

64 exps = [0] * dim 

65 for n in range(dim - 1): 

66 while (rest % facs[n]) == 0: 

67 exps[n] = exps[n] + 1 

68 rest = rest // facs[n] 

69 if rest > 1: 

70 facs[dim - 1] = rest 

71 exps[dim - 1] = 1 

72 

73 nBlocks = [1] * dim 

74 for n in range(dim - 1, -1, -1): 

75 while exps[n] > 0: 

76 dummymax = -1 

77 dmax = 0 

78 for d, nPts in enumerate(gridSizes): 

79 dummy = (nPts + nBlocks[d] - 1) // nBlocks[d] 

80 if dummy >= dummymax: 

81 dummymax = dummy 

82 dmax = d 

83 nBlocks[dmax] = nBlocks[dmax] * facs[n] 

84 exps[n] = exps[n] - 1 

85 

86 else: 

87 raise NotImplementedError(f"algo={algo}") 

88 

89 # Store attributes 

90 self.dim = dim 

91 self.nBlocks = nBlocks 

92 self.gridSizes = gridSizes 

93 

94 # Used for rank block distribution 

95 self.gRank = gRank 

96 self.order = order 

97 

98 @property 

99 def ranks(self): 

100 """ 

101 Coordinates of the block of `gRank` in the cartesian grid of blocks, which uses `order`. 

102 

103 Returns 

104 ------- 

105 list[int] 

106 Block index of `gRank` in each dimension. 

107 """ 

108 assert self.gRank is not None, "gRank attribute needs to be set" 

109 cart = np.arange(np.prod(self.nBlocks)).reshape(self.nBlocks, order=self.order) 

110 return list(np.argwhere(cart == self.gRank)[0]) 

111 

112 @property 

113 def localBounds(self): 

114 """ 

115 Bounds of the block of `gRank`. The points left over by an even split go one each to the first blocks. 

116 

117 Returns 

118 ------- 

119 iLocList : list[int] 

120 Global index of the first local point in each dimension. 

121 nLocList : list[int] 

122 Number of local points in each dimension. 

123 """ 

124 iLocList, nLocList = [], [] 

125 for rank, nPoints, nBlocks in zip(self.ranks, self.gridSizes, self.nBlocks, strict=True): 

126 n0 = nPoints // nBlocks 

127 nRest = nPoints - nBlocks * n0 

128 nLoc = n0 + 1 * (rank < nRest) 

129 iLoc = rank * n0 + nRest * (rank >= nRest) + rank * (rank < nRest) 

130 

131 iLocList.append(iLoc) 

132 nLocList.append(nLoc) 

133 return iLocList, nLocList 

134 

135 

136if __name__ == "__main__": 

137 # Base usage of this module for a 2D decomposition 

138 from mpi4py import MPI 

139 from time import sleep 

140 

141 comm: MPI.Intracomm = MPI.COMM_WORLD 

142 MPI_SIZE = comm.Get_size() 

143 MPI_RANK = comm.Get_rank() 

144 

145 blocks = BlockDecomposition(MPI_SIZE, [256, 64], gRank=MPI_RANK) 

146 if MPI_RANK == 0: 

147 print(f"nBlocks : {blocks.nBlocks}") 

148 

149 ranks = blocks.ranks 

150 bounds = blocks.localBounds 

151 

152 comm.Barrier() 

153 sleep(0.01 * MPI_RANK) 

154 print(f"[Rank {MPI_RANK}] pRankX={ranks}, bounds={bounds}")