Coverage for gpaw/test/gpu/test_kernels.py: 19%

1import pytest 

2import numpy as np 

3 

4from gpaw.utilities import as_real_dtype 

5from gpaw.gpu import cupy as cp, cupy_is_fake 

6 

7seed = 42 

8 

9 

10@pytest.mark.gpu 

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

12@pytest.mark.parametrize("dtype", [np.float32, np.float64, 

13 np.complex64, np.complex128]) 

14def test_dH_aii_times_P_ani(dtype): 

15 from _gpaw import dH_aii_times_P_ani_gpu as kernel_call 

16 from gpaw.purepython import dH_aii_times_P_ani_gpu as cupy_call 

17 assert cupy_call is not kernel_call 

18 

19 rng = cp.random.RandomState(seed) 

20 aN = 3 

21 nN = 5 

22 iN = 12 

23 IN = aN * iN 

24 

25 rdtype = as_real_dtype(dtype) 

26 dH_aii = rng.randn(aN * iN * iN, dtype=as_real_dtype(dtype)) 

27 P_nI = rng.randn(nN, IN, dtype=as_real_dtype(dtype)) 

28 if np.issubdtype(dtype, np.complexfloating): 

29 P_nI = P_nI + 1j * rng.randn(nN, IN, dtype=rdtype) 

30 out_kernel_nI = cp.zeros((nN, IN), dtype=dtype) 

31 out_cupy_nI = cp.zeros((nN, IN), dtype=dtype) 

32 ni_a = cp.ones(aN, dtype=np.int32) * iN 

33 

34 kernel_call(dH_aii, ni_a, P_nI, out_kernel_nI) 

35 cupy_call(dH_aii, ni_a, P_nI, out_cupy_nI) 

36 assert out_kernel_nI.get() == \ 

37 pytest.approx(out_cupy_nI.get(), abs=1e-5) 

38 

39 

40@pytest.mark.gpu 

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

42@pytest.mark.parametrize("dtype", [np.float32, np.float64, 

43 np.complex64, np.complex128]) 

44@pytest.mark.parametrize("cc", [True, False]) 

45def test_pwlfc_expand(dtype, cc): 

46 from _gpaw import pwlfc_expand_gpu as kernel_call 

47 from gpaw.purepython import pwlfc_expand_gpu as cupy_call 

48 assert cupy_call is not kernel_call 

49 

50 rng = cp.random.RandomState(seed) 

51 GN = 100 

52 aN = 3 

53 LN = 25 

54 sN = 4 

55 JN = 8 

56 IN = 26 

57 

58 rdtype = as_real_dtype(dtype) 

59 f_Gs = rng.randn(GN, sN, dtype=rdtype) 

60 Gk_Gv = rng.randn(GN, 3, dtype=rdtype) 

61 pos_av = rng.randn(aN, 3, dtype=rdtype) 

62 eikR_a = rng.randn(aN, dtype=rdtype) \ 

63 + 1j * rng.randn(aN, dtype=rdtype) 

64 Y_GL = rng.randn(GN, LN, dtype=rdtype) 

65 

66 gN = GN if np.issubdtype(dtype, np.complexfloating) else 2 * GN 

67 f_kernel_GI = cp.zeros((gN, IN), dtype=dtype) 

68 f_cupy_GI = cp.zeros((gN, IN), dtype=dtype) 

69 l_s = cp.arange(sN, dtype=np.int32) 

70 a_J = cp.ones(JN, dtype=np.int32) 

71 a_J[0] = 0 

72 a_J[-1] = aN - 1 

73 s_J = cp.ones(JN, dtype=np.int32) 

74 s_J[0] = 0 

75 s_J[-1] = sN - 1 

76 I_J = cp.zeros(JN, dtype=np.int32) 

77 I1 = 0 

78 for J, (a, s) in enumerate(zip(a_J, s_J)): 

79 l = l_s[s] 

80 I2 = I1 + 2 * l + 1 

81 I_J[J] = I1 

82 I1 = I2 

83 assert I2 == IN 

84 

85 kernel_call(f_Gs, Gk_Gv, pos_av, 

86 eikR_a, Y_GL, 

87 l_s, a_J, s_J, 

88 cc, f_kernel_GI, I_J) 

89 cupy_call(f_Gs, Gk_Gv, pos_av, 

90 eikR_a, Y_GL, 

91 l_s, a_J, s_J, 

92 cc, f_cupy_GI, I_J) 

93 

94 assert f_kernel_GI.get() == pytest.approx(f_cupy_GI.get(), abs=1e-6) 

95 

96 

97@pytest.mark.gpu 

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

99@pytest.mark.parametrize("dtype", [np.float32, np.float64]) 

100def test_pw_amend_insert_realwf(dtype): 

101 from _gpaw import pw_amend_insert_realwf_gpu as kernel_call 

102 from gpaw.purepython import pw_amend_insert_realwf_gpu as cupy_call 

103 assert cupy_call is not kernel_call 

104 

105 rng = cp.random.RandomState(seed) 

106 nN = 10 

107 QN = (10, 15, 20) 

108 n = 3 

109 m = 7 

110 

111 array_nQ = rng.randn(nN, *QN, dtype=dtype) \ 

112 + 1j * rng.randn(nN, *QN, dtype=dtype) 

113 

114 array_kernel_nQ = array_nQ.copy() 

115 array_cupy_nQ = array_nQ.copy() 

116 

117 kernel_call(array_kernel_nQ, n, m) 

118 cupy_call(array_cupy_nQ, n, m) 

119 

120 assert array_kernel_nQ.get() != pytest.approx(array_nQ.get()) 

121 assert array_kernel_nQ.get() == pytest.approx(array_cupy_nQ.get()) 

122 

123 

124@pytest.mark.gpu 

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

126@pytest.mark.parametrize("dtype", [np.float32, np.float64, 

127 np.complex64, np.complex128]) 

128def test_calculate_residuals(dtype): 

129 from _gpaw import calculate_residuals_gpu as kernel_call 

130 from gpaw.purepython import calculate_residuals_gpu as cupy_call 

131 assert cupy_call is not kernel_call 

132 

133 rng = cp.random.RandomState(seed) 

134 nN = 10 

135 GN = 100 

136 

137 rdtype = as_real_dtype(dtype) 

138 wfs_nG = rng.randn(nN, GN, dtype=rdtype) 

139 if np.issubdtype(dtype, np.complexfloating): 

140 wfs_nG = wfs_nG + 1j * rng.randn(nN, GN, dtype=rdtype) 

141 eps_n = rng.randn(nN, dtype=rdtype) 

142 residual_kernel_nG = cp.zeros((nN, GN), dtype=dtype) 

143 residual_cupy_nG = cp.zeros((nN, GN), dtype=dtype) 

144 

145 kernel_call(residual_kernel_nG, eps_n, wfs_nG) 

146 cupy_call(residual_cupy_nG, eps_n, wfs_nG) 

147 

148 assert residual_kernel_nG.get() == pytest.approx(residual_cupy_nG.get()) 

149 

150 

151@pytest.mark.gpu 

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

153@pytest.mark.parametrize("dtype", [np.float32, np.float64, 

154 np.complex64, np.complex128]) 

155def test_add_to_density(dtype): 

156 from _gpaw import add_to_density_gpu as kernel_call 

157 from gpaw.purepython import add_to_density_gpu as cupy_call 

158 assert cupy_call is not kernel_call 

159 

160 rng = cp.random.RandomState(seed) 

161 nN = 10 

162 RN = (10, 15, 20) 

163 

164 rdtype = as_real_dtype(dtype) 

165 psit_nR = rng.randn(nN, *RN, dtype=rdtype) 

166 if np.issubdtype(dtype, np.complexfloating): 

167 psit_nR = psit_nR + 1j * rng.randn(nN, *RN, dtype=rdtype) 

168 

169 weight_n = rng.randn(nN, dtype=np.float64) 

170 nt_kernel_R = cp.zeros(RN, dtype=np.float64) 

171 nt_cupy_R = cp.zeros(RN, dtype=np.float64) 

172 

173 kernel_call(weight_n, psit_nR, nt_kernel_R) 

174 cupy_call(weight_n, psit_nR, nt_cupy_R) 

175 

176 assert nt_kernel_R.get() == \ 

177 pytest.approx(nt_cupy_R.get(), abs=1e-5) # Not exact? 

178 

179 

180@pytest.mark.gpu 

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

182@pytest.mark.parametrize("dtype", [np.complex64, np.complex128]) 

183def test_pw_norm(dtype): 

184 from _gpaw import pw_norm_gpu as kernel_call 

185 from gpaw.purepython import pw_norm_gpu as cupy_call 

186 assert cupy_call is not kernel_call 

187 

188 rng = cp.random.RandomState(seed) 

189 rdtype = as_real_dtype(dtype) 

190 

191 xN = 20 

192 GN = 420 

193 

194 C_xG = rng.randn(xN, GN, dtype=rdtype) 

195 if np.issubdtype(dtype, np.complexfloating): 

196 C_xG = C_xG + 1j * rng.randn(xN, GN, dtype=rdtype) 

197 result_kernel_x = cp.empty(xN, dtype=rdtype) 

198 result_cupy_x = cp.empty(xN, dtype=rdtype) 

199 

200 kernel_call(result_kernel_x, C_xG) 

201 cupy_call(result_cupy_x, C_xG) 

202 

203 assert result_kernel_x.get() == pytest.approx(result_cupy_x.get()) 

204 

205 

206@pytest.mark.gpu 

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

208@pytest.mark.parametrize("dtype", [np.complex64, np.complex128]) 

209def test_pw_norm_kinetic(dtype): 

210 from _gpaw import pw_norm_kinetic_gpu as kernel_call 

211 from gpaw.purepython import pw_norm_kinetic_gpu as cupy_call 

212 assert cupy_call is not kernel_call 

213 

214 rng = cp.random.RandomState(seed) 

215 rdtype = as_real_dtype(dtype) 

216 

217 xN = 20 

218 GN = 420 

219 

220 C_xG = rng.randn(xN, GN, dtype=rdtype) 

221 kin_G = rng.randn(GN, dtype=rdtype) 

222 if np.issubdtype(dtype, np.complexfloating): 

223 C_xG = C_xG + 1j * rng.randn(xN, GN, dtype=rdtype) 

224 result_kernel_x = cp.zeros(xN, dtype=rdtype) 

225 result_cupy_x = cp.zeros(xN, dtype=rdtype) 

226 

227 kernel_call(result_kernel_x, C_xG, kin_G) 

228 cupy_call(result_cupy_x, C_xG, kin_G) 

229 

230 assert result_kernel_x.get() == \ 

231 pytest.approx(result_cupy_x.get(), abs=1e-4) # Not exact? 

232 

233 

234@pytest.mark.gpu 

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

236@pytest.mark.parametrize("dtype", [np.complex64, np.complex128]) 

237def test_pw_insert(dtype): 

238 from _gpaw import pw_insert_gpu as kernel_call 

239 from gpaw.purepython import pw_insert_gpu as cupy_call 

240 assert cupy_call is not kernel_call 

241 

242 rng = cp.random.RandomState(seed) 

243 rdtype = as_real_dtype(dtype) 

244 

245 bN = 20 

246 nN = bN 

247 GN = 240 

248 QN = 360 

249 scale = 1.0 

250 nx = 6 

251 ny = 6 

252 nz = 10 

253 

254 Q_G = cp.arange(GN, dtype=np.int32) 

255 Q_G[-1] = QN - 1 

256 psit_nG = rng.randn(nN, GN, dtype=rdtype) 

257 if np.issubdtype(dtype, np.complexfloating): 

258 psit_nG = psit_nG + 1j * rng.randn(nN, GN, dtype=rdtype) 

259 psit_kernel_bQ = cp.zeros((bN, QN), dtype=dtype) 

260 psit_cupy_bQ = cp.zeros((bN, QN), dtype=dtype) 

261 

262 kernel_call(psit_nG, Q_G, scale, psit_kernel_bQ, nx, ny, nz) 

263 cupy_call(psit_nG, Q_G, scale, psit_cupy_bQ, nx, ny, nz) 

264 

265 assert psit_kernel_bQ.get() == pytest.approx(psit_cupy_bQ.get())