Coverage for gamdpy/interactions/make_fixed_interactions.py: 11%

1import numba 

2from numba import cuda 

3 

4 

5def make_fixed_interactions(configuration, fixed_potential, compute_plan, verbose=True, ): 

6 """ Generate a kernel for fixed interactions between particles. """ 

7 # Unpack parameters from configuration and compute_plan 

8 D, num_part = configuration.D, configuration.N 

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

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

11 

12 if verbose: 

13 print(f'Generating fixed interactions for {num_part} particles in {D} dimensions:') 

14 print(f'\tpb: {pb}, tp:{tp}, num_blocks:{num_blocks}') 

15 print(f'\tNumber (virtual) particles: {num_blocks * pb}') 

16 print(f'\tNumber of threads {num_blocks * pb * tp}') 

17 

18 # Prepare user-specified functions for inclusion in kernel(s) 

19 # NOTE: Include check they can be called with right parameters and return the right number and type of parameters  

20 

21 potential_calculator = numba.njit(fixed_potential) 

22 

23 def fixed_interactions(grid, vectors, scalars, ptype, sim_box, interaction_parameters): 

24 indices, values = interaction_parameters 

25 num_interactions = indices.shape[0] 

26 num_threads = num_blocks * pb * tp 

27 

28 my_block = cuda.blockIdx.x 

29 local_id = cuda.threadIdx.x 

30 my_t = cuda.threadIdx.y 

31 #global_id = (my_block*pb + local_id)*tp + my_t 

32 global_id = (my_block * pb + local_id) + my_t * cuda.blockDim.x * cuda.gridDim.x # Faster 

33 

34 for index in range(global_id, num_interactions, num_threads): 

35 potential_calculator(vectors, scalars, ptype, sim_box, indices[index], values) 

36 return 

37 

38 fixed_interactions = cuda.jit(device=gridsync)(fixed_interactions) 

39 

40 if gridsync: 

41 return fixed_interactions # return device function 

42 else: 

43 return fixed_interactions[num_blocks, (pb, tp)] # return kernel, incl. launch parameters