Coverage for gamdpy/simulation/Simulation.py: 40%

1import numpy as np 

2import numba 

3import math 

4import sys 

5from numba import cuda 

6 

7# gamdpy 

8import gamdpy as gp 

9 

10# For type annotation 

11from gamdpy.integrators import Integrator 

12from gamdpy.interactions import Interaction 

13from gamdpy.runtime_actions import RuntimeAction 

14 

15# TODO: to remove import above you need to add a lot of lines are the following 

16#from ..simulation.get_default_compute_plan import get_default_compute_plan 

17#from ..configuration.Configuration import Configuration 

18 

19# IO 

20import h5py 

21 

22 

23class Simulation(): 

24 """ Class for running a simulation. 

25 

26 Parameters 

27 ---------- 

28 

29 configuration : gamdpy.Configuration 

30 The configuration to simulate. 

31 

32 interactions : list of interactions 

33 Interactions such as pair potentials, bonds, external fields, etc. 

34 

35 integrator : an integrator 

36 The integrator to use for the simulation. 

37 

38 runtime_actions : list of runtime actions 

39 Runtime actions such as ScalarSaver, TrajectorySaver or MomentumReset 

40 

41 num_timeblocks : int 

42 Number of timeblocks to run the simulation. If not 0, then steps_per_timeblock should be set. 

43 

44 steps_per_timeblock : int 

45 Number of steps per timeblock. 

46 

47 compute_plan : dict 

48 A dictionary with the compute plan for the simulation. If None, a default compute plan is used. 

49 

50 storage : str 

51 Storage for the simulation output. Can be 'memory' or a filename with extension '.h5'. 

52 

53 verbose : bool 

54 If True, print verbose output. 

55 

56 timing : bool 

57 If True, timing information is saved. 

58 

59 

60 See also 

61 -------- 

62 

63 :func:`gamdpy.get_default_sim` 

64 

65 """ 

66 

67 def __init__(self, configuration: gp.Configuration, 

68 interactions: Interaction|list[Interaction], 

69 integrator: Integrator, 

70 runtime_actions: list[RuntimeAction], 

71 num_timeblocks, steps_per_timeblock, 

72 storage: str, 

73 compute_plan=None, 

74 verbose=False, 

75 timing=True, 

76 steps_in_kernel_test=1): 

77 

78 self.configuration = configuration 

79 if compute_plan == None: 

80 self.compute_plan = gp.get_default_compute_plan(self.configuration) 

81 else: 

82 self.compute_plan = compute_plan 

83 

84 # Integrator 

85 if type(interactions) == list: 

86 self.interactions = interactions 

87 else: 

88 self.interactions = [interactions, ] 

89 self.integrator = integrator 

90 self.dt = self.integrator.dt 

91 

92 self.num_blocks = num_timeblocks 

93 self.current_block = -1 

94 

95 self.steps_per_block = steps_per_timeblock 

96 self.storage = storage 

97 self.timing = timing 

98 self.steps_in_kernel_test = steps_in_kernel_test 

99 

100 # Close output object if there 

101 # Check https://stackoverflow.com/questions/610883/how-to-check-if-an-object-has-an-attribute 

102 # Create output objects 

103 if self.storage == 'memory': 

104 # Creates a memory h5 file with named id(self).h5; id(self) is ensured to be unique 

105 self.memory = h5py.File(f"{id(self)}.h5", "w", driver='core', backing_store=False) 

106 elif isinstance(self.storage, str) and self.storage[-3:] == '.h5': 

107 # The append is important for repeated istances of sim with same self.storage 

108 self.memory = h5py.File(self.storage, "w") 

109 else: 

110 raise ValueError(f"storage needs to be either 'memory' or an hdf5 filename ending in '.h5' (got: {storage})") 

111 

112 # Save setup info 

113 self.memory.attrs['dt'] = self.dt 

114 if 'script_name' not in self.memory.keys(): 

115 script_name = sys.argv[0] 

116 self.memory.attrs['script_name'] = script_name 

117 if isinstance(script_name,str) and script_name != '': 

118 with open(script_name, 'r') as file: 

119 script_content = file.read() 

120 self.memory.attrs['script_content'] = script_content 

121 

122 # Saving initial configuration 

123 self.configuration.save(output=self.memory, group_name="initial_configuration", mode="w", include_topology=True) 

124 

125 self.runtime_actions = runtime_actions 

126 

127 compute_flags = None 

128 for runtime_action in self.runtime_actions: 

129 if runtime_action.get_compute_flags() is not None: 

130 if compute_flags is not None: 

131 raise ValueError('Can not handle more than one compute_flags in runtime_actions') 

132 else: 

133 compute_flags = runtime_action.get_compute_flags() 

134 

135 self.compute_flags = gp.get_default_compute_flags() # configuration.compute_flags 

136 if compute_flags is not None: 

137 # only keys present in the default are processed 

138 for k in compute_flags: 

139 if k in self.compute_flags: 

140 self.compute_flags[k] = compute_flags[k] 

141 else: 

142 raise ValueError('Unknown key in compute_flags:%s' %k) 

143 for k in self.compute_flags: 

144 if self.compute_flags[k] and not configuration.compute_flags[k]: 

145 raise ValueError('compute_flags["%s]" set for Simulation but not in Configuration' % k) 

146 

147 for runtime_action in self.runtime_actions: 

148 runtime_action.setup(configuration=self.configuration, num_timeblocks=num_timeblocks, 

149 steps_per_timeblock=steps_per_timeblock, output=self.memory ) 

150 

151 self.vectors_list = [] 

152 self.scalars_list = [] 

153 self.simbox_data_list = [] 

154 

155 self.JIT_and_test_kernel() 

156 for interaction in self.interactions: # Attempt to catch (eg.) nblist errors before actually doing any simulation 

157 interaction.check_datastructure_validity() 

158 

159 

160 if self.storage and self.storage[-3:] == '.h5': 

161 self.memory.close() 

162 

163 # __del__ is supposed to work also if __init__ fails. This means you can't use attributed defined in __init__ 

164 # https://www.algorithm.co.il/programming/python-gotchas-1-__del__-is-not-the-opposite-of-__init__/ 

165 

166 def get_output(self, mode="r"): 

167 if self.storage and self.storage[-3:] == '.h5': 

168 output = h5py.File(self.storage, mode) 

169 elif self.storage == 'memory': 

170 output = self.memory 

171 else: 

172 #print("Warning: self.output can't recognize self.storage option, returning None") 

173 output = None 

174 return output 

175 

176 # might be worth looking into cache_property https://www.reddit.com/r/learnpython/comments/184kqzp/in_class_properties_defined_in_functions/ 

177 @property 

178 def output(self): 

179 return self.get_output() 

180 

181 def JIT_and_test_kernel(self, adjust_compute_plan=True): 

182 while True: 

183 try: 

184 self.get_kernels_and_params() 

185 self.integrate = self.make_integrator(self.configuration, self.integrator_kernel, self.interactions_kernel, 

186 #self.output_calculator_kernel, 

187 self.runtime_actions_prestep_kernel, 

188 self.runtime_actions_poststep_kernel, 

189 self.compute_plan, True) 

190 self.configuration.copy_to_device() # By _not_ copying back to host later we dont change configuration 

191 self.integrate_self(0.0, self.steps_in_kernel_test) 

192 break 

193 except numba.cuda.cudadrv.driver.CudaAPIError as e: 

194 if not adjust_compute_plan: 

195 self.compute_plan['tp'] = 0 # Signal failure to autotuner 

196 break 

197 #print('Failed compute_plan : ', self.compute_plan) 

198 if self.compute_plan['tp'] > 1: # Most common problem tp is too big 

199 self.compute_plan['tp'] -= 1 # ... so we reduce it and try again 

200 elif self.compute_plan['gridsync'] == True: # Last resort: turn off gridsync 

201 self.compute_plan['gridsync'] = False 

202 else: 

203 print(f'FAILURE. Can not handle cuda error {e}') 

204 exit() 

205 print('Trying adjusted compute_plan :', self.compute_plan) 

206 

207 def get_kernels_and_params(self, verbose=False): 

208 # Interactions 

209 self.interactions_kernel, self.interactions_params = gp.add_interactions_list(self.configuration, 

210 self.interactions, 

211 compute_plan=self.compute_plan, 

212 compute_flags=self.compute_flags) 

213 

214 # Runtime actions 

215 if self.runtime_actions: 

216 self.runtime_actions_prestep_kernel, self.runtime_actions_poststep_kernel, self.runtime_actions_params = gp.add_runtime_actions_list(self.configuration, 

217 self.runtime_actions, 

218 compute_plan=self.compute_plan) 

219 else: 

220 self.runtime_actions_prestep_kernel = None 

221 self.runtime_actions_poststep_kernel = None 

222 self.runtime_actions_params = (0,) 

223 

224 # Integrator 

225 self.integrator_params = self.integrator.get_params(self.configuration, self.interactions_params) 

226 self.integrator_kernel = self.integrator.get_kernel(self.configuration, self.compute_plan, self.compute_flags, self.interactions_kernel) 

227 

228 return 

229 

230 def update_params(self, verbose=False): 

231 # Interactions 

232 _, self.interactions_params = gp.add_interactions_list(self.configuration, 

233 self.interactions, 

234 compute_plan=self.compute_plan, 

235 compute_flags=self.compute_flags) 

236 

237 # Runtime actions 

238 if self.runtime_actions: 

239 _, _, self.runtime_actions_params = gp.add_runtime_actions_list(self.configuration, 

240 self.runtime_actions, 

241 compute_plan=self.compute_plan) 

242 else: 

243 #self.runtime_actions_kernel = None 

244 self.runtime_actions_params = (0,) 

245 

246 # Integrator 

247 self.integrator_params = self.integrator.get_params(self.configuration, self.interactions_params, verbose) 

248 #self.integrator_kernel = self.integrator.get_kernel(self.configuration, self.compute_plan, verbose) 

249 

250 return 

251 

252 def integrate_self(self, time_zero, steps): 

253 self.integrate(self.configuration.d_vectors, 

254 self.configuration.d_scalars, 

255 self.configuration.d_ptype, 

256 self.configuration.d_r_im, 

257 self.configuration.simbox.d_data, 

258 self.interactions_params, 

259 self.integrator_params, 

260 self.runtime_actions_params, 

261 np.float32(time_zero), 

262 steps) 

263 return 

264 

265 

266 def make_integrator(self, configuration, integration_step, compute_interactions,# output_calculator_kernel, 

267 runtime_actions_prestep_kernel, runtime_actions_poststep_kernel, compute_plan, verbose=True): 

268 

269 # Unpack parameters from configuration and compute_plan 

270 D, num_part = configuration.D, configuration.N 

271 pb, tp, gridsync = [compute_plan[key] for key in ['pb', 'tp', 'gridsync']] 

272 num_blocks = (num_part - 1) // pb + 1 

273 

274 get_integrator_setup = getattr(self.integrator, "get_setup_kernel", None) 

275 if get_integrator_setup is not None: 

276 integrator_setup = get_integrator_setup(self.configuration, self.compute_plan, self.interactions_kernel) 

277 self.configuration.copy_to_device() 

278 integrator_setup( 

279 self.configuration.d_vectors, 

280 self.configuration.d_scalars, 

281 self.configuration.d_r_im, 

282 self.configuration.simbox.d_data, 

283 self.integrator_params, 

284 self.configuration.d_ptype, 

285 ) 

286 

287 if gridsync: 

288 # Return a kernel that does 'steps' timesteps, using grid.sync to syncronize  

289 @cuda.jit 

290 def integrator(vectors, scalars, ptype, r_im, sim_box, interaction_params, integrator_params, 

291 runtime_actions_params, time_zero, steps): 

292 grid = cuda.cg.this_grid() 

293 for step in range(steps + 1): # make extra step without integration, so that interactions and run_time actions called for final configuration 

294 time = time_zero + step * integrator_params[0] 

295 compute_interactions(grid, vectors, scalars, ptype, sim_box, interaction_params) 

296 grid.sync() 

297 if runtime_actions_prestep_kernel != None: 

298 runtime_actions_prestep_kernel(grid, vectors, scalars, r_im, sim_box, step, runtime_actions_params) 

299 grid.sync() 

300 if step<steps: 

301 integration_step(grid, vectors, scalars, r_im, sim_box, integrator_params, time, ptype) 

302 grid.sync() 

303 if runtime_actions_poststep_kernel != None: 

304 runtime_actions_poststep_kernel(grid, vectors, scalars, r_im, sim_box, step, runtime_actions_params) 

305 grid.sync() 

306 

307 return 

308 

309 return integrator[num_blocks, (pb, tp)] 

310 

311 else: 

312 

313 # Return a Python function that does 'steps' timesteps, using kernel calls to syncronize  

314 def integrator(vectors, scalars, ptype, r_im, sim_box, interaction_params, integrator_params, 

315 runtime_actions_params, time_zero, steps): 

316 for step in range(steps + 1): # make extra step without integration, so that interactions and run_time actions called for final configuration 

317 time = time_zero + step * integrator_params[0] 

318 compute_interactions(0, vectors, scalars, ptype, sim_box, interaction_params) 

319 if runtime_actions_prestep_kernel != None: 

320 runtime_actions_prestep_kernel(0, vectors, scalars, r_im, sim_box, step, runtime_actions_params) 

321 if step<steps: 

322 integration_step(0, vectors, scalars, r_im, sim_box, integrator_params, time, ptype) 

323 if runtime_actions_poststep_kernel != None: 

324 runtime_actions_poststep_kernel(0, vectors, scalars, r_im, sim_box, step, runtime_actions_params) 

325 return 

326 

327 return integrator 

328 return 

329 

330 # simple run function 

331 

332 def run(self, verbose=True): 

333 """ Run the simulation. 

334 

335 See also 

336 -------- 

337 

338 :func:`gamdpy.Simulation.timeblocks` 

339 

340 """ 

341 for _ in self.run_timeblocks(): 

342 if verbose: 

343 print(self.status(per_particle=True)) 

344 if verbose: 

345 print(self.summary()) 

346 

347 # generator for running simulation one block at a time 

348 def run_timeblocks(self, num_timeblocks=-1): 

349 """ Generator for running the simulation one block at a time. 

350 

351 Parameters 

352 ---------- 

353 

354 num_timeblocks : int 

355 Number of blocks to run. If -1, all blocks are run. 

356 

357 Examples 

358 -------- 

359 

360 >>> import gamdpy as gp 

361 >>> sim = gp.get_default_sim() 

362 >>> for block in sim.run_timeblocks(num_timeblocks=3): 

363 ... print(f'{block=}') # Replace with code to analyze the current configuration 

364 block=0 

365 block=1 

366 block=2 

367 

368 See also 

369 -------- 

370 

371 :func:`gamdpy.Simulation.run` 

372 

373 """ 

374 if num_timeblocks == -1: 

375 num_timeblocks = self.num_blocks 

376 self.last_num_blocks = num_timeblocks 

377 assert (num_timeblocks <= self.num_blocks) # Could be made OK with more blocks 

378 

379 self.configuration.copy_to_device() 

380 self.vectors_list = [] 

381 self.scalars_list = [] 

382 self.simbox_data_list = [] 

383 self.scalars_t = [] 

384 

385 if self.timing: 

386 start = cuda.event() 

387 end = cuda.event() 

388 start_block = cuda.event() 

389 end_block = cuda.event() 

390 block_times = [] 

391 start.record() 

392 

393 zero = np.float32(0.0) 

394 

395 for block in range(num_timeblocks): 

396 

397 self.current_block = block 

398 for runtime_action in self.runtime_actions: 

399 runtime_action.initialize_before_timeblock(block, self.get_output(mode="a")) 

400 

401 if self.timing: 

402 start_block.record() 

403 

404 self.integrate_self(np.float32(block * self.steps_per_block * self.dt),self.steps_per_block) 

405 

406 if self.timing: 

407 end_block.record() 

408 end_block.synchronize() 

409 block_times.append(cuda.event_elapsed_time(start_block, end_block)) 

410 

411 self.configuration.copy_to_host() 

412 

413 for interaction in self.interactions: 

414 interaction.check_datastructure_validity() 

415 

416 for runtime_action in self.runtime_actions: 

417 runtime_action.update_at_end_of_timeblock(block, self.get_output(mode="a")) 

418 

419 #if self.storage: 

420 # self.configuration.save(output=self.get_output(mode="a"), group_name=f"/restarts/restart{block:04d}", mode="w", include_topology=True) 

421 

422 if self.storage and self.storage[-3:] == '.h5': 

423 self.output.close() 

424 

425 yield block 

426 

427 # Finalizing run 

428 if self.timing: 

429 end.record() 

430 end.synchronize() 

431 

432 self.timing_numba = cuda.event_elapsed_time(start, end) 

433 self.timing_numba_blocks = np.array(block_times) 

434 #self.nbflag = self.interactions[0].nblist.d_nbflag.copy_to_host() 

435 self.scalars_list = np.array(self.scalars_list) 

436 

437 def status(self, per_particle=False) -> str: 

438 """ String with the current status 

439 Should be executed during the simulation run, see :func:`gamdpy.Simulation.timeblocks` 

440 

441 Parameters 

442 ---------- 

443 

444 per_particle : bool 

445 If True, the values are divided by the number of particles in the configuration. 

446 

447 Returns 

448 ------- 

449 

450 str 

451 A string with the current status of the simulation. 

452 

453 """ 

454 time = (self.current_block+1) * self.steps_per_block * self.dt 

455 st = f'timeblock= {self.current_block :<6}' 

456 st += f'{time= :<12.3f}' 

457 for name in self.configuration.sid: 

458 if name in self.configuration.compute_flags and self.configuration.compute_flags[name]: 

459 data = np.sum(self.configuration[name]) 

460 if per_particle: 

461 data /= self.configuration.N 

462 st += f'{name}= {data:<10.3f}' 

463 return st 

464 

465 def summary(self) -> str: 

466 """ Returns a summary string of the simulation run. 

467 Should be called after the simulation has been run, 

468 see :func:`gamdpy.Simulation.timeblocks` or :func:`gamdpy.Simulation.run` 

469 """ 

470 if self.timing: 

471 time_total = self.timing_numba / 1000 

472 tps_total = self.last_num_blocks * self.steps_per_block / time_total 

473 time_sim = np.sum(self.timing_numba_blocks) / 1000 

474 tps_sim = self.last_num_blocks * self.steps_per_block / time_sim 

475 

476 st = f'Particles : {self.configuration.N} \n' 

477 st += f'Steps : {self.last_num_blocks} * {self.steps_per_block} = ' 

478 st += f'{self.last_num_blocks * self.steps_per_block:_} \n' 

479 if self.timing: 

480 st += f'Total run time (incl. time spent between timeblocks): {time_total:.2f} s ' 

481 st += f'( TPS: {tps_total:.2e} )\n' 

482 st += f'Simulation time (excl. time spent between timeblocks): {time_sim:.2f} s ' 

483 st += f'( TPS: {tps_sim:.2e} )\n' 

484 return st 

485 

486 def autotune_bruteforce(self, pbs='auto', skins='auto', tps='auto', timesteps=0, repeats=1, verbose=False): 

487 if verbose: 

488 print('compute_plan :', self.compute_plan) 

489 if timesteps==0: 

490 timesteps = self.steps_per_block 

491 assert timesteps<=self.steps_per_block 

492 

493 pb = self.compute_plan['pb'] 

494 if pbs=='auto': 

495 pbs = [pb//2, pb, pb*2] 

496 if pbs=='default': 

497 pbs = [pb,] 

498 if verbose: 

499 print('pbs :', pbs) 

500 

501 tp = self.compute_plan['tp'] 

502 if tps=='auto': 

503 tps = [tp - 3, tp - 2, tp - 1, tp, tp + 1, tp + 2, tp + 3,] 

504 if tps=='default': 

505 tps = [tp,] 

506 if verbose: 

507 print('tps :', tps) 

508 

509 skin = self.compute_plan['skin'] 

510 if skins=='auto': 

511 skins = [skin - 0.6, skin - 0.4, skin - 0.2, skin, skin + 0.2, skin + 0.4, skin + 0.6, skin + 0.8, skin + 1.0] 

512 elif skins=='default': 

513 skins = [skin, ] 

514 if verbose: 

515 print('skins :', skins) 

516 

517 nblists = [] 

518 if self.configuration.N < 32000: 

519 nblists.append('N squared') 

520 if self.configuration.N > 2000: 

521 nblists.append('linked lists') 

522 if verbose: 

523 print('nblists:', nblists) 

524 

525 gridsyncs = [] 

526 if self.configuration.N < 200000: 

527 gridsyncs.append(True) 

528 if self.configuration.N > 10000: 

529 gridsyncs.append(False) 

530 if verbose: 

531 print('gridsyncs :', gridsyncs) 

532 

533 UtilizeNIIIs = [False, True] 

534 if verbose: 

535 print('UtilizeNIIIs :', UtilizeNIIIs) 

536 

537 flag = cuda.config.CUDA_LOW_OCCUPANCY_WARNINGS 

538 cuda.config.CUDA_LOW_OCCUPANCY_WARNINGS = False 

539 

540 skin_times = [] 

541 total_min_time = 1e9 

542 optimal_compute_plan = gp.get_default_compute_plan(self.configuration) # Overwritten below 

543 

544 for nblist in nblists: 

545 self.compute_plan['nblist'] = nblist 

546 for gridsync in gridsyncs: 

547 self.compute_plan['gridsync'] = gridsync 

548 for UtilizeNIII in UtilizeNIIIs: 

549 self.compute_plan['UtilizeNIII'] = UtilizeNIII 

550 print(f'\n {nblist}, {gridsync=}, {UtilizeNIII=}: ', end='') 

551 local_min_time = 1e9 

552 for pb in pbs: 

553 if pb <= 512: 

554 self.compute_plan['pb'] = pb 

555 for tp in tps: 

556 if tp>0: 

557 self.compute_plan['tp'] = tp 

558 gridsync = self.compute_plan['gridsync'] 

559 self.JIT_and_test_kernel(adjust_compute_plan=False) 

560 # does kernel run without adjustment? 

561 if self.compute_plan['tp'] != tp or self.compute_plan['gridsync'] != gridsync: 

562 break 

563 #print('Seems to work, so looping over skins...') 

564 total_min_time, local_min_time = self.autotune_scan_skin(self.compute_plan, skins, timesteps, repeats, total_min_time, local_min_time, optimal_compute_plan, verbose) 

565 

566 self.compute_plan = optimal_compute_plan 

567 if verbose: 

568 print('\nFinal compute_plan :', self.compute_plan) 

569 else: 

570 print('') 

571 #print('\nFinal compute_plan :', self.compute_plan) 

572 self.JIT_and_test_kernel() 

573 

574 cuda.config.CUDA_LOW_OCCUPANCY_WARNINGS = flag 

575 

576 

577 def autotune_scan_skin(self, compute_plan, skins, timesteps, repeats, total_min_time, local_min_time, optimal_compute_plan, verbose=False): 

578 min_time = 1e9 

579 skin_times = [] 

580 pb = compute_plan['pb'] 

581 tp = compute_plan['tp'] 

582 for skin in skins: 

583 if 0 < skin < 1.2: 

584 self.compute_plan['skin'] = skin 

585 self.update_params() 

586 self.configuration.copy_to_device() # By _not_ copying back to host later we dont change configuration 

587 start = cuda.event() 

588 end = cuda.event() 

589 start.record() 

590 for i in range(repeats): 

591 self.integrate_self(0.0, timesteps) 

592 end.record() 

593 end.synchronize() 

594 time_elapsed = cuda.event_elapsed_time(start, end) 

595 if time_elapsed < min_time: 

596 min_time = time_elapsed 

597 min_skin = skin 

598 skin_times.append(time_elapsed) 

599 if verbose: 

600 print(f"({skin}, {skin_times[-1]:.3})", end=' ', flush=True) 

601 max_TPS = repeats * timesteps / min_time * 1000 

602 if verbose: 

603 print('\n', pb, tp, min_skin, min_time, max_TPS) 

604 else: 

605 print('.', end='', flush=True) 

606 if min_time < local_min_time: 

607 local_min_time = min_time 

608 print(f' ({pb}x{tp},{min_skin:.2f}):{max_TPS:.2f}', end=' ', flush=True) 

609 if min_time < total_min_time: 

610 total_min_time = min_time 

611 optimal_compute_plan['UtilizeNIII'] = self.compute_plan['UtilizeNIII'] 

612 optimal_compute_plan['nblist'] = self.compute_plan['nblist'] 

613 optimal_compute_plan['gridsync'] = self.compute_plan['gridsync'] 

614 optimal_compute_plan['skin'] = min_skin 

615 optimal_compute_plan['pb'] = pb 

616 optimal_compute_plan['tp'] = tp 

617 return total_min_time, local_min_time 

618 

619 

620 def autotune(self): 

621 flag = cuda.config.CUDA_LOW_OCCUPANCY_WARNINGS 

622 cuda.config.CUDA_LOW_OCCUPANCY_WARNINGS = False 

623 

624 initial_compute_plan = self.compute_plan.copy() 

625 timesteps = self.steps_per_block 

626 repeats = 1 

627 

628 # Binary choises: Take self.compute_plan as starting point, and add alternatives if appropiate  

629 gridsyncs = [initial_compute_plan['gridsync'], ] 

630 if gridsyncs[0] != True: 

631 gridsyncs.append(True) 

632 if gridsyncs[0] != False and self.configuration.N > 10000: 

633 gridsyncs.append(False) 

634 

635 nblists = [initial_compute_plan['nblist'], ] 

636 if nblists[0] != 'N squared' and self.configuration.N < 32000: 

637 nblists.append('N squared') 

638 if nblists[0] != 'linked lists' and self.configuration.N > 2000: 

639 nblists.append('linked lists') 

640 

641 UtilizeNIIIs = [initial_compute_plan['UtilizeNIII'], ] 

642 if UtilizeNIIIs[0] != False: 

643 UtilizeNIIIs.append(False) 

644 if UtilizeNIIIs[0] != True: 

645 UtilizeNIIIs.append(True) 

646 

647 

648 pb = self.compute_plan['pb'] 

649 pbs = [pb//2, pb, pb*2] 

650 

651 optimal_compute_plan = initial_compute_plan.copy() 

652 results = [] 

653 # Loop over binary parameters 

654 total_min_time = 1e9 

655 binary_min_time = 1e9 

656 for nblist in nblists: 

657 self.compute_plan['nblist'] = nblist 

658 for gridsync in gridsyncs: 

659 self.compute_plan['gridsync'] = gridsync 

660 for UtilizeNIII in UtilizeNIIIs: 

661 self.compute_plan['UtilizeNIII'] = UtilizeNIII 

662 

663 # Get time for default 'pb', 'tp' to check if its worth to proceede with these choises 

664 self.compute_plan['pb'] = initial_compute_plan['pb'] 

665 self.compute_plan['tp'] = initial_compute_plan['tp'] 

666 print(f'\n {nblist+",":12}\t{gridsync=}\t{UtilizeNIII=}:\t', end='') 

667 self.JIT_and_test_kernel(adjust_compute_plan=False) 

668 if self.compute_plan['tp'] == 0 or self.compute_plan['gridsync'] != gridsync: 

669 continue 

670 

671 local_min_time = 1e9 

672 total_min_time, local_min_time, min_time = self.autotune_skin(initial_compute_plan['skin'], 0.2, 

673 timesteps, repeats, total_min_time, local_min_time, 

674 optimal_compute_plan, verbose=False) 

675 self.compute_plan['min_time'] = min_time 

676 if min_time < binary_min_time: 

677 binary_min_time = min_time 

678 results.append(self.compute_plan.copy()) 

679 

680 for compute_plan in results: 

681 if compute_plan['min_time'] < 1.15 * binary_min_time: 

682 

683 local_min_time = 1e9 

684 self.compute_plan = compute_plan.copy() 

685 nblist = compute_plan['nblist'] 

686 gridsync = compute_plan['gridsync'] 

687 UtilizeNIII = compute_plan['UtilizeNIII'] 

688 print(f'\n {nblist+",":12}\t{gridsync=}\t{UtilizeNIII=}:\t', end='') 

689 

690 pb = compute_plan['pb'] 

691 pb_min_time = 1e9 

692 initial_min_time = -1e9 

693 while 4 <= pb <= 1024: 

694 self.compute_plan['pb'] = pb 

695 print(f' {pb=} ', end='') 

696 total_min_time, local_min_time, min_time = self.autotune_tp(compute_plan['tp'], 1, compute_plan, timesteps, repeats, optimal_compute_plan, total_min_time, local_min_time) 

697 if initial_min_time<0: 

698 initial_min_time = min_time 

699 #print(f'[{min_time:.3}, {pb_min_time:.3}]') 

700 if min_time > 1.01 * pb_min_time: 

701 break 

702 if min_time < pb_min_time: 

703 pb_min_time = min_time 

704 pb *= 2 

705 pb = compute_plan['pb'] // 2 

706 if initial_min_time < 1.01*pb_min_time: # Is it worth trying smaller pb? 

707 while 4 <= pb <= 1024: 

708 self.compute_plan['pb'] = pb 

709 print(f' {pb=} ', end='') 

710 total_min_time, local_min_time, min_time = self.autotune_tp(compute_plan['tp'], 1, compute_plan, timesteps, repeats, optimal_compute_plan, total_min_time, local_min_time) 

711 #print(f'[{min_time:.3}, {pb_min_time:.3}]') 

712 if min_time > 1.01 * pb_min_time: 

713 break 

714 if min_time < pb_min_time: 

715 pb_min_time = min_time 

716 pb = pb // 2 

717 

718 self.compute_plan = optimal_compute_plan 

719 self.JIT_and_test_kernel() 

720 print() 

721 cuda.config.CUDA_LOW_OCCUPANCY_WARNINGS = flag 

722 

723 

724 def autotune_tp(self, initial_tp, delta_tp, initial_compute_plan, timesteps, repeats, optimal_compute_plan, total_min_time, local_min_time): 

725 tp = initial_tp 

726 tp_min_time = 1e9 

727 while 0 < tp <= 64: 

728 #print(self.compute_plan['pb'], tp, (self.compute_plan['pb']*tp) % 32) 

729 if (self.compute_plan['pb']*tp) % 32 == 0: 

730 self.compute_plan['tp'] = tp 

731 self.JIT_and_test_kernel(adjust_compute_plan=False) 

732 if self.compute_plan['tp'] != tp: #or self.compute_plan['gridsync'] != gridsync:  

733 break 

734 total_min_time, local_min_time, min_time = self.autotune_skin(initial_compute_plan['skin'], 0.2, 

735 timesteps, repeats, total_min_time, local_min_time, 

736 optimal_compute_plan, verbose=False) 

737 if min_time > 1.05 * tp_min_time: 

738 break 

739 if min_time < tp_min_time: 

740 tp_min_time = min_time 

741 tp += delta_tp 

742 tp = initial_tp - delta_tp 

743 while 0 < tp <= 64: 

744 if (self.compute_plan['pb']*tp) % 32 == 0: 

745 self.compute_plan['tp'] = tp 

746 self.JIT_and_test_kernel(adjust_compute_plan=False) 

747 if self.compute_plan['tp'] != tp: #or self.compute_plan['gridsync'] != gridsync:  

748 break 

749 total_min_time, local_min_time, min_time = self.autotune_skin(initial_compute_plan['skin'], 0.2, 

750 timesteps, repeats, total_min_time, local_min_time, 

751 optimal_compute_plan, verbose=False) 

752 if min_time > 1.05 * tp_min_time: 

753 break 

754 if min_time < tp_min_time: 

755 tp_min_time = min_time 

756 tp -= delta_tp 

757 return total_min_time, local_min_time, tp_min_time 

758 

759 

760 def autotune_skin(self, initial_skin, delta_skin, timesteps, repeats, total_min_time, local_min_time, optimal_compute_plan, verbose=False): 

761 min_time = 1e9 

762 min_skin = -0.1 

763 

764 # Upscan