Coverage for gamdpy/runtime_actions/runtime_action.py: 80%

1from numba import cuda 

2from abc import ABC, abstractmethod 

3from gamdpy import Configuration 

4from typing import Callable 

5 

6 

7class RuntimeAction(ABC): 

8 """ 

9 Abstract Base Class specifying the requirements for a runtime_action, i.e. an action to compiled into to innner MD kernel 

10 """ 

11 

12 def get_compute_flags(self): 

13 return None 

14 

15 def setup(self, configuration: Configuration, num_timeblocks: int, steps_per_timeblock: int, output, verbose=False) -> None: 

16 pass 

17 

18 @abstractmethod 

19 def get_prestep_kernel(self, configuration: Configuration, compute_plan: dict) -> Callable: 

20 """ 

21 Get a kernel (or python function depending on compute_plan["gridsync"]) that implements the runtime_action. 

22 The generated kernel is called after evaluation of interactions, before intergration step is performed, see class Simulation 

23 """ 

24 

25 pass 

26 

27 @abstractmethod 

28 def get_poststep_kernel(self, configuration: Configuration, compute_plan: dict) -> Callable: 

29 """ 

30 Get a kernel (or python function depending on compute_plan["gridsync"]) that implements the runtime_action 

31 The generated kernel is called immediately after evaluation intergration step is performed, see class Simulation 

32 """ 

33 

34 pass 

35 

36 @abstractmethod 

37 def get_params(self, configuration: Configuration, compute_plan: dict) -> tuple : 

38 """ 

39 Get a tuple with the parameters expected by the associated kernel 

40 """ 

41 

42 pass 

43 

44 def update_at_end_of_timeblock(self, timeblock: int, output_reference): 

45 """ 

46 Method to be called at the end of a timeblock, for e.g. saving data to a file if needed 

47 """ 

48 

49 pass 

50 

51 def initialize_before_timeblock(self, timeblock: int, output_reference): 

52 """ 

53 Method to be called before each timeblock  

54 """ 

55 

56 pass 

57 

58def merge_runtime_actions(configuration: Configuration, prestep_kernelA: Callable, poststep_kernelA: Callable, paramsA: tuple, actionB: RuntimeAction, compute_plan: dict) -> tuple[Callable, Callable, tuple] : 

59 paramsB = actionB.get_params(configuration, compute_plan) 

60 prestep_kernelB = actionB.get_prestep_kernel(configuration, compute_plan) 

61 poststep_kernelB = actionB.get_poststep_kernel(configuration, compute_plan) 

62 

63 if compute_plan['gridsync']: 

64 # A device function, calling a number of device functions, using gridsync to syncronize 

65 @cuda.jit( device=compute_plan['gridsync']) 

66 def prestep_kernel(grid, vectors, scalars, r_im, sim_box, step, runtime_actions_params): 

67 prestep_kernelA(grid, vectors, scalars, r_im, sim_box, step, runtime_actions_params[0]) 

68 grid.sync() # Not always necessary !!! 

69 prestep_kernelB(grid, vectors, scalars, r_im, sim_box, step, runtime_actions_params[1]) 

70 return 

71 

72 @cuda.jit( device=compute_plan['gridsync']) 

73 def poststep_kernel(grid, vectors, scalars, r_im, sim_box, step, runtime_actions_params): 

74 poststep_kernelA(grid, vectors, scalars, r_im, sim_box, step, runtime_actions_params[0]) 

75 grid.sync() # Not always necessary !!! 

76 poststep_kernelB(grid, vectors, scalars, r_im, sim_box, step, runtime_actions_params[1]) 

77 return 

78 

79 return prestep_kernel, poststep_kernel, (paramsA, paramsB, ) 

80 else: 

81 # Two python function, making several kernel calls to syncronize 

82 def prestep_kernel(grid, vectors, scalars, r_im, sim_box, step, runtime_actions_params): 

83 prestep_kernelA(grid, vectors, scalars, r_im, sim_box, step, runtime_actions_params[0]) 

84 prestep_kernelB(grid, vectors, scalars, r_im, sim_box, step, runtime_actions_params[1]) 

85 return 

86 

87 def poststep_kernel(grid, vectors, scalars, r_im, sim_box, step, runtime_actions_params): 

88 poststep_kernelA(grid, vectors, scalars, r_im, sim_box, step, runtime_actions_params[0]) 

89 poststep_kernelB(grid, vectors, scalars, r_im, sim_box, step, runtime_actions_params[1]) 

90 return 

91 

92 return prestep_kernel, poststep_kernel, (paramsA, paramsB, ) 

93 

94def add_runtime_actions_list(configuration: Configuration, runtime_actions_list: list[RuntimeAction], compute_plan: dict, verbose: bool = False) -> tuple[Callable, Callable, tuple]: 

95 

96 # Setup first interaction and cuda.jit it if gridsync is used for syncronization 

97 params = runtime_actions_list[0].get_params(configuration, compute_plan) 

98 prestep_kernel = runtime_actions_list[0].get_prestep_kernel(configuration, compute_plan) 

99 poststep_kernel = runtime_actions_list[0].get_poststep_kernel(configuration, compute_plan) 

100 

101 if compute_plan['gridsync']: 

102 prestep_kernel: Callable = cuda.jit( device=compute_plan['gridsync'] )(prestep_kernel) 

103 poststep_kernel: Callable = cuda.jit( device=compute_plan['gridsync'] )(poststep_kernel) 

104 

105 # Merge in the rest of the runtime_actions (maximum recursion depth might set a maximum for number of interactions) 

106 for i in range(1, len(runtime_actions_list)): 

107 prestep_kernel, poststep_kernel, params = merge_runtime_actions(configuration, prestep_kernel, poststep_kernel, params, runtime_actions_list[i], compute_plan) 

108 

109 return prestep_kernel, poststep_kernel, params