Coverage for pySDC/implementations/hooks/log_timings.py: 95%

108 statements  

« prev     ^ index     » next       coverage.py v7.6.9, created at 2024-12-20 14:51 +0000

1import time 

2from pySDC.core.hooks import Hooks 

3 

4try: 

5 import cupy as cp 

6except ImportError: 

7 cp = None 

8 

9 

10class Timings(Hooks): 

11 """ 

12 Abstract base class for recoding timings 

13 

14 Attributes: 

15 __t0_setup (float): private variable to get starting time of setup 

16 __t0_run (float): private variable to get starting time of the run 

17 __t0_predict (float): private variable to get starting time of the predictor 

18 __t0_step (float): private variable to get starting time of the step 

19 __t0_iteration (float): private variable to get starting time of the iteration 

20 __t0_sweep (float): private variable to get starting time of the sweep 

21 __t0_comm (list): private variable to get starting time of the communication 

22 __t1_run (float): private variable to get end time of the run 

23 __t1_predict (float): private variable to get end time of the predictor 

24 __t1_step (float): private variable to get end time of the step 

25 __t1_iteration (float): private variable to get end time of the iteration 

26 __t1_sweep (float): private variable to get end time of the sweep 

27 __t1_setup (float): private variable to get end time of setup 

28 __t1_comm (list): private variable to hold timing of the communication (!) 

29 """ 

30 

31 prefix = '' 

32 

33 def _compute_time_elapsed(self, event_after, event_before): 

34 raise NotImplementedError 

35 

36 def _get_event(self): 

37 raise NotImplementedError 

38 

39 def __init__(self): 

40 super().__init__() 

41 self.__t0_setup = None 

42 self.__t0_run = None 

43 self.__t0_predict = None 

44 self.__t0_step = None 

45 self.__t0_iteration = None 

46 self.__t0_sweep = None 

47 self.__t0_comm = [] 

48 self.__t1_run = None 

49 self.__t1_predict = None 

50 self.__t1_step = None 

51 self.__t1_iteration = None 

52 self.__t1_sweep = None 

53 self.__t1_setup = None 

54 self.__t1_comm = [] 

55 

56 def pre_setup(self, step, level_number): 

57 """ 

58 Default routine called before setup starts 

59 

60 Args: 

61 step (pySDC.Step.step): the current step 

62 level_number (int): the current level number 

63 """ 

64 super().pre_setup(step, level_number) 

65 self.__t0_setup = self._get_event() 

66 

67 def pre_run(self, step, level_number): 

68 """ 

69 Default routine called before time-loop starts 

70 

71 Args: 

72 step (pySDC.Step.step): the current step 

73 level_number (int): the current level number 

74 """ 

75 super().pre_run(step, level_number) 

76 self.__t0_run = self._get_event() 

77 

78 def pre_predict(self, step, level_number): 

79 """ 

80 Default routine called before predictor starts 

81 

82 Args: 

83 step (pySDC.Step.step): the current step 

84 level_number (int): the current level number 

85 """ 

86 super().pre_predict(step, level_number) 

87 self.__t0_predict = self._get_event() 

88 

89 def pre_step(self, step, level_number): 

90 """ 

91 Hook called before each step 

92 

93 Args: 

94 step (pySDC.Step.step): the current step 

95 level_number (int): the current level number 

96 """ 

97 super().pre_step(step, level_number) 

98 self.__t0_step = self._get_event() 

99 

100 def pre_iteration(self, step, level_number): 

101 """ 

102 Default routine called before iteration starts 

103 

104 Args: 

105 step (pySDC.Step.step): the current step 

106 level_number (int): the current level number 

107 """ 

108 super().pre_iteration(step, level_number) 

109 self.__t0_iteration = self._get_event() 

110 

111 def pre_sweep(self, step, level_number): 

112 """ 

113 Default routine called before sweep starts 

114 

115 Args: 

116 step (pySDC.Step.step): the current step 

117 level_number (int): the current level number 

118 """ 

119 super().pre_sweep(step, level_number) 

120 self.__t0_sweep = self._get_event() 

121 

122 def pre_comm(self, step, level_number): 

123 """ 

124 Default routine called before communication starts 

125 

126 Args: 

127 step (pySDC.Step.step): the current step 

128 level_number (int): the current level number 

129 """ 

130 super().pre_comm(step, level_number) 

131 if len(self.__t0_comm) >= level_number + 1: 

132 self.__t0_comm[level_number] = self._get_event() 

133 else: 

134 while len(self.__t0_comm) < level_number: 

135 self.__t0_comm.append(None) 

136 self.__t0_comm.append(self._get_event()) 

137 while len(self.__t1_comm) <= level_number: 

138 self.__t1_comm.append(0.0) 

139 assert len(self.__t0_comm) == level_number + 1 

140 assert len(self.__t1_comm) == level_number + 1 

141 

142 def post_comm(self, step, level_number, add_to_stats=False): 

143 """ 

144 Default routine called after each communication 

145 

146 Args: 

147 step (pySDC.Step.step): the current step 

148 level_number (int): the current level number 

149 add_to_stats (bool): set if result should go to stats object 

150 """ 

151 super().post_comm(step, level_number) 

152 assert len(self.__t1_comm) >= level_number + 1 

153 self.__t1_comm[level_number] += self._compute_time_elapsed(self._get_event(), self.__t0_comm[level_number]) 

154 

155 if add_to_stats: 

156 L = step.levels[level_number] 

157 

158 self.add_to_stats( 

159 process=step.status.slot, 

160 process_sweeper=L.sweep.rank, 

161 time=L.time, 

162 level=L.level_index, 

163 iter=step.status.iter, 

164 sweep=L.status.sweep, 

165 type=f'{self.prefix}timing_comm', 

166 value=self.__t1_comm[level_number], 

167 ) 

168 self.__t1_comm[level_number] = 0.0 

169 

170 def post_sweep(self, step, level_number): 

171 """ 

172 Default routine called after each sweep 

173 

174 Args: 

175 step (pySDC.Step.step): the current step 

176 level_number (int): the current level number 

177 """ 

178 super().post_sweep(step, level_number) 

179 self.__t1_sweep = self._get_event() 

180 

181 L = step.levels[level_number] 

182 

183 self.add_to_stats( 

184 process=step.status.slot, 

185 process_sweeper=L.sweep.rank, 

186 time=L.time, 

187 level=L.level_index, 

188 iter=step.status.iter, 

189 sweep=L.status.sweep, 

190 type=f'{self.prefix}timing_sweep', 

191 value=self._compute_time_elapsed(self.__t1_sweep, self.__t0_sweep), 

192 ) 

193 

194 def post_iteration(self, step, level_number): 

195 """ 

196 Default routine called after each iteration 

197 

198 Args: 

199 step (pySDC.Step.step): the current step 

200 level_number (int): the current level number 

201 """ 

202 super().post_iteration(step, level_number) 

203 self.__t1_iteration = self._get_event() 

204 

205 L = step.levels[level_number] 

206 

207 self.add_to_stats( 

208 process=step.status.slot, 

209 process_sweeper=L.sweep.rank, 

210 time=L.time, 

211 level=L.level_index, 

212 iter=step.status.iter, 

213 sweep=L.status.sweep, 

214 type=f'{self.prefix}timing_iteration', 

215 value=self._compute_time_elapsed(self.__t1_iteration, self.__t0_iteration), 

216 ) 

217 

218 def post_step(self, step, level_number): 

219 """ 

220 Default routine called after each step or block 

221 

222 Args: 

223 step (pySDC.Step.step): the current step 

224 level_number (int): the current level number 

225 """ 

226 super().post_step(step, level_number) 

227 self.__t1_step = self._get_event() 

228 

229 L = step.levels[level_number] 

230 

231 self.add_to_stats( 

232 process=step.status.slot, 

233 process_sweeper=L.sweep.rank, 

234 time=L.time, 

235 level=L.level_index, 

236 iter=step.status.iter, 

237 sweep=L.status.sweep, 

238 type=f'{self.prefix}timing_step', 

239 value=self._compute_time_elapsed(self.__t1_step, self.__t0_step), 

240 ) 

241 

242 def post_predict(self, step, level_number): 

243 """ 

244 Default routine called after each predictor 

245 

246 Args: 

247 step (pySDC.Step.step): the current step 

248 level_number (int): the current level number 

249 """ 

250 super().post_predict(step, level_number) 

251 self.__t1_predict = self._get_event() 

252 

253 L = step.levels[level_number] 

254 

255 self.add_to_stats( 

256 process=step.status.slot, 

257 process_sweeper=L.sweep.rank, 

258 time=L.time, 

259 level=L.level_index, 

260 iter=step.status.iter, 

261 sweep=L.status.sweep, 

262 type=f'{self.prefix}timing_predictor', 

263 value=self._compute_time_elapsed(self.__t1_predict, self.__t0_predict), 

264 ) 

265 

266 def post_run(self, step, level_number): 

267 """ 

268 Default routine called after each run 

269 

270 Args: 

271 step (pySDC.Step.step): the current step 

272 level_number (int): the current level number 

273 """ 

274 super().post_run(step, level_number) 

275 self.__t1_run = self._get_event() 

276 t_run = self._compute_time_elapsed(self.__t1_run, self.__t0_run) 

277 

278 L = step.levels[level_number] 

279 

280 self.add_to_stats( 

281 process=step.status.slot, 

282 process_sweeper=L.sweep.rank, 

283 time=L.time, 

284 level=L.level_index, 

285 iter=step.status.iter, 

286 sweep=L.status.sweep, 

287 type=f'{self.prefix}timing_run', 

288 value=t_run, 

289 ) 

290 self.logger.info(f'Finished run after {t_run:.2e}s') 

291 

292 def post_setup(self, step, level_number): 

293 """ 

294 Default routine called after setup 

295 

296 Args: 

297 step (pySDC.Step.step): the current step 

298 level_number (int): the current level number 

299 """ 

300 super().post_setup(step, level_number) 

301 self.__t1_setup = self._get_event() 

302 

303 self.add_to_stats( 

304 process=-1, 

305 process_sweeper=-1, 

306 time=-1, 

307 level=-1, 

308 iter=-1, 

309 sweep=-1, 

310 type=f'{self.prefix}timing_setup', 

311 value=self._compute_time_elapsed(self.__t1_setup, self.__t0_setup), 

312 ) 

313 

314 

315class CPUTimings(Timings): 

316 """ 

317 Hook for recording CPU timings of important operations during a pySDC run. 

318 """ 

319 

320 def _compute_time_elapsed(self, event_after, event_before): 

321 return event_after - event_before 

322 

323 def _get_event(self): 

324 return time.perf_counter() 

325 

326 

327class GPUTimings(Timings): 

328 """ 

329 Hook for recording GPU timings of important operations during a pySDC run. 

330 """ 

331 

332 prefix = 'GPU_' 

333 

334 def _compute_time_elapsed(self, event_after, event_before): 

335 event_after.synchronize() 

336 return cp.cuda.get_elapsed_time(event_before, event_after) / 1e3 

337 

338 def _get_event(self): 

339 event = cp.cuda.Event() 

340 event.record() 

341 return event