gpu4pyscf 1.0.2__tar.gz → 1.2.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {gpu4pyscf-1.0.2/gpu4pyscf.egg-info → gpu4pyscf-1.2.0}/PKG-INFO +1 -1
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/README.md +9 -6
- gpu4pyscf-1.2.0/gpu4pyscf/__config__.py +34 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/__init__.py +3 -3
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/cc/ccsd_incore.py +8 -8
- gpu4pyscf-1.2.0/gpu4pyscf/df/df.py +353 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/df_jk.py +234 -113
- gpu4pyscf-1.2.0/gpu4pyscf/df/grad/jk.py +82 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/grad/rhf.py +58 -73
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/grad/rks.py +6 -7
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/grad/uhf.py +59 -44
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/grad/uks.py +2 -3
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/hessian/rhf.py +73 -81
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/hessian/rks.py +20 -17
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/hessian/uhf.py +60 -73
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/hessian/uks.py +26 -24
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/int3c2e.py +628 -543
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/__init__.py +2 -3
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/gks.py +1 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/libxc.py +14 -28
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/numint.py +629 -523
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/rks.py +27 -17
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/roks.py +24 -20
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/uks.py +23 -17
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/xc_deriv.py +1 -1
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/drivers/dft_driver.py +3 -3
- gpu4pyscf-1.2.0/gpu4pyscf/grad/__init__.py +6 -0
- gpu4pyscf-1.2.0/gpu4pyscf/grad/rhf.py +376 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/grad/rks.py +171 -100
- gpu4pyscf-1.2.0/gpu4pyscf/grad/uhf.py +135 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/grad/uks.py +158 -112
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/gto/mole.py +13 -10
- gpu4pyscf-1.2.0/gpu4pyscf/gto/moleintor.py +482 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/hessian/__init__.py +2 -8
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/hessian/dispersion.py +1 -2
- gpu4pyscf-1.2.0/gpu4pyscf/hessian/rhf.py +804 -0
- gpu4pyscf-1.2.0/gpu4pyscf/hessian/rks.py +702 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/hessian/uhf.py +72 -239
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/hessian/uks.py +106 -132
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/lib/cublas.py +2 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/lib/cupy_helper.py +203 -7
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/lib/cusolver.py +70 -16
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/lib/cutensor.py +21 -21
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/lib/logger.py +3 -3
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/mp/dfmp2.py +48 -18
- gpu4pyscf-1.2.0/gpu4pyscf/pbc/__init__.py +19 -0
- gpu4pyscf-1.2.0/gpu4pyscf/pbc/df/__init__.py +25 -0
- gpu4pyscf-1.2.0/gpu4pyscf/pbc/df/fft.py +272 -0
- gpu4pyscf-1.2.0/gpu4pyscf/pbc/df/fft_jk.py +346 -0
- gpu4pyscf-1.2.0/gpu4pyscf/pbc/dft/__init__.py +31 -0
- gpu4pyscf-1.2.0/gpu4pyscf/pbc/dft/gen_grid.py +79 -0
- gpu4pyscf-1.2.0/gpu4pyscf/pbc/dft/numint.py +433 -0
- gpu4pyscf-1.2.0/gpu4pyscf/pbc/dft/rks.py +228 -0
- gpu4pyscf-1.2.0/gpu4pyscf/pbc/scf/__init__.py +32 -0
- gpu4pyscf-1.2.0/gpu4pyscf/pbc/scf/hf.py +268 -0
- gpu4pyscf-1.2.0/gpu4pyscf/pbc/tools/__init__.py +18 -0
- gpu4pyscf-1.2.0/gpu4pyscf/pbc/tools/pbc.py +103 -0
- gpu4pyscf-1.2.0/gpu4pyscf/pop/esp.py +221 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/properties/shielding.py +25 -31
- gpu4pyscf-1.2.0/gpu4pyscf/qmmm/chelpg.py +127 -0
- gpu4pyscf-1.2.0/gpu4pyscf/qmmm/pbc/__init__.py +16 -0
- gpu4pyscf-1.2.0/gpu4pyscf/qmmm/pbc/itrf.py +1171 -0
- gpu4pyscf-1.2.0/gpu4pyscf/qmmm/pbc/mm_mole.py +425 -0
- gpu4pyscf-1.2.0/gpu4pyscf/qmmm/pbc/tools.py +307 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/scf/_response_functions.py +65 -38
- gpu4pyscf-1.2.0/gpu4pyscf/scf/dispersion.py +22 -0
- gpu4pyscf-1.2.0/gpu4pyscf/scf/hf.py +538 -0
- gpu4pyscf-1.2.0/gpu4pyscf/scf/hf_symm.py +22 -0
- gpu4pyscf-1.2.0/gpu4pyscf/scf/int2c2e.py +141 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/scf/int4c2e.py +310 -4
- gpu4pyscf-1.2.0/gpu4pyscf/scf/j_engine.py +230 -0
- gpu4pyscf-1.2.0/gpu4pyscf/scf/jk.py +777 -0
- gpu4pyscf-1.2.0/gpu4pyscf/scf/rohf.py +229 -0
- gpu4pyscf-1.2.0/gpu4pyscf/scf/soscf.py +704 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/scf/uhf.py +46 -16
- gpu4pyscf-1.2.0/gpu4pyscf/scf/uhf_symm.py +20 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/__init__.py +2 -2
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/_attach_solvent.py +1 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/grad/pcm.py +48 -48
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/grad/smd.py +5 -10
- gpu4pyscf-1.2.0/gpu4pyscf/solvent/hessian/__init__.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/hessian/pcm.py +1 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/pcm.py +4 -1
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/smd.py +6 -2
- gpu4pyscf-1.2.0/gpu4pyscf/tdscf/__init__.py +22 -0
- gpu4pyscf-1.2.0/gpu4pyscf/tdscf/_lr_eig.py +994 -0
- gpu4pyscf-1.2.0/gpu4pyscf/tdscf/_uhf_resp_sf.py +217 -0
- gpu4pyscf-1.2.0/gpu4pyscf/tdscf/rhf.py +368 -0
- gpu4pyscf-1.2.0/gpu4pyscf/tdscf/rks.py +156 -0
- gpu4pyscf-1.2.0/gpu4pyscf/tdscf/uhf.py +777 -0
- gpu4pyscf-1.2.0/gpu4pyscf/tdscf/uks.py +197 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0/gpu4pyscf.egg-info}/PKG-INFO +1 -1
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf.egg-info/SOURCES.txt +34 -1
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf.egg-info/requires.txt +1 -1
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/setup.py +1 -1
- gpu4pyscf-1.0.2/gpu4pyscf/__config__.py +0 -38
- gpu4pyscf-1.0.2/gpu4pyscf/df/df.py +0 -299
- gpu4pyscf-1.0.2/gpu4pyscf/grad/__init__.py +0 -6
- gpu4pyscf-1.0.2/gpu4pyscf/grad/rhf.py +0 -702
- gpu4pyscf-1.0.2/gpu4pyscf/grad/uhf.py +0 -353
- gpu4pyscf-1.0.2/gpu4pyscf/hessian/rhf.py +0 -695
- gpu4pyscf-1.0.2/gpu4pyscf/hessian/rks.py +0 -701
- gpu4pyscf-1.0.2/gpu4pyscf/qmmm/chelpg.py +0 -271
- gpu4pyscf-1.0.2/gpu4pyscf/scf/dispersion.py +0 -32
- gpu4pyscf-1.0.2/gpu4pyscf/scf/hf.py +0 -1063
- gpu4pyscf-1.0.2/gpu4pyscf/scf/rohf.py +0 -76
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/LICENSE +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/MANIFEST.in +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/cc/__init__.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/__init__.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/cderi.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/grad/__init__.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/df/hessian/__init__.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/gen_grid.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/libxc_structs.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/radi.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/dft/xc_alias.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/drivers/__init__.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/drivers/opt_driver.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/fci/__init__.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/fci/direct_spin1.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/grad/dispersion.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/gto/__init__.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/lib/__init__.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/lib/diis.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/lib/utils.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/mp/__init__.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/mp/mp2.py +0 -0
- {gpu4pyscf-1.0.2/gpu4pyscf/solvent/hessian → gpu4pyscf-1.2.0/gpu4pyscf/pop}/__init__.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/properties/__init__.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/properties/ir.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/properties/polarizability.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/qmmm/__init__.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/scf/__init__.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/scf/cphf.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/scf/diis.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/scf/ghf.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/scf/ucphf.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/grad/__init__.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/grad/smd_experiment.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/hessian/smd.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/hessian/smd_experiment.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf/solvent/smd_experiment.py +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf.egg-info/dependency_links.txt +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/gpu4pyscf.egg-info/top_level.txt +0 -0
- {gpu4pyscf-1.0.2 → gpu4pyscf-1.2.0}/setup.cfg +0 -0
|
@@ -43,6 +43,7 @@ Features
|
|
|
43
43
|
- Density fitting scheme and direct SCF scheme;
|
|
44
44
|
- SCF, analytical Gradient, and analytical Hessian calculations for Hartree-Fock and DFT;
|
|
45
45
|
- LDA, GGA, mGGA, hybrid, and range-separated functionals via [libXC](https://gitlab.com/libxc/libxc/-/tree/master/);
|
|
46
|
+
- Spin-conserved and spin-flip TDA and TDDFT for excitated states
|
|
46
47
|
- Geometry optimization and transition state search via [geomeTRIC](https://geometric.readthedocs.io/en/latest/);
|
|
47
48
|
- Dispersion corrections via [DFTD3](https://github.com/dftd3/simple-dftd3) and [DFTD4](https://github.com/dftd4/dftd4);
|
|
48
49
|
- Nonlocal functional correction (vv10) for SCF and gradient;
|
|
@@ -50,7 +51,10 @@ Features
|
|
|
50
51
|
- PCM models, SMD model, their analytical gradients, and semi-analytical Hessian matrix;
|
|
51
52
|
- Unrestricted Hartree-Fock and Unrestricted DFT, gradient, and Hessian;
|
|
52
53
|
- MP2/DF-MP2 and CCSD (experimental);
|
|
53
|
-
- Polarizability, IR, and NMR shielding
|
|
54
|
+
- Polarizability, IR, and NMR shielding (experimental);
|
|
55
|
+
- QM/MM with PBC;
|
|
56
|
+
- CHELPG, ESP, and RESP atomic charge;
|
|
57
|
+
- Multi-GPU for density fitting (experimental)
|
|
54
58
|
|
|
55
59
|
Limitations
|
|
56
60
|
--------
|
|
@@ -58,7 +62,6 @@ Limitations
|
|
|
58
62
|
- Atomic basis up to g orbitals;
|
|
59
63
|
- Auxiliary basis up to i orbitals;
|
|
60
64
|
- Density fitting scheme up to ~168 atoms with def2-tzvpd basis, bounded by CPU memory;
|
|
61
|
-
- Hessian is unavailable for Direct SCF yet;
|
|
62
65
|
- meta-GGA without density laplacian;
|
|
63
66
|
- Double hybrid functionals are not supported;
|
|
64
67
|
|
|
@@ -132,22 +135,22 @@ References
|
|
|
132
135
|
---------
|
|
133
136
|
```
|
|
134
137
|
@misc{li2024introducting,
|
|
135
|
-
title={Introducing GPU-acceleration into the Python-based Simulations of Chemistry Framework},
|
|
138
|
+
title={Introducing GPU-acceleration into the Python-based Simulations of Chemistry Framework},
|
|
136
139
|
author={Rui Li and Qiming Sun and Xing Zhang and Garnet Kin-Lic Chan},
|
|
137
140
|
year={2024},
|
|
138
141
|
eprint={2407.09700},
|
|
139
142
|
archivePrefix={arXiv},
|
|
140
143
|
primaryClass={physics.comp-ph},
|
|
141
|
-
url={https://arxiv.org/abs/2407.09700},
|
|
144
|
+
url={https://arxiv.org/abs/2407.09700},
|
|
142
145
|
}
|
|
143
146
|
|
|
144
147
|
@misc{wu2024enhancing,
|
|
145
|
-
title={Enhancing GPU-acceleration in the Python-based Simulations of Chemistry Framework},
|
|
148
|
+
title={Enhancing GPU-acceleration in the Python-based Simulations of Chemistry Framework},
|
|
146
149
|
author={Xiaojie Wu and Qiming Sun and Zhichen Pu and Tianze Zheng and Wenzhi Ma and Wen Yan and Xia Yu and Zhengxiao Wu and Mian Huo and Xiang Li and Weiluo Ren and Sheng Gong and Yumin Zhang and Weihao Gao},
|
|
147
150
|
year={2024},
|
|
148
151
|
eprint={2404.09452},
|
|
149
152
|
archivePrefix={arXiv},
|
|
150
153
|
primaryClass={physics.comp-ph},
|
|
151
|
-
url={https://arxiv.org/abs/2404.09452},
|
|
154
|
+
url={https://arxiv.org/abs/2404.09452},
|
|
152
155
|
}
|
|
153
156
|
```
|
|
@@ -0,0 +1,34 @@
|
|
|
1
|
+
import cupy
|
|
2
|
+
|
|
3
|
+
_num_devices = cupy.cuda.runtime.getDeviceCount()
|
|
4
|
+
|
|
5
|
+
# TODO: switch to non_blocking stream (currently blocked by libxc)
|
|
6
|
+
_streams = [None] * _num_devices
|
|
7
|
+
for device_id in range(_num_devices):
|
|
8
|
+
with cupy.cuda.Device(device_id):
|
|
9
|
+
_streams[device_id] = cupy.cuda.stream.Stream(non_blocking=False)
|
|
10
|
+
|
|
11
|
+
props = cupy.cuda.runtime.getDeviceProperties(0)
|
|
12
|
+
GB = 1024*1024*1024
|
|
13
|
+
min_ao_blksize = 128
|
|
14
|
+
min_grid_blksize = 128*128
|
|
15
|
+
ao_aligned = 32
|
|
16
|
+
grid_aligned = 256
|
|
17
|
+
|
|
18
|
+
# Use smaller blksize for old gaming GPUs
|
|
19
|
+
if props['totalGlobalMem'] < 16 * GB:
|
|
20
|
+
min_ao_blksize = 64
|
|
21
|
+
min_grid_blksize = 64*64
|
|
22
|
+
|
|
23
|
+
# Use 90% of the global memory for CuPy memory pool
|
|
24
|
+
mem_fraction = 0.9
|
|
25
|
+
cupy.get_default_memory_pool().set_limit(fraction=mem_fraction)
|
|
26
|
+
|
|
27
|
+
# Check P2P data transfer is available
|
|
28
|
+
_p2p_access = True
|
|
29
|
+
if _num_devices > 1:
|
|
30
|
+
for src in range(_num_devices):
|
|
31
|
+
for dst in range(_num_devices):
|
|
32
|
+
if src != dst:
|
|
33
|
+
can_access_peer = cupy.cuda.runtime.deviceCanAccessPeer(src, dst)
|
|
34
|
+
_p2p_access &= can_access_peer
|
|
@@ -1,7 +1,7 @@
|
|
|
1
|
-
|
|
2
|
-
|
|
3
|
-
__version__ = '1.0.2'
|
|
1
|
+
__version__ = '1.2.0'
|
|
4
2
|
|
|
5
3
|
# monkey patch libxc reference due to a bug in nvcc
|
|
6
4
|
from pyscf.dft import libxc
|
|
7
5
|
libxc.__reference__ = 'unable to decode the reference due to https://github.com/NVIDIA/cuda-python/issues/29'
|
|
6
|
+
|
|
7
|
+
from . import lib, grad, hessian, solvent, scf, dft
|
|
@@ -31,7 +31,7 @@ from pyscf.ao2mo import _ao2mo
|
|
|
31
31
|
from pyscf.cc import ccsd
|
|
32
32
|
from pyscf.cc import _ccsd
|
|
33
33
|
from pyscf import __config__
|
|
34
|
-
from gpu4pyscf.scf import
|
|
34
|
+
from gpu4pyscf.scf import int4c2e
|
|
35
35
|
from gpu4pyscf.lib.cupy_helper import load_library
|
|
36
36
|
from gpu4pyscf.lib import logger
|
|
37
37
|
|
|
@@ -171,7 +171,7 @@ def _direct_ovvv_vvvv(mycc, t1, t2):
|
|
|
171
171
|
((mem_avail-Ht2_mem)*.5/8/nao_cart**2)**.5)))
|
|
172
172
|
logger.debug1(mycc, 'blksize %d nao %d', blksize, nao_cart)
|
|
173
173
|
|
|
174
|
-
vhfopt =
|
|
174
|
+
vhfopt = int4c2e._VHFOpt(mycc.mol, 'int2e')
|
|
175
175
|
vhfopt.build(group_size=blksize, diag_block_with_triu=True)
|
|
176
176
|
mol = vhfopt.mol
|
|
177
177
|
|
|
@@ -219,7 +219,7 @@ def _direct_ovvv_vvvv(mycc, t1, t2):
|
|
|
219
219
|
log_qs = vhfopt.log_qs
|
|
220
220
|
cp_idx, cp_jdx = np.tril_indices(len(vhfopt.uniq_l_ctr))
|
|
221
221
|
|
|
222
|
-
if vhfopt.uniq_l_ctr[:,0].max() <=
|
|
222
|
+
if vhfopt.uniq_l_ctr[:,0].max() <= int4c2e.LMAX_ON_GPU:
|
|
223
223
|
# Computing ERIs on GPU
|
|
224
224
|
idx, idy = cupy.tril_indices(nao)
|
|
225
225
|
#eribuf = cupy.empty(blksize**2*nao**2)
|
|
@@ -279,7 +279,7 @@ def _direct_ovvv_vvvv(mycc, t1, t2):
|
|
|
279
279
|
cpj = cp_jdx[cp_ij_id]
|
|
280
280
|
li = vhfopt.uniq_l_ctr[cpi,0]
|
|
281
281
|
lj = vhfopt.uniq_l_ctr[cpj,0]
|
|
282
|
-
if li >
|
|
282
|
+
if li > int4c2e.LMAX_ON_GPU or lj > int4c2e.LMAX_ON_GPU or log_q_ij.size == 0:
|
|
283
283
|
continue
|
|
284
284
|
|
|
285
285
|
ish0 = l_ctr_offsets[cpi]
|
|
@@ -375,7 +375,7 @@ def _fill_eri_block(eri, strides, ao_offsets, vhfopt, group_id):
|
|
|
375
375
|
cpl = cp_jdx[cp_kl_id]
|
|
376
376
|
lk = vhfopt.uniq_l_ctr[cpk,0]
|
|
377
377
|
ll = vhfopt.uniq_l_ctr[cpl,0]
|
|
378
|
-
if lk >
|
|
378
|
+
if lk > int4c2e.LMAX_ON_GPU or ll > int4c2e.LMAX_ON_GPU:
|
|
379
379
|
raise NotImplementedError
|
|
380
380
|
|
|
381
381
|
stream = cupy.cuda.get_current_stream()
|
|
@@ -394,7 +394,7 @@ def _fill_eri_block(eri, strides, ao_offsets, vhfopt, group_id):
|
|
|
394
394
|
cpj = cp_jdx[cp_ij_id]
|
|
395
395
|
li = vhfopt.uniq_l_ctr[cpi,0]
|
|
396
396
|
lj = vhfopt.uniq_l_ctr[cpj,0]
|
|
397
|
-
if li >
|
|
397
|
+
if li > int4c2e.LMAX_ON_GPU or lj > int4c2e.LMAX_ON_GPU or log_q_ij.size == 0:
|
|
398
398
|
continue
|
|
399
399
|
|
|
400
400
|
t0 = time.perf_counter()
|
|
@@ -446,7 +446,7 @@ def _make_eris_incore(mycc, mo_coeff=None):
|
|
|
446
446
|
(mem_avail*.5/8/nao_cart**2)**.5)))
|
|
447
447
|
logger.debug1(mycc, 'blksize %d nao %d', blksize, nao_cart)
|
|
448
448
|
|
|
449
|
-
vhfopt =
|
|
449
|
+
vhfopt = int4c2e._VHFOpt(mycc.mol, 'int2e')
|
|
450
450
|
vhfopt.build(group_size=blksize, diag_block_with_triu=True)
|
|
451
451
|
mol = vhfopt.mol
|
|
452
452
|
mo = vhfopt.coeff.dot(mo_coeff)
|
|
@@ -470,7 +470,7 @@ def _make_eris_incore(mycc, mo_coeff=None):
|
|
|
470
470
|
cpj = cp_jdx[cp_ij_id]
|
|
471
471
|
li = vhfopt.uniq_l_ctr[cpi,0]
|
|
472
472
|
lj = vhfopt.uniq_l_ctr[cpj,0]
|
|
473
|
-
if li >
|
|
473
|
+
if li > int4c2e.LMAX_ON_GPU or lj > int4c2e.LMAX_ON_GPU or log_q_ij.size == 0:
|
|
474
474
|
continue
|
|
475
475
|
|
|
476
476
|
ish0 = l_ctr_offsets[cpi]
|
|
@@ -0,0 +1,353 @@
|
|
|
1
|
+
# Copyright 2023 The GPU4PySCF Authors. All Rights Reserved.
|
|
2
|
+
#
|
|
3
|
+
# This program is free software: you can redistribute it and/or modify
|
|
4
|
+
# it under the terms of the GNU General Public License as published by
|
|
5
|
+
# the Free Software Foundation, either version 3 of the License, or
|
|
6
|
+
# (at your option) any later version.
|
|
7
|
+
#
|
|
8
|
+
# This program is distributed in the hope that it will be useful,
|
|
9
|
+
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
10
|
+
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
11
|
+
# GNU General Public License for more details.
|
|
12
|
+
#
|
|
13
|
+
# You should have received a copy of the GNU General Public License
|
|
14
|
+
# along with this program. If not, see <http://www.gnu.org/licenses/>.
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
import copy
|
|
18
|
+
import cupy
|
|
19
|
+
from concurrent.futures import ThreadPoolExecutor
|
|
20
|
+
import numpy as np
|
|
21
|
+
from cupyx.scipy.linalg import solve_triangular
|
|
22
|
+
from pyscf import lib
|
|
23
|
+
from pyscf.df import df, addons, incore
|
|
24
|
+
from gpu4pyscf.lib.cupy_helper import cholesky, tag_array, get_avail_mem, cart2sph, p2p_transfer
|
|
25
|
+
from gpu4pyscf.df import int3c2e, df_jk
|
|
26
|
+
from gpu4pyscf.lib import logger
|
|
27
|
+
from gpu4pyscf import __config__
|
|
28
|
+
from gpu4pyscf.__config__ import _streams, _num_devices
|
|
29
|
+
|
|
30
|
+
MIN_BLK_SIZE = getattr(__config__, 'min_ao_blksize', 128)
|
|
31
|
+
ALIGNED = getattr(__config__, 'ao_aligned', 32)
|
|
32
|
+
GB = 1024*1024*1024
|
|
33
|
+
|
|
34
|
+
LINEAR_DEP_THR = incore.LINEAR_DEP_THR
|
|
35
|
+
GROUP_SIZE = 256
|
|
36
|
+
|
|
37
|
+
class DF(lib.StreamObject):
|
|
38
|
+
from gpu4pyscf.lib.utils import to_gpu, device
|
|
39
|
+
|
|
40
|
+
_keys = {'intopt', 'mol', 'auxmol', 'use_gpu_memory'}
|
|
41
|
+
|
|
42
|
+
def __init__(self, mol, auxbasis=None):
|
|
43
|
+
self.mol = mol
|
|
44
|
+
self.stdout = mol.stdout
|
|
45
|
+
self.verbose = mol.verbose
|
|
46
|
+
self.max_memory = mol.max_memory
|
|
47
|
+
self.use_gpu_memory = True
|
|
48
|
+
self._auxbasis = auxbasis
|
|
49
|
+
|
|
50
|
+
self.auxmol = None
|
|
51
|
+
self.intopt = None
|
|
52
|
+
self.nao = None
|
|
53
|
+
self.naux = None
|
|
54
|
+
self.cd_low = None
|
|
55
|
+
self._cderi = None
|
|
56
|
+
self._rsh_df = {}
|
|
57
|
+
|
|
58
|
+
@property
|
|
59
|
+
def auxbasis(self):
|
|
60
|
+
return self._auxbasis
|
|
61
|
+
@auxbasis.setter
|
|
62
|
+
def auxbasis(self, x):
|
|
63
|
+
if self._auxbasis != x:
|
|
64
|
+
self.reset()
|
|
65
|
+
self._auxbasis = x
|
|
66
|
+
|
|
67
|
+
def to_cpu(self):
|
|
68
|
+
from gpu4pyscf.lib.utils import to_cpu
|
|
69
|
+
obj = to_cpu(self)
|
|
70
|
+
return obj.reset()
|
|
71
|
+
|
|
72
|
+
def build(self, direct_scf_tol=1e-14, omega=None):
|
|
73
|
+
mol = self.mol
|
|
74
|
+
auxmol = self.auxmol
|
|
75
|
+
self.nao = mol.nao
|
|
76
|
+
log = logger.new_logger(mol, mol.verbose)
|
|
77
|
+
t0 = log.init_timer()
|
|
78
|
+
if auxmol is None:
|
|
79
|
+
self.auxmol = auxmol = addons.make_auxmol(mol, self.auxbasis)
|
|
80
|
+
|
|
81
|
+
if omega and omega > 1e-10:
|
|
82
|
+
with auxmol.with_range_coulomb(omega):
|
|
83
|
+
j2c_cpu = auxmol.intor('int2c2e', hermi=1)
|
|
84
|
+
else:
|
|
85
|
+
j2c_cpu = auxmol.intor('int2c2e', hermi=1)
|
|
86
|
+
j2c = cupy.asarray(j2c_cpu, order='C')
|
|
87
|
+
t0 = log.timer_debug1('2c2e', *t0)
|
|
88
|
+
intopt = int3c2e.VHFOpt(mol, auxmol, 'int2e')
|
|
89
|
+
intopt.build(direct_scf_tol, diag_block_with_triu=False, aosym=True,
|
|
90
|
+
group_size=GROUP_SIZE, group_size_aux=GROUP_SIZE)
|
|
91
|
+
log.timer_debug1('prepare intopt', *t0)
|
|
92
|
+
self.j2c = j2c.copy()
|
|
93
|
+
|
|
94
|
+
j2c = intopt.sort_orbitals(j2c, aux_axis=[0,1])
|
|
95
|
+
try:
|
|
96
|
+
self.cd_low = cholesky(j2c)
|
|
97
|
+
self.cd_low = tag_array(self.cd_low, tag='cd')
|
|
98
|
+
except Exception:
|
|
99
|
+
w, v = cupy.linalg.eigh(j2c)
|
|
100
|
+
idx = w > LINEAR_DEP_THR
|
|
101
|
+
self.cd_low = (v[:,idx] / cupy.sqrt(w[idx]))
|
|
102
|
+
self.cd_low = tag_array(self.cd_low, tag='eig')
|
|
103
|
+
|
|
104
|
+
v = w = None
|
|
105
|
+
naux = self.naux = self.cd_low.shape[1]
|
|
106
|
+
log.debug('size of aux basis %d', naux)
|
|
107
|
+
|
|
108
|
+
self._cderi = cholesky_eri_gpu(intopt, mol, auxmol, self.cd_low,
|
|
109
|
+
omega=omega, use_gpu_memory=self.use_gpu_memory)
|
|
110
|
+
log.timer_debug1('cholesky_eri', *t0)
|
|
111
|
+
self.intopt = intopt
|
|
112
|
+
return self
|
|
113
|
+
|
|
114
|
+
def get_jk(self, dm, hermi=1, with_j=True, with_k=True,
|
|
115
|
+
direct_scf_tol=getattr(__config__, 'scf_hf_SCF_direct_scf_tol', 1e-13),
|
|
116
|
+
omega=None):
|
|
117
|
+
if omega is None:
|
|
118
|
+
return df_jk.get_jk(self, dm, hermi, with_j, with_k, direct_scf_tol)
|
|
119
|
+
assert omega >= 0.0
|
|
120
|
+
|
|
121
|
+
# A temporary treatment for RSH-DF integrals
|
|
122
|
+
key = '%.6f' % omega
|
|
123
|
+
if key in self._rsh_df:
|
|
124
|
+
rsh_df = self._rsh_df[key]
|
|
125
|
+
else:
|
|
126
|
+
rsh_df = self._rsh_df[key] = self.copy().reset()
|
|
127
|
+
logger.info(self, 'Create RSH-DF object %s for omega=%s', rsh_df, omega)
|
|
128
|
+
|
|
129
|
+
return df_jk.get_jk(rsh_df, dm, hermi, with_j, with_k, direct_scf_tol, omega=omega)
|
|
130
|
+
|
|
131
|
+
def get_blksize(self, extra=0, nao=None):
|
|
132
|
+
'''
|
|
133
|
+
extra for pre-calculated space for other variables
|
|
134
|
+
'''
|
|
135
|
+
if nao is None: nao = self.nao
|
|
136
|
+
mem_avail = get_avail_mem()
|
|
137
|
+
blksize = int(mem_avail*0.2/8/(nao*nao + extra) / ALIGNED) * ALIGNED
|
|
138
|
+
blksize = min(blksize, MIN_BLK_SIZE)
|
|
139
|
+
log = logger.new_logger(self.mol, self.mol.verbose)
|
|
140
|
+
device_id = cupy.cuda.Device().id
|
|
141
|
+
log.debug(f"{mem_avail/1e9:.3f} GB memory available on Device {device_id}, block size = {blksize}")
|
|
142
|
+
if blksize < ALIGNED:
|
|
143
|
+
raise RuntimeError("Not enough GPU memory")
|
|
144
|
+
return blksize
|
|
145
|
+
|
|
146
|
+
def loop(self, blksize=None, unpack=True):
|
|
147
|
+
''' loop over cderi for the current device
|
|
148
|
+
and unpack the CDERI in (Lij) format
|
|
149
|
+
'''
|
|
150
|
+
device_id = cupy.cuda.Device().id
|
|
151
|
+
cderi_sparse = self._cderi[device_id]
|
|
152
|
+
if blksize is None:
|
|
153
|
+
blksize = self.get_blksize()
|
|
154
|
+
nao = self.nao
|
|
155
|
+
naux_slice = cderi_sparse.shape[0]
|
|
156
|
+
rows = self.intopt.cderi_row
|
|
157
|
+
cols = self.intopt.cderi_col
|
|
158
|
+
buf_prefetch = None
|
|
159
|
+
buf_cderi = cupy.zeros([blksize,nao,nao])
|
|
160
|
+
for p0, p1 in lib.prange(0, naux_slice, blksize):
|
|
161
|
+
p2 = min(naux_slice, p1+blksize)
|
|
162
|
+
if isinstance(cderi_sparse, cupy.ndarray):
|
|
163
|
+
buf = cderi_sparse[p0:p1,:]
|
|
164
|
+
if isinstance(cderi_sparse, np.ndarray):
|
|
165
|
+
# first block
|
|
166
|
+
if buf_prefetch is None:
|
|
167
|
+
buf = cupy.asarray(cderi_sparse[p0:p1,:])
|
|
168
|
+
buf_prefetch = cupy.empty([p2-p1,cderi_sparse.shape[1]])
|
|
169
|
+
if isinstance(cderi_sparse, np.ndarray) and p1 < p2:
|
|
170
|
+
buf_prefetch.set(cderi_sparse[p1:p2,:])
|
|
171
|
+
if unpack:
|
|
172
|
+
buf_cderi[:p1-p0,rows,cols] = buf
|
|
173
|
+
buf_cderi[:p1-p0,cols,rows] = buf
|
|
174
|
+
buf2 = buf_cderi[:p1-p0]
|
|
175
|
+
else:
|
|
176
|
+
buf2 = None
|
|
177
|
+
yield buf2, buf.T
|
|
178
|
+
if isinstance(cderi_sparse, np.ndarray):
|
|
179
|
+
cupy.cuda.Device().synchronize()
|
|
180
|
+
|
|
181
|
+
if buf_prefetch is not None:
|
|
182
|
+
buf = buf_prefetch
|
|
183
|
+
|
|
184
|
+
def reset(self, mol=None):
|
|
185
|
+
'''Reset mol and clean up relevant attributes for scanner mode'''
|
|
186
|
+
if mol is not None:
|
|
187
|
+
self.mol = mol
|
|
188
|
+
self.auxmol = None
|
|
189
|
+
self._cderi = None
|
|
190
|
+
self._vjopt = None
|
|
191
|
+
self._rsh_df = {}
|
|
192
|
+
self.intopt = None
|
|
193
|
+
self.nao = None
|
|
194
|
+
self.naux = None
|
|
195
|
+
self.cd_low = None
|
|
196
|
+
return self
|
|
197
|
+
|
|
198
|
+
get_ao_eri = get_eri = NotImplemented
|
|
199
|
+
get_mo_eri = ao2mo = NotImplemented
|
|
200
|
+
|
|
201
|
+
def cholesky_eri_gpu(intopt, mol, auxmol, cd_low,
|
|
202
|
+
omega=None, sr_only=False, use_gpu_memory=True):
|
|
203
|
+
'''
|
|
204
|
+
Returns:
|
|
205
|
+
2D array of (naux,nao*(nao+1)/2) in C-contiguous
|
|
206
|
+
'''
|
|
207
|
+
naux = cd_low.shape[1]
|
|
208
|
+
npairs = len(intopt.cderi_row)
|
|
209
|
+
log = logger.new_logger(mol, mol.verbose)
|
|
210
|
+
|
|
211
|
+
# Available memory on Device 0.
|
|
212
|
+
avail_mem = get_avail_mem()
|
|
213
|
+
|
|
214
|
+
if use_gpu_memory:
|
|
215
|
+
# CDERI will be equally distributed to the devices
|
|
216
|
+
# Other devices usually have more memory available than Device 0
|
|
217
|
+
# CDERI will use up to 40% of the available memory
|
|
218
|
+
use_gpu_memory = naux * npairs * 8 < 0.4 * avail_mem * _num_devices
|
|
219
|
+
|
|
220
|
+
if use_gpu_memory:
|
|
221
|
+
log.debug("Saving CDERI on GPU")
|
|
222
|
+
else:
|
|
223
|
+
log.debug("Saving CDERI on CPU")
|
|
224
|
+
|
|
225
|
+
_cderi = {}
|
|
226
|
+
blksize = (naux + _num_devices - 1) // _num_devices
|
|
227
|
+
for device_id, (p0,p1) in enumerate(lib.prange(0, naux, blksize)):
|
|
228
|
+
if use_gpu_memory:
|
|
229
|
+
with cupy.cuda.Device(device_id), _streams[device_id]:
|
|
230
|
+
_cderi[device_id] = cupy.empty([p1-p0, npairs])
|
|
231
|
+
log.debug(f"CDERI size {_cderi[device_id].nbytes/GB:.3f} on Device {device_id}")
|
|
232
|
+
else:
|
|
233
|
+
mem = cupy.cuda.alloc_pinned_memory((p1-p0) * npairs * 8)
|
|
234
|
+
cderi_blk = np.ndarray([p1-p0, npairs], dtype=np.float64, order='C', buffer=mem)
|
|
235
|
+
_cderi[device_id] = cderi_blk
|
|
236
|
+
|
|
237
|
+
npairs_per_ctr = [len(intopt.ao_pairs_row[cp_ij_id]) for cp_ij_id in range(len(intopt.log_qs))]
|
|
238
|
+
|
|
239
|
+
npairs_per_ctr = np.array(npairs_per_ctr)
|
|
240
|
+
total_task_list = np.argsort(npairs_per_ctr)
|
|
241
|
+
task_list_per_device = []
|
|
242
|
+
for device_id in range(_num_devices):
|
|
243
|
+
task_list_per_device.append(total_task_list[device_id::_num_devices])
|
|
244
|
+
|
|
245
|
+
cd_low_f = cupy.array(cd_low, order='F', copy=False)
|
|
246
|
+
cd_low_f = tag_array(cd_low_f, tag=cd_low.tag)
|
|
247
|
+
|
|
248
|
+
cupy.cuda.get_current_stream().synchronize()
|
|
249
|
+
futures = []
|
|
250
|
+
with ThreadPoolExecutor(max_workers=_num_devices) as executor:
|
|
251
|
+
for device_id in range(_num_devices):
|
|
252
|
+
task_list = task_list_per_device[device_id]
|
|
253
|
+
future = executor.submit(_cderi_task, intopt, cd_low_f, task_list, _cderi,
|
|
254
|
+
omega=omega, sr_only=sr_only, device_id=device_id)
|
|
255
|
+
futures.append(future)
|
|
256
|
+
|
|
257
|
+
for future in futures:
|
|
258
|
+
future.result()
|
|
259
|
+
|
|
260
|
+
if not use_gpu_memory:
|
|
261
|
+
cupy.cuda.Device().synchronize()
|
|
262
|
+
|
|
263
|
+
return _cderi
|
|
264
|
+
|
|
265
|
+
def _cderi_task(intopt, cd_low, task_list, _cderi, omega=None, sr_only=False, device_id=0):
|
|
266
|
+
''' Execute CDERI tasks on one device
|
|
267
|
+
'''
|
|
268
|
+
nq = len(intopt.log_qs)
|
|
269
|
+
mol = intopt.mol
|
|
270
|
+
naux = cd_low.shape[1]
|
|
271
|
+
naoaux = cd_low.shape[0]
|
|
272
|
+
npairs = [len(intopt.ao_pairs_row[cp_ij]) for cp_ij in range(len(intopt.log_qs))]
|
|
273
|
+
pairs_loc = np.append(0, np.cumsum(npairs))
|
|
274
|
+
blksize = (naux + _num_devices - 1) // _num_devices
|
|
275
|
+
with cupy.cuda.Device(device_id), _streams[device_id]:
|
|
276
|
+
log = logger.new_logger(mol, mol.verbose)
|
|
277
|
+
t1 = log.init_timer()
|
|
278
|
+
cd_low_tag = cd_low.tag
|
|
279
|
+
cd_low = cupy.asarray(cd_low)
|
|
280
|
+
|
|
281
|
+
cart_ao_loc = intopt.cart_ao_loc
|
|
282
|
+
aux_ao_loc = intopt.aux_ao_loc
|
|
283
|
+
ao_loc = intopt.ao_loc
|
|
284
|
+
for cp_ij_id in task_list:
|
|
285
|
+
cpi = intopt.cp_idx[cp_ij_id]
|
|
286
|
+
cpj = intopt.cp_jdx[cp_ij_id]
|
|
287
|
+
li = intopt.angular[cpi]
|
|
288
|
+
lj = intopt.angular[cpj]
|
|
289
|
+
i0, i1 = cart_ao_loc[cpi], cart_ao_loc[cpi+1]
|
|
290
|
+
j0, j1 = cart_ao_loc[cpj], cart_ao_loc[cpj+1]
|
|
291
|
+
ni = i1 - i0
|
|
292
|
+
nj = j1 - j0
|
|
293
|
+
if sr_only:
|
|
294
|
+
# TODO: in-place implementation or short-range kernel
|
|
295
|
+
ints_slices = cupy.zeros([naoaux, nj, ni], order='C')
|
|
296
|
+
for cp_kl_id, _ in enumerate(intopt.aux_log_qs):
|
|
297
|
+
k0 = aux_ao_loc[cp_kl_id]
|
|
298
|
+
k1 = aux_ao_loc[cp_kl_id+1]
|
|
299
|
+
int3c2e.get_int3c2e_slice(intopt, cp_ij_id, cp_kl_id, out=ints_slices[k0:k1])
|
|
300
|
+
if omega is not None:
|
|
301
|
+
ints_slices_lr = cupy.zeros([naoaux, nj, ni], order='C')
|
|
302
|
+
for cp_kl_id, _ in enumerate(intopt.aux_log_qs):
|
|
303
|
+
k0 = aux_ao_loc[cp_kl_id]
|
|
304
|
+
k1 = aux_ao_loc[cp_kl_id+1]
|
|
305
|
+
int3c2e.get_int3c2e_slice(intopt, cp_ij_id, cp_kl_id, out=ints_slices[k0:k1], omega=omega)
|
|
306
|
+
ints_slices -= ints_slices_lr
|
|
307
|
+
else:
|
|
308
|
+
# Initialization is required due to cutensor operations later
|
|
309
|
+
ints_slices = cupy.zeros([naoaux, nj, ni], order='C')
|
|
310
|
+
for cp_kl_id, _ in enumerate(intopt.aux_log_qs):
|
|
311
|
+
k0 = aux_ao_loc[cp_kl_id]
|
|
312
|
+
k1 = aux_ao_loc[cp_kl_id+1]
|
|
313
|
+
|
|
314
|
+
int3c2e.get_int3c2e_slice(intopt, cp_ij_id, cp_kl_id, out=ints_slices[k0:k1], omega=omega)
|
|
315
|
+
if lj>1 and not mol.cart: ints_slices = cart2sph(ints_slices, axis=1, ang=lj)
|
|
316
|
+
if li>1 and not mol.cart: ints_slices = cart2sph(ints_slices, axis=2, ang=li)
|
|
317
|
+
|
|
318
|
+
i0, i1 = ao_loc[cpi], ao_loc[cpi+1]
|
|
319
|
+
j0, j1 = ao_loc[cpj], ao_loc[cpj+1]
|
|
320
|
+
|
|
321
|
+
row = intopt.ao_pairs_row[cp_ij_id] - i0
|
|
322
|
+
col = intopt.ao_pairs_col[cp_ij_id] - j0
|
|
323
|
+
|
|
324
|
+
ints_slices_f= cupy.empty([naoaux,len(row)], order='F')
|
|
325
|
+
ints_slices_f[:] = ints_slices[:,col,row]
|
|
326
|
+
ints_slices = None
|
|
327
|
+
if cd_low_tag == 'eig':
|
|
328
|
+
cderi_block = cupy.dot(cd_low.T, ints_slices_f)
|
|
329
|
+
ints_slices = None
|
|
330
|
+
elif cd_low_tag == 'cd':
|
|
331
|
+
cderi_block = solve_triangular(cd_low, ints_slices_f, lower=True, overwrite_b=True)
|
|
332
|
+
else:
|
|
333
|
+
RuntimeError('Tag is not found in lower triangular matrix.')
|
|
334
|
+
t1 = log.timer_debug1(f'solve {cp_ij_id} / {nq} on Device {device_id}', *t1)
|
|
335
|
+
|
|
336
|
+
# TODO:
|
|
337
|
+
# 1) async data transfer
|
|
338
|
+
# 2) auxiliary basis in the last dimension
|
|
339
|
+
|
|
340
|
+
# if CDERI is saved on CPU
|
|
341
|
+
ij0 = pairs_loc[cp_ij_id]
|
|
342
|
+
ij1 = pairs_loc[cp_ij_id+1]
|
|
343
|
+
if isinstance(_cderi[0], np.ndarray):
|
|
344
|
+
for slice_id, (p0,p1) in enumerate(lib.prange(0, naux, blksize)):
|
|
345
|
+
for i in range(p0,p1):
|
|
346
|
+
cderi_block[i].get(out=_cderi[slice_id][i-p0,ij0:ij1])
|
|
347
|
+
else:
|
|
348
|
+
# Copy data to other Devices
|
|
349
|
+
for slice_id, (p0,p1) in enumerate(lib.prange(0, naux, blksize)):
|
|
350
|
+
#_cderi[slice_id][:,ij0:ij1] = cderi_block[p0:p1]
|
|
351
|
+
p2p_transfer(_cderi[slice_id][:,ij0:ij1], cderi_block[p0:p1])
|
|
352
|
+
t1 = log.timer_debug1(f'transfer data for {cp_ij_id} / {nq} on Device {device_id}', *t1)
|
|
353
|
+
return
|