Coverage for gamdpy/interactions/nblist.py: 36%

1import numpy as np 

2import numba 

3from numba import cuda 

4 

5class NbList2(): 

6 

7 def __init__(self, configuration, exclusions, max_num_nbs): 

8 self.nblist = np.zeros((configuration.N, max_num_nbs+1), dtype=np.int32) 

9 self.nbflag = np.zeros(3, dtype=np.int32) 

10 self.r_ref = np.zeros_like(configuration['r']) # Inherits also data type 

11 self.exclusions = exclusions # Should be able to be a list (eg from bonds, angles, etc), and merge  

12 self.d_simbox_last_rebuild = cuda.to_device(np.zeros(configuration.simbox.len_sim_box_data, dtype=np.float32)) 

13 

14 def copy_to_device(self): 

15 self.d_nblist = cuda.to_device(self.nblist) 

16 self.d_nbflag = cuda.to_device(self.nbflag) 

17 self.d_r_ref = cuda.to_device(self.r_ref) 

18 self.d_exclusions = cuda.to_device(self._exclusions) 

19 

20 def get_params(self, max_cut, compute_plan, verbose=True): 

21 self.max_cut = max_cut 

22 self.skin, = [compute_plan[key] for key in ['skin']] 

23 

24 #print('NbList.exclusions:\n', self.exclusions) 

25 if type(self.exclusions) == np.ndarray: # Don't change user-set properties 

26 self._exclusions = self.exclusions.copy() 

27 else: 

28 self._exclusions = np.zeros((self.r_ref.shape[0], 2), dtype=np.int32) 

29 self.copy_to_device() 

30 return (np.float32(self.max_cut), np.float32(self.skin), self.d_nbflag, self.d_r_ref, self.d_exclusions, self.d_simbox_last_rebuild) 

31 

32 def get_kernel(self, configuration, compute_plan, compute_flags, verbose=False, force_update=False): 

33 

34 # Unpack parameters from configuration and compute_plan 

35 D, num_part = configuration.D, configuration.N 

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

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

38 compute_stresses = compute_flags['stresses'] 

39 

40 # Unpack indices for vectors and scalars to be compiled into kernel 

41 r_id, f_id = [configuration.vectors.indices[key] for key in ['r', 'f']] 

42 if compute_stresses: 

43 sx_id = configuration.vectors.indices['sx'] 

44 if D > 1: 

45 sy_id = configuration.vectors.indices['sy'] 

46 if D > 2: 

47 sz_id = configuration.vectors.indices['sz'] 

48 if D > 3: 

49 sw_id = configuration.vectors.indices['sw'] 

50 

51 # JIT compile functions to be compiled into kernel 

52 dist_sq_function = numba.njit(configuration.simbox.get_dist_sq_function()) 

53 dist_moved_exceeds_limit_function = numba.njit(configuration.simbox.get_dist_moved_exceeds_limit_function()) 

54 

55 

56 @cuda.jit( device=gridsync ) 

57 def nblist_check(vectors, sim_box, skin, r_ref, nbflag, simbox_last_rebuild, cut): 

58 """ Check validity of nblist, i.e. did any particle mode more than skin/2 since last nblist update? 

59 Each thread-block checks the assigned particles (global_id) 

60 nbflag[0] = 0 : No update needed 

61 nbflag[0] = num_blocks : Update needed 

62 Kernel configuration: [num_blocks, (pb, tp)] 

63 """ 

64 

65 global_id, my_t = cuda.grid(2) 

66 if force_update: # nblist update forced (for benchmark or similar) 

67 if global_id==0 and my_t==0: 

68 nbflag[0]=num_blocks 

69 

70 if global_id < num_part and my_t==0: 

71 if dist_moved_exceeds_limit_function(vectors[r_id][global_id], r_ref[global_id], sim_box, simbox_last_rebuild, skin, cut): 

72 nbflag[0] = num_blocks 

73 

74 

75 return 

76 

77 @cuda.jit(device=gridsync) 

78 def nblist_update(vectors, sim_box, cut_plus_skin, nbflag, nblist, r_ref, exclusions, simbox_last_rebuild): 

79 """ N^2 Update neighbor-list using numba.cuda  

80 Kernel configuration: [num_blocks, (pb, tp)] 

81 """ 

82 

83 my_block = cuda.blockIdx.x 

84 local_id = cuda.threadIdx.x 

85 global_id = my_block*pb + local_id 

86 my_t = cuda.threadIdx.y 

87 

88 if nbflag[0] > 0: 

89 max_nbs = nblist.shape[1]-1 # Last index is used for storing number of neighbors 

90 

91 if global_id < num_part and my_t==0: 

92 nblist[global_id, max_nbs] = 0 # Set number of neighbors (stored at index max_nbs) to zero 

93 

94 cuda.syncthreads() # wait for nblist[global_id, max_nbs] to be ready 

95 

96 if global_id < num_part: 

97 my_num_exclusions = exclusions[global_id,-1] 

98 for i in range(0, num_part, pb*tp): # Loop over blocks 

99 for j in range(pb): # Loop over particles the pb in block 

100 other_global_id = j + i + my_t*pb 

101 if UtilizeNIII: 

102 TwodN = 2*(other_global_id - global_id) 

103 flag = other_global_id < num_part and (0 < TwodN <= num_part or TwodN < -num_part) 

104 else: 

105 flag = other_global_id != global_id and other_global_id < num_part 

106 if flag: 

107 dist_sq = dist_sq_function(vectors[r_id][other_global_id], vectors[r_id][global_id], sim_box) 

108 if dist_sq < cut_plus_skin*cut_plus_skin: 

109 not_excluded = True # Check exclusion list 

110 for k in range(my_num_exclusions): 

111 if exclusions[global_id, k] == other_global_id: 

112 not_excluded = False 

113 if not_excluded: 

114 my_num_nbs = cuda.atomic.add(nblist, (global_id, max_nbs), 1) # Find next free index in nblist 

115 if my_num_nbs < max_nbs: 

116 nblist[global_id, my_num_nbs] = other_global_id # Last entry is number of neighbors 

117 

118 # Various house-keeping 

119 if global_id < num_part and my_t==0: 

120 for k in range(D): 

121 r_ref[global_id, k] = vectors[r_id][global_id, k] # Store positions for wich nblist was updated ( used in nblist_check() )  

122 if local_id == 0 and my_t==0: 

123 cuda.atomic.add(nbflag, 0, -1) # nbflag[0] = 0 by when all blocks are done 

124 if global_id == 0 and my_t==0: 

125 cuda.atomic.add(nbflag, 2, 1) # Count how many updates are done in nbflag[2] 

126 for k in range(len(simbox_last_rebuild)): 

127 simbox_last_rebuild[k] = sim_box[k] 

128 

129 if my_num_nbs >= max_nbs: # Overflow detected, nbflag[1] should be checked later, and then 

130 cuda.atomic.max(nbflag, 1, my_num_nbs) # re-allocate larger nb-list, and redo computations from last safe state 

131 

132 return 

133 

134 if gridsync==True: 

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

136 @cuda.jit( device=gridsync ) 

137 def check_and_update(grid, vectors, scalars, ptype, sim_box, nblist, nblist_parameters): 

138 max_cut, skin, nbflag, r_ref, exclusions, simbox_last_rebuild = nblist_parameters 

139 nblist_check(vectors, sim_box, skin, r_ref, nbflag, simbox_last_rebuild, max_cut) 

140 grid.sync() 

141 nblist_update(vectors, sim_box, max_cut+skin, nbflag, nblist, r_ref, exclusions, simbox_last_rebuild) 

142 return 

143 return check_and_update 

144 

145 else: 

146 # A python function, making several kernel calls to syncronize  

147 def check_and_update(grid, vectors, scalars, ptype, sim_box, nblist, nblist_parameters): 

148 max_cut, skin, nbflag, r_ref, exclusions, simbox_last_rebuild = nblist_parameters 

149 nblist_check[num_blocks, (pb, 1)](vectors, sim_box, skin, r_ref, nbflag, simbox_last_rebuild, max_cut) 

150 nblist_update[num_blocks, (pb, tp)](vectors, sim_box, max_cut+skin, nbflag, nblist, r_ref, exclusions, simbox_last_rebuild) 

151 return 

152 return check_and_update