Coverage for gpaw/test/gpu/test_blas.py: 10%

1import numpy as np 

2import pytest 

3 

4from gpaw.gpu import cupy as cp, cupy_is_fake 

5from gpaw.utilities.blas import (gpu_axpy, gpu_dotc, gpu_dotu, gpu_gemm, 

6 gpu_gemv, gpu_mmm, gpu_r2k, gpu_rk, gpu_scal, 

7 mmm, r2k, rk) 

8 

9 

10@pytest.mark.gpu 

11@pytest.mark.skipif(cupy_is_fake, reason='No cupy') 

12@pytest.mark.parametrize('dtype', [float, complex]) 

13def test_blas(dtype): 

14 N = 100 

15 rng = np.random.default_rng(seed=42) 

16 a = np.zeros((N, N), dtype=dtype) 

17 b = np.zeros_like(a) 

18 c = np.zeros_like(a) 

19 x = np.zeros((N,), dtype=dtype) 

20 y = np.zeros_like(x) 

21 if dtype == float: 

22 a[:] = rng.random((N, N)) 

23 b[:] = rng.random((N, N)) 

24 c[:] = rng.random((N, N)) 

25 x[:] = rng.random((N,)) 

26 y[:] = rng.random((N,)) 

27 else: 

28 a.real = rng.random((N, N)) 

29 a.imag = rng.random((N, N)) 

30 b.real = rng.random((N, N)) 

31 b.imag = rng.random((N, N)) 

32 c.real = rng.random((N, N)) 

33 c.imag = rng.random((N, N)) 

34 x.real = rng.random((N,)) 

35 x.imag = rng.random((N,)) 

36 y.real = rng.random((N,)) 

37 y.imag = rng.random((N,)) 

38 

39 a_gpu = cp.asarray(a) 

40 b_gpu = cp.asarray(b) 

41 c_gpu = cp.asarray(c) 

42 x_gpu = cp.asarray(x) 

43 y_gpu = cp.asarray(y) 

44 

45 # axpy 

46 y += 0.5 * x 

47 check_cpu = y.sum() 

48 

49 gpu_axpy(0.5, x_gpu, y_gpu) 

50 check_gpu = y_gpu.sum().get() 

51 

52 assert check_cpu == pytest.approx(check_gpu, rel=1e-14) 

53 

54 # mmm 

55 mmm(0.5, a, 'N', b, 'N', 0.2, c) 

56 check_cpu = c.sum() 

57 

58 gpu_mmm(0.5, a_gpu, 'N', b_gpu, 'N', 0.2, c_gpu) 

59 check_gpu = c_gpu.sum().get() 

60 

61 assert check_cpu == pytest.approx(check_gpu, rel=1e-14) 

62 

63 # gemm 

64 c *= 0.2 

65 c += 0.5 * b @ a 

66 check_cpu = c.sum() 

67 

68 gpu_gemm(0.5, a_gpu, b_gpu, 0.2, c_gpu) 

69 check_gpu = c_gpu.sum().get() 

70 

71 assert check_cpu == pytest.approx(check_gpu, rel=1e-14) 

72 

73 # gemv 

74 y *= 0.2 

75 y += 0.5 * a @ x 

76 check_cpu = y.sum() 

77 

78 gpu_gemv(0.5, a_gpu, x_gpu, 0.2, y_gpu) 

79 check_gpu = y_gpu.sum().get() 

80 

81 assert check_cpu == pytest.approx(check_gpu, rel=1e-14) 

82 

83 # rk 

84 rk(0.5, a, 0.2, c) 

85 check_cpu = c.sum() 

86 

87 gpu_rk(0.5, a_gpu, 0.2, c_gpu) 

88 check_gpu = c_gpu.sum().get() 

89 

90 assert check_cpu == pytest.approx(check_gpu, rel=1e-14) 

91 

92 # r2k 

93 r2k(0.5, a, b, 0.2, c) 

94 check_cpu = c.sum() 

95 

96 gpu_r2k(0.5, a_gpu, b_gpu, 0.2, c_gpu) 

97 check_gpu = c_gpu.sum().get() 

98 

99 assert check_cpu == pytest.approx(check_gpu, rel=1e-14) 

100 

101 # dotc 

102 check_cpu = x.conj() @ y 

103 

104 check_gpu = gpu_dotc(x_gpu, y_gpu) 

105 

106 assert check_cpu == pytest.approx(check_gpu, rel=1e-14) 

107 

108 # dotu 

109 check_cpu = x @ y 

110 

111 check_gpu = gpu_dotu(x_gpu, y_gpu) 

112 

113 assert check_cpu == pytest.approx(check_gpu, rel=1e-14) 

114 

115 # scal 

116 a *= 0.5 

117 check_cpu = a.sum() 

118 

119 gpu_scal(0.5, a_gpu) 

120 check_gpu = a_gpu.sum().get() 

121 

122 assert check_cpu == pytest.approx(check_gpu, rel=1e-14)