Coverage for gamdpy/simulation/get_default_compute_plan.py: 90%

49 statements  

« prev     ^ index     » next       coverage.py v7.9.1, created at 2025-06-14 15:55 +0200

1import numpy as np 

2import numba 

3from numba import cuda 

4import math 

5import matplotlib.pyplot as plt 

6import os 

7 

8def get_default_compute_plan(configuration): 

9 """ Return a default compute_plan 

10 The default compute_plan is a dictionary with a set of parameters specifying how computations are done on the GPU. 

11 The returned plan depends on the number of particles, and properties of the GPU. The keys of the dictionary are: 

12 

13 - 'pb': particle per thread block 

14 - 'tp': threads per particle 

15 - 'gridsync': Boolean indicating if syncronization should be done by grid.sync() calls 

16 - 'skin': used when updating nblist 

17 - 'UtilizeNIII': Boolean indicating if Newton's third law (NIII) should be utilized (see pairpotential_calculator). 

18 - 'nblist' : 'N squared' (default) or 'linked lists'. Determines algorithm updating nblist 

19 

20 """ 

21 N = configuration.N 

22 

23 # Get relevant info about the device. At some point we should be able to deal with no device (GPU) available 

24 if os.getenv("NUMBA_ENABLE_CUDASIM") != "1": 

25 # Trying to handle no device (GPU) case 

26 # NUMBA_ENABLE_CUDASIM environment variable is set to "1" if the cuda simulator is used. 

27 # See: https://numba.pydata.org/numba-doc/dev/cuda/simulator.html 

28 device = cuda.get_current_device() 

29 

30 # Apperently we can't ask the device about how many cores it has, neither in total or per SM (Streaming Processor), 

31 # so we read the latter from a stored dictionary dependent on the compute capability. 

32 from gamdpy.cc_cores_per_SM_dict import cc_cores_per_SM_dict 

33 if device.compute_capability in cc_cores_per_SM_dict: 

34 cc_cores_per_SM = cc_cores_per_SM_dict[device.compute_capability] 

35 else: 

36 print('gamdpy WARNING: Could not find cc_cores_per_SM for this compute_capability. Guessing: 128') 

37 cc_cores_per_SM = 128 

38 num_SM = device.MULTIPROCESSOR_COUNT 

39 num_cc_cores = cc_cores_per_SM * num_SM 

40 warpsize = device.WARP_SIZE 

41 else: # Sets up the behaviour in case the GPU simulator is active and set num_cc_cores = number of threads 

42 num_SM = 1 

43 num_cc_cores = numba.get_num_threads() 

44 warpsize = 1 

45 

46 # pb: particle per (thread) block 

47 pb = 512 

48 while N // pb < 2 * num_SM and pb >= 8: # Performance heuristic 

49 pb = pb // 2 

50 if pb < 8: 

51 pb = 8 

52 if pb > 256: 

53 pb = 256 

54 

55 # tp: threads per particle 

56 tp = 1 

57 while N * tp < 2 * num_cc_cores: # Performance heuristic (conservative) 

58 tp += 1 

59 

60 while (pb * tp) % warpsize != 0: # Number of threads per thread-block should be multiplum of warpsize 

61 tp += 1 

62 

63 if tp > 16: 

64 tp = 16 

65 

66 # skin: used when updating nblist 

67 skin = 0.5 

68 if N > 6 * 1024: 

69 skin = np.float32( 1.0) # make the nblist be valid for many steps for large N. 

70 

71 # UtilizeNIII: Boolean flag indicating if Newton's third law (NIII) should be utilized (see pairpotential_calculator). 

72 # Utilization of NIII is implemented by using atomic add's to the force array,  

73 # so it is inefficient at small system sizes where a lot of conflicts occur. 

74 UtilizeNIII = True 

75 if N < 16 * 1024: 

76 UtilizeNIII = False 

77 

78 # gridsync: Bolean flag indicating whether synchronization should be done via grid.sync() 

79 gridsync = True 

80 if N * tp > 4 * num_cc_cores: # Heuristic 

81 gridsync = False 

82 

83 nblist = 'N squared' 

84 if N > 6_000: # Heuristic 

85 nblist = 'linked lists' 

86 skin = 0.5 

87 

88 return {'pb': pb, 'tp': tp, 'skin': skin, 

89 'UtilizeNIII': UtilizeNIII, 'gridsync': gridsync, 'nblist': nblist} 

90 

91