Coverage for pySDC/tutorial/step_7/G_pySDC_on_GPU.py: 100%

1from pathlib import Path 

2 

3from mpi4py import MPI 

4 

5from pySDC.helpers.stats_helper import get_sorted 

6from pySDC.implementations.controller_classes.controller_MPI import controller_MPI 

7from pySDC.implementations.controller_classes.controller_nonMPI import controller_nonMPI 

8from pySDC.implementations.problem_classes.HeatEquation_ND_FD import heatNd_unforced 

9from pySDC.implementations.sweeper_classes.generic_implicit import generic_implicit 

10from pySDC.implementations.transfer_classes.TransferMesh import mesh_to_mesh 

11 

12 

13def get_description(useGPU, ml): 

14 """ 

15 Set up the heat equation, on a GPU or not, on one space level or two. 

16 

17 The only thing that differs between the CPU and the GPU version is `useGPU`. The problem class 

18 is the same one either way: `setup_GPU` swaps the array library, the sparse library and the 

19 datatypes, so the body of the class goes on calling `self.xp.sin` and does not care. 

20 

21 Args: 

22 useGPU (bool): Run on a GPU 

23 ml (bool): Use two space levels rather than one 

24 

25 Returns: 

26 dict: description of the problem to be solved 

27 """ 

28 level_params = {'restol': 1e-10, 'dt': 1e-2} 

29 

30 sweeper_params = {'quad_type': 'RADAU-RIGHT', 'num_nodes': 3, 'QI': 'LU'} 

31 

32 problem_params = { 

33 'nu': 0.1, 

34 'freq': 2, 

35 'bc': 'periodic', 

36 'nvars': [128, 64] if ml else 128, 

37 'useGPU': useGPU, 

38 } 

39 

40 step_params = {'maxiter': 20} 

41 

42 description = { 

43 'problem_class': heatNd_unforced, 

44 'problem_params': problem_params, 

45 'sweeper_class': generic_implicit, 

46 'sweeper_params': sweeper_params, 

47 'level_params': level_params, 

48 'step_params': step_params, 

49 } 

50 

51 if ml: 

52 # the space transfer works on GPU arrays as well: the interpolation and restriction 

53 # matrices are assembled with SciPy and then moved to the device once 

54 description['space_transfer_class'] = mesh_to_mesh 

55 description['space_transfer_params'] = {'rorder': 2, 'iorder': 4, 'periodic': True} 

56 

57 return description 

58 

59 

60def run(description, comm=None, num_procs=1, Tend=8e-2): 

61 """ 

62 Run to `Tend` and report how it went. 

63 

64 Args: 

65 description (dict): description of the problem to be solved 

66 comm (mpi4py.Intracomm): time communicator, for the parallel-in-time run 

67 num_procs (int): number of time steps to treat in parallel, for the serial controller 

68 Tend (float): time to run to 

69 

70 Returns: 

71 float: error against the exact solution 

72 int: total number of iterations 

73 """ 

74 controller_params = {'logger_level': 30} 

75 

76 if comm is None: 

77 controller = controller_nonMPI( 

78 num_procs=num_procs, controller_params=controller_params, description=description 

79 ) 

80 prob = controller.MS[0].levels[0].prob 

81 else: 

82 # the parallel-in-time controller sends the solution from one time rank to the next as a 

83 # GPU array, which needs MPI to have been told to expect device pointers -- see the README 

84 controller = controller_MPI(controller_params=controller_params, description=description, comm=comm) 

85 prob = controller.S.levels[0].prob 

86 

87 uinit = prob.u_exact(0.0) 

88 uend, stats = controller.run(u0=uinit, t0=0.0, Tend=Tend) 

89 

90 error = abs(prob.u_exact(Tend) - uend) 

91 iterations = sum(count for _, count in get_sorted(stats, type='niter', comm=comm)) 

92 

93 return error, iterations 

94 

95 

96def main(): 

97 """ 

98 Solve the same heat equation with SDC, MLSDC and PFASST, all of it on GPUs. 

99 

100 The three differ only in what they are given: one space level and one time step at a time is 

101 SDC, two space levels is MLSDC, and two space levels spread over several time ranks is PFASST. 

102 """ 

103 comm = MPI.COMM_WORLD 

104 

105 # every rank runs the two serial variants -- they are the reference the parallel one is judged 

106 # against, and running them everywhere keeps the ranks in step 

107 runs = [ 

108 ('SDC ', *run(get_description(useGPU=True, ml=False))), 

109 ('MLSDC ', *run(get_description(useGPU=True, ml=True))), 

110 ('PFASST', *run(get_description(useGPU=True, ml=True), comm=comm)), 

111 ] 

112 

113 if comm.rank == 0: 

114 Path('data').mkdir(parents=True, exist_ok=True) 

115 with open('data/step_7_G_out.txt', 'a') as f: 

116 for name, error, iterations in runs: 

117 out = f'{name} on {comm.size} GPU(s): error {error:.4e}, {iterations} iterations in total' 

118 f.write(out + '\n') 

119 print(out) 

120 

121 # all three solve the same problem, so they had better agree on the answer 

122 errors = [error for _, error, _ in runs] 

123 assert max(errors) < 1e-8, f'Some run was not accurate enough: {errors}' 

124 assert abs(errors[1] - errors[2]) < 1e-10, 'PFASST and MLSDC disagree, which they should not' 

125 

126 

127if __name__ == '__main__': 

128 main()