Coverage for strongcoca / compute_config.py: 100%

1from dataclasses import dataclass 

2from typing import Optional 

3 

4_VALID_BACKENDS = ('none', 'torch', 'cupy') 

5_VALID_PRECISIONS = ('float32', 'float64') 

6 

7 

8@dataclass(frozen=True) 

9class ComputeConfig: 

10 """Global runtime configuration for the O(N³) linear-algebra operations 

11 in :class:`~strongcoca.calculators.PolarizabilityCalculator`. 

12 

13 Parameters 

14 ---------- 

15 backend 

16 ``'none'`` (CPU only), ``'torch'`` (PyTorch/CUDA), or ``'cupy'`` (CuPy/CUDA). 

17 precision 

18 Floating-point precision for GPU operations. ``'float32'`` casts to 

19 complex64 before the GPU kernel and back to complex128 afterwards 

20 (roughly 10x faster on consumer GPUs where FP64 throughput is 

21 restricted); ``'float64'`` keeps complex128 throughout. 

22 max_solve_mem 

23 Memory budget in MiB for intermediate matrices in the CPU linear 

24 solve step. The frequency axis is chunked so that the A matrix 

25 stays below this limit. 

26 gpu_batched_max 

27 Maximum value of N3 = 3 x N_particles for which the cuBLAS batched 

28 solver is used in GPU kernels. Above this threshold, the code loops 

29 over the batch dimension to invoke the non-batched cuSOLVER path, 

30 which is faster for large matrices on consumer GPUs. Increase this 

31 value (e.g. to 1024 or higher) on A100/H100/H200 cards, where the 

32 batched cuSOLVER path performs well at large N. 

33 """ 

34 

35 backend: str 

36 precision: str 

37 max_solve_mem: float 

38 gpu_batched_max: int 

39 

40 

41_config = ComputeConfig(backend='none', precision='float32', 

42 max_solve_mem=80.0, gpu_batched_max=32) 

43 

44 

45def set_compute_config(backend: Optional[str] = None, 

46 precision: Optional[str] = None, 

47 max_solve_mem: Optional[float] = None, 

48 gpu_batched_max: Optional[int] = None) -> None: 

49 """Update one or more fields of the global compute configuration. 

50 

51 Any argument left as ``None`` keeps its current value. 

52 All arguments are validated before anything is written, so an invalid 

53 call leaves the configuration completely unchanged. 

54 

55 Parameters 

56 ---------- 

57 backend 

58 See :class:`ComputeConfig`. 

59 precision 

60 See :class:`ComputeConfig`. 

61 max_solve_mem 

62 See :class:`ComputeConfig`. 

63 gpu_batched_max 

64 See :class:`ComputeConfig`. 

65 

66 Examples 

67 -------- 

68 >>> from strongcoca import set_compute_config, get_compute_config 

69 >>> set_compute_config(max_solve_mem=200) 

70 >>> get_compute_config().max_solve_mem 

71 200.0 

72 >>> set_compute_config(max_solve_mem=80) # reset to the default 

73 """ 

74 global _config 

75 

76 new_backend = _config.backend if backend is None else backend 

77 new_precision = _config.precision if precision is None else precision 

78 new_max_solve_mem = _config.max_solve_mem if max_solve_mem is None else float(max_solve_mem) 

79 new_gpu_batched_max = _config.gpu_batched_max 

80 if gpu_batched_max is not None: 

81 new_gpu_batched_max = int(gpu_batched_max) 

82 

83 if new_backend not in _VALID_BACKENDS: 

84 raise ValueError(f'backend must be one of {_VALID_BACKENDS}, got {new_backend!r}') 

85 if new_precision not in _VALID_PRECISIONS: 

86 raise ValueError(f'precision must be one of {_VALID_PRECISIONS}, got {new_precision!r}') 

87 if new_max_solve_mem <= 0: 

88 raise ValueError(f'max_solve_mem must be positive, got {new_max_solve_mem}') 

89 if new_gpu_batched_max <= 0: 

90 raise ValueError(f'gpu_batched_max must be positive, got {new_gpu_batched_max}') 

91 

92 _config = ComputeConfig(backend=new_backend, 

93 precision=new_precision, 

94 max_solve_mem=new_max_solve_mem, 

95 gpu_batched_max=new_gpu_batched_max) 

96 

97 

98def get_compute_config() -> ComputeConfig: 

99 """Return the current global compute configuration.""" 

100 return _config