Coverage for tests/test_gridsync.py: 6%

78 statements  

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

1""" Check CUDA availability, versions, and test if gridsync is supported. """ 

2 

3 

4def check_cuda(verbose=True): 

5 """ Check CUDA availability, versions, and test if gridsync is supported. Returns True if gridsync is supported.""" 

6 import numba 

7 from numba import cuda 

8 

9 if verbose: 

10 print(' ..:: CUDA information ::..') 

11 print(f'{numba.__version__ = }') 

12 print(f'{numba.cuda.is_available() = }') 

13 print(f'{numba.cuda.is_supported_version() = }') 

14 print(f'{cuda.runtime.get_version() = }') 

15 

16 

17def gridsync_example(): 

18 """ Example from https://numba.readthedocs.io/en/stable/cuda/cooperative_groups.html """ 

19 from numba import cuda, int32, config 

20 import numpy as np 

21 

22 config.CUDA_LOW_OCCUPANCY_WARNINGS = False 

23 config.CUDA_WARN_ON_IMPLICIT_COPY = False 

24 

25 sig = (int32[:, ::1],) 

26 

27 @cuda.jit(sig) 

28 def sequential_rows(M): 

29 col = cuda.grid(1) 

30 g = cuda.cg.this_grid() 

31 

32 rows = M.shape[0] 

33 cols = M.shape[1] 

34 

35 for row in range(1, rows): 

36 opposite = cols - col - 1 

37 M[row, col] = M[row - 1, opposite] + 1 

38 g.sync() 

39 

40 A = np.zeros((1024, 1024), dtype=np.int32) 

41 blockdim = 32 

42 griddim = A.shape[1] // blockdim 

43 sequential_rows[griddim, blockdim](A) 

44 

45 

46 overload = sequential_rows.overloads[(int32[:, ::1],)] 

47 max_blocks = overload.max_cooperative_grid_blocks(blockdim) 

48 return max_blocks 

49 

50def check_gridsync(verbose=True): 

51 """ Check if gridsync is supported. Returns True if gridsync is supported.""" 

52 import numba 

53 from numba import cuda 

54 # See if gridsync is working 

55 max_blocks = None 

56 try: 

57 max_blocks = gridsync_example() 

58 except numba.cuda.cudadrv.driver.LinkerError as e: 

59 if verbose: 

60 print('Warning: gridsync is not supported. Try this hack:') 

61 print('Find where libcudadevrt.a is located, and write something like this') 

62 print(' ln -s /usr/lib/x86_64-linux-gnu/libcudadevrt.a .') 

63 print('in the directory where you run the code.') 

64 return False 

65 

66 if verbose: 

67 print(' ..:: Gridsync check ::..') 

68 print('Confirmed that gridsync is supported by executing test code.') 

69 print(f'{max_blocks = }') 

70 return True 

71 

72def check_gpu(device_id=None): 

73 """ Print some information about the GPU. """ 

74 from numba import cuda 

75 

76 # Get the device 

77 if device_id==None: 

78 device = cuda.get_current_device() 

79 else: 

80 device = cuda.select_device(device_id) 

81 

82 from gamdpy.cc_cores_per_SM_dict import cc_cores_per_SM_dict 

83 if device.compute_capability in cc_cores_per_SM_dict: 

84 cc_cores_per_SM = cc_cores_per_SM_dict[device.compute_capability] 

85 else: 

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

87 cc_cores_per_SM=128 

88 

89 # Print relevant attributes 

90 print(' ..:: GPU information ::..') 

91 print("Device Name:", device.name) 

92 print("Compute Capability:", device.compute_capability) 

93 print("Number of Streaming Multiprocessors:", device.MULTIPROCESSOR_COUNT) 

94 print("Total number of cores:", cc_cores_per_SM*device.MULTIPROCESSOR_COUNT) 

95 print("Max Threads Per Block:", device.MAX_THREADS_PER_BLOCK) 

96 print("Max Block Dimensions (x, y, z):", 

97 device.MAX_BLOCK_DIM_X, device.MAX_BLOCK_DIM_Y, device.MAX_BLOCK_DIM_Z) 

98 print("Max Grid Dimensions (x, y, z):", 

99 device.MAX_GRID_DIM_X, device.MAX_GRID_DIM_Y, device.MAX_GRID_DIM_Z) 

100 print("Max Shared Memory Per Block:", device.MAX_SHARED_MEMORY_PER_BLOCK) 

101 print("Total Constant Memory:", device.TOTAL_CONSTANT_MEMORY) 

102 print("Warp Size:", device.WARP_SIZE) 

103 print("L2 cache size:", device.L2_CACHE_SIZE) 

104 print("Max registers per block:", device.MAX_REGISTERS_PER_BLOCK) 

105 print("Single to double performance ratio:", device.SINGLE_TO_DOUBLE_PRECISION_PERF_RATIO) 

106 

107 

108if __name__ == '__main__': 

109 check_cuda() 

110 check_gridsync() 

111 check_gpu()