Coverage for gpaw/test/gpu/test_hip_patch.py: 32%

1import pytest 

2from gpaw.fftw import rfftn_patch, irfftn_patch 

3from gpaw.gpu import is_hip, cupy_is_fake 

4 

5 

6@pytest.mark.gpu 

7def test_hip_patch_on_cuda(): 

8 if is_hip: 

9 pytest.skip('Test runs only on CUDA (have hip)') 

10 if cupy_is_fake: 

11 pytest.skip('Run runs only on CUDA (have fake cupy)') 

12 

13 import cupy 

14 import cupyx 

15 

16 def matrices(): 

17 for sx in [1, 5, 6, 32, 101, 102]: 

18 for sy in [1, 5, 6, 32, 101, 102]: 

19 for sz in [5, 6, 32, 101, 102]: 

20 yield cupy.random.rand(sx, sy, sz) 

21 for tmp_R in matrices(): 

22 print(tmp_R.shape) 

23 tmp2_G = cupyx.scipy.fft.rfftn(tmp_R) 

24 tmp_G = rfftn_patch(tmp_R) 

25 assert cupy.allclose(tmp_G, tmp2_G, rtol=1e-10, atol=1e-10) 

26 back_R = irfftn_patch(tmp_G, tmp_R.shape) 

27 assert cupy.allclose(tmp_R, back_R, rtol=1e-10, atol=1e-10) 

28 back2_R = cupyx.scipy.fft.irfftn(tmp_G, tmp_R.shape) 

29 assert cupy.allclose(back_R, back2_R, rtol=1e-10, atol=1e-10)