Skip to content

Commit 542dc82

Browse files
puzhichenpuzhichen
authored
Move 3rd-derivative of xc to GPU (#529)
* test libxc in gpu bug for 3rd derivatives. * use libxc-0.7.0 * using libxc 0.7.0 * change the unittest.yml * remove some codes * change the unite test settings for multi-gpu * add sudo uninstall gpu4pyscf * fix some typo * print more info * change some thresholds * use gpu to calculate 3-rd functionals * remove some comments * add environment varible to switch functional calculations on CPU or GPU
1 parent ca3c85c commit 542dc82

File tree

3 files changed

+66
-31
lines changed

3 files changed

+66
-31
lines changed

‎gpu4pyscf/dft/tests/test_libxc.py‎

Lines changed: 34 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
from gpu4pyscf.dft.numint import NumInt as numint_gpu
2121
from pyscf.dft.numint import NumInt as numint_cpu
2222
import cupy
23+
import os
2324

2425
def setUpModule():
2526
global mol, dm1, dm0
@@ -51,7 +52,7 @@ def _diff(dat, ref):
5152
return np.min((abs(d/(ref+1e-300)), abs(d)), axis=0)
5253

5354
class KnownValues(unittest.TestCase):
54-
def _check_xc(self, xc, spin=0, fxc_tol=1e-10, kxc_tol=1e-10):
55+
def _check_xc(self, xc, spin=0, deriv=2, fxc_tol=1e-10, kxc_tol=1e-10):
5556
ni_cpu = numint_cpu()
5657
ni_gpu = numint_gpu()
5758
xctype = ni_cpu._xc_type(xc)
@@ -66,8 +67,15 @@ def _check_xc(self, xc, spin=0, fxc_tol=1e-10, kxc_tol=1e-10):
6667
if spin != 0:
6768
rho = (rho, rho)
6869

69-
exc_cpu, vxc_cpu, fxc_cpu, kxc_cpu = ni_cpu.eval_xc_eff(xc, rho, deriv=2, xctype=xctype)
70-
exc_gpu, vxc_gpu, fxc_gpu, kxc_gpu = ni_gpu.eval_xc_eff(xc, cupy.array(rho), deriv=2, xctype=xctype)
70+
exc_cpu, vxc_cpu, fxc_cpu, kxc_cpu = ni_cpu.eval_xc_eff(xc, rho, deriv=deriv, xctype=xctype)
71+
exc_gpu, vxc_gpu, fxc_gpu, kxc_gpu = ni_gpu.eval_xc_eff(xc, cupy.array(rho), deriv=deriv, xctype=xctype)
72+
73+
print(f"{xc} {spin} exc", _diff(exc_gpu[:,0].get(), exc_cpu).max())
74+
print(f"{xc} {spin} vxc", _diff(vxc_gpu.get(), vxc_cpu).max())
75+
if fxc_gpu is not None:
76+
print(f"{xc} {spin} fxc", _diff(fxc_gpu.get(), fxc_cpu).max())
77+
if kxc_gpu is not None:
78+
print(f"{xc} {spin} kxc", _diff(kxc_gpu.get(), kxc_cpu).max())
7179

7280
assert _diff(exc_gpu[:,0].get(), exc_cpu).max() < 1e-10
7381
assert _diff(vxc_gpu.get(), vxc_cpu).max() < 1e-10
@@ -77,15 +85,32 @@ def _check_xc(self, xc, spin=0, fxc_tol=1e-10, kxc_tol=1e-10):
7785
assert _diff(kxc_gpu.get(), kxc_cpu).max() < kxc_tol
7886

7987
def test_LDA(self):
80-
self._check_xc('LDA_C_VWN')
88+
whether_use_gpu = os.environ.get('LIBXC_ON_GPU', '0') == '1'
89+
if whether_use_gpu:
90+
deriv = 3
91+
print("test LDA with deriv 3")
92+
else:
93+
deriv = 2
94+
print("test LDA with deriv 2")
95+
self._check_xc('LDA_C_VWN', deriv=deriv)
8196

8297
def test_GGA(self):
83-
self._check_xc('HYB_GGA_XC_B3LYP')
84-
self._check_xc('GGA_X_B88', fxc_tol=1e-10)
85-
self._check_xc('GGA_C_PBE', fxc_tol=1e-4)
98+
whether_use_gpu = os.environ.get('LIBXC_ON_GPU', '0') == '1'
99+
if whether_use_gpu:
100+
deriv = 3
101+
else:
102+
deriv = 2
103+
self._check_xc('HYB_GGA_XC_B3LYP', deriv=deriv, kxc_tol=1e-9)
104+
self._check_xc('GGA_X_B88', fxc_tol=1e-10, deriv=deriv, kxc_tol=1e-8)
105+
self._check_xc('GGA_C_PBE', fxc_tol=1e-4, deriv=deriv, kxc_tol=3e2)
86106

87107
def test_mGGA(self):
88-
self._check_xc('MGGA_C_M06', fxc_tol=1e-4)
108+
whether_use_gpu = os.environ.get('LIBXC_ON_GPU', '0') == '1'
109+
if whether_use_gpu:
110+
deriv = 3
111+
else:
112+
deriv = 2
113+
self._check_xc('MGGA_C_M06', fxc_tol=1e-4, deriv=deriv, kxc_tol=1e-2)
89114

90115
def test_u_LDA(self):
91116
self._check_xc('LDA_C_VWN', spin=1)
@@ -101,4 +126,4 @@ def test_u_mGGA(self):
101126

102127
if __name__ == "__main__":
103128
print("Full Tests for xc fun")
104-
unittest.main()
129+
unittest.main()

‎gpu4pyscf/grad/tdrks.py‎

Lines changed: 20 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
from gpu4pyscf.grad import rks as rks_grad
2828
from gpu4pyscf.grad import tdrhf
2929
from gpu4pyscf import tdscf
30+
import os
3031

3132

3233
#
@@ -372,13 +373,16 @@ def _contract_xc_kernel(td_grad, xc_code, dmvo, dmoo=None,
372373
rho = ni.eval_rho2(_sorted_mol, ao0, mo_coeff_mask, mo_occ, mask, xctype, with_lapl=False)
373374
# quick fix
374375
if deriv > 2:
375-
ni_cpu = numint_cpu()
376-
# TODO: If the libxc is stablized, this should be gpulized
377-
# vxc, fxc, kxc = ni.eval_xc_eff(xc_code, rho, deriv, xctype=xctype)[1:]
378-
vxc, fxc, kxc = ni_cpu.eval_xc_eff(xc_code, rho.get(), deriv, xctype=xctype)[1:]
379-
if isinstance(vxc,np.ndarray): vxc = cp.asarray(vxc)
380-
if isinstance(fxc,np.ndarray): fxc = cp.asarray(fxc)
381-
if isinstance(kxc,np.ndarray): kxc = cp.asarray(kxc)
376+
whether_use_gpu = os.environ.get('LIBXC_ON_GPU', '0') == '1'
377+
if not whether_use_gpu:
378+
ni_cpu = numint_cpu()
379+
# TODO: If the libxc is stablized, this should be gpulized
380+
vxc, fxc, kxc = ni_cpu.eval_xc_eff(xc_code, rho.get(), deriv, xctype=xctype)[1:]
381+
if isinstance(vxc,np.ndarray): vxc = cp.asarray(vxc)
382+
if isinstance(fxc,np.ndarray): fxc = cp.asarray(fxc)
383+
if isinstance(kxc,np.ndarray): kxc = cp.asarray(kxc)
384+
else:
385+
vxc, fxc, kxc = ni.eval_xc_eff(xc_code, rho, deriv, xctype=xctype)[1:]
382386
else:
383387
vxc, fxc, kxc = ni.eval_xc_eff(xc_code, rho, deriv, xctype=xctype)[1:]
384388
dmvo_mask = dmvo[mask[:, None], mask]
@@ -433,14 +437,15 @@ def _contract_xc_kernel(td_grad, xc_code, dmvo, dmoo=None,
433437
rho *= 0.5
434438
rho = cp.repeat(rho[cp.newaxis], 2, axis=0)
435439
# quick fix
436-
if deriv > 2:
437-
ni_cpu = numint_cpu()
438-
vxc, fxc, kxc = ni_cpu.eval_xc_eff(xc_code, rho.get(), deriv, xctype=xctype)[1:]
439-
if isinstance(vxc,np.ndarray): vxc = cp.asarray(vxc)
440-
if isinstance(fxc,np.ndarray): fxc = cp.asarray(fxc)
441-
if isinstance(kxc,np.ndarray): kxc = cp.asarray(kxc)
442-
else:
443-
vxc, fxc, kxc = ni.eval_xc_eff(xc_code, rho, deriv, xctype=xctype)[1:]
440+
# if deriv > 2:
441+
# ni_cpu = numint_cpu()
442+
# vxc, fxc, kxc = ni_cpu.eval_xc_eff(xc_code, rho.get(), deriv, xctype=xctype)[1:]
443+
# if isinstance(vxc,np.ndarray): vxc = cp.asarray(vxc)
444+
# if isinstance(fxc,np.ndarray): fxc = cp.asarray(fxc)
445+
# if isinstance(kxc,np.ndarray): kxc = cp.asarray(kxc)
446+
# else:
447+
# vxc, fxc, kxc = ni.eval_xc_eff(xc_code, rho, deriv, xctype=xctype)[1:]
448+
vxc, fxc, kxc = ni.eval_xc_eff(xc_code, rho, deriv, xctype=xctype)[1:]
444449
# fxc_t couples triplet excitation amplitudes
445450
# 1/2 int (tia - tIA) fxc (tjb - tJB) = tia fxc_t tjb
446451
fxc_t = fxc[:, :, 0] - fxc[:, :, 1]

‎gpu4pyscf/grad/tduks.py‎

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
from gpu4pyscf.grad import tdrks
2727
from gpu4pyscf.scf import ucphf
2828
from gpu4pyscf import tdscf
29+
import os
2930

3031

3132
#
@@ -453,13 +454,17 @@ def _contract_xc_kernel(td_grad, xc_code, dmvo, dmoo=None, with_vxc=True, with_k
453454
ni.eval_rho2(_sorted_mol, ao0, mo_coeff_mask_a, mo_occ[0], mask, xctype,with_lapl=False),
454455
ni.eval_rho2(_sorted_mol, ao0, mo_coeff_mask_b, mo_occ[1], mask, xctype, with_lapl=False)))
455456
if deriv > 2:
456-
ni_cpu = numint_cpu()
457-
# TODO: If the libxc is stablized, this should be gpulized
458-
# vxc, fxc, kxc = ni.eval_xc_eff(xc_code, rho, deriv, xctype=xctype)[1:]
459-
vxc, fxc, kxc = ni_cpu.eval_xc_eff(xc_code, rho.get(), deriv, xctype=xctype)[1:]
460-
if isinstance(vxc, np.ndarray): vxc = cp.asarray(vxc)
461-
if isinstance(fxc, np.ndarray): fxc = cp.asarray(fxc)
462-
if isinstance(kxc, np.ndarray): kxc = cp.asarray(kxc)
457+
whether_use_gpu = os.environ.get('LIBXC_ON_GPU', '0') == '1'
458+
if not whether_use_gpu:
459+
ni_cpu = numint_cpu()
460+
# TODO: If the libxc is stablized, this should be gpulized
461+
# vxc, fxc, kxc = ni.eval_xc_eff(xc_code, rho, deriv, xctype=xctype)[1:]
462+
vxc, fxc, kxc = ni_cpu.eval_xc_eff(xc_code, rho.get(), deriv, xctype=xctype)[1:]
463+
if isinstance(vxc, np.ndarray): vxc = cp.asarray(vxc)
464+
if isinstance(fxc, np.ndarray): fxc = cp.asarray(fxc)
465+
if isinstance(kxc, np.ndarray): kxc = cp.asarray(kxc)
466+
else:
467+
vxc, fxc, kxc = ni.eval_xc_eff(xc_code, rho, deriv, xctype=xctype)[1:]
463468
else:
464469
vxc, fxc, kxc = ni.eval_xc_eff(xc_code, rho, deriv, xctype=xctype)[1:]
465470
dmvo_mask_a = dmvo[0, mask[:, None], mask]

0 commit comments

Comments
 (0)