PBC Coulomb matrix using MD J engine (#547)

* Reduce shared memory requirements in MD J engine

* Bugfixes

* Optimize md_j by caching Rt2 indices

* Add pbc_md_contract_j cuda kernel

* Update j_engine

* Bugfix for multi-dm md-j

* Off-diagonal elements were incorrectly scaled by 0.5

* Debug pbc md-j

* Fix small errors in q_cond

* almost work

* dtype changes cause rsjk crash

* Fix kpts Coulomb matrix; Add tests

* Lint

* j_engine_3c2e bug due to syncthreads

* pickle serialization error

Author: sunqm

gpu4pyscf/df/j_engine_3c2e.py

@@ -66,12 +66,13 @@ def build(self, cutoff=1e-14):
         else:
             q_cond = np.empty((nbas,nbas))
             intor = prim_mol._add_suffix('int2e')
-            _vhf.libcvhf.CVHFnr_int2e_q_cond(
-                getattr(_vhf.libcvhf, intor), lib.c_null_ptr(),
-                q_cond.ctypes, ao_loc.ctypes,
-                prim_mol._atm.ctypes, ctypes.c_int(prim_mol.natm),
-                prim_mol._bas.ctypes, ctypes.c_int(prim_mol.nbas),
-                prim_mol._env.ctypes)
+            with prim_mol.with_integral_screen(1e-26):
+                _vhf.libcvhf.CVHFnr_int2e_q_cond(
+                    getattr(_vhf.libcvhf, intor), lib.c_null_ptr(),
+                    q_cond.ctypes, ao_loc.ctypes,
+                    prim_mol._atm.ctypes, ctypes.c_int(prim_mol.natm),
+                    prim_mol._bas.ctypes, ctypes.c_int(prim_mol.nbas),
+                    prim_mol._env.ctypes)
             q_cond = np.log(q_cond + 1e-300).astype(np.float32)
         log.timer('Initialize q_cond', *cput0)
 

gpu4pyscf/df/tests/test_df_int3c2e.py

@@ -159,7 +159,7 @@ def test_contract_int3c2e():
         H   0.757    4.   -0.4696
         C   1.      1.    0.
         ''',
-        basis=('ccpvtz', [[1, [3.7, 1]], [1, [2., 1]], [1, [.8, 1]]])
+        basis=('ccpvtz', [[1, [3.7, 1, .1]], [1, [2., .5, .3]], [1, [.8, .5, .8]]])
     )
     auxmol = mol.copy()
     auxmol.basis = ('weigend', [[3, [2, 1, .5], [1, .2, 1]]])

gpu4pyscf/hessian/tests/test_rhf_hessian.py

@@ -180,7 +180,6 @@ def test_jk_mix(self):
         C    H
             0.4        1
             ''',
-            output = '/dev/null'
         )
         nao = mol1.nao
         mo_coeff = cupy.random.rand(nao, nao)
@@ -197,7 +196,6 @@ def test_jk_mix(self):
         vk_cpu = (mo_coeff.T @ vk @ mocc).reshape(1,-1)
         assert cupy.linalg.norm(vj_cpu - vj_mo) < 1e-5
         assert cupy.linalg.norm(vk_cpu - vk_mo) < 1e-5
-        mol1.stdout.close()
 
     @unittest.skipIf(num_devices > 1, '')
     def test_ecp_hess(self):

gpu4pyscf/lib/cupy_helper.py

@@ -520,7 +520,7 @@ def takebak(out, a, indices, axis=-1):
 
 def transpose_sum(a, stream=None):
     '''
-    return a + a.transpose(0,2,1)
+    perform a + a.transpose(0,2,1) inplace
     '''
     assert isinstance(a, cupy.ndarray)
     assert a.flags.c_contiguous

gpu4pyscf/lib/gvhf-md/CMakeLists.txt

@@ -4,6 +4,7 @@ add_library(gvhf_md SHARED
   md_contract_j.cu unrolled_md_j.cu unrolled_md_j_4dm.cu
   md_indices.cu md_j_driver.cu md_pairdata.c 
   contract_int3c2e.cu
+  pbc_md_contract_j.cu
 )
 
 #option(BUILD_SHARED_LIBS "build shared libraries" 1)

gpu4pyscf/lib/gvhf-md/contract_int3c2e.cu

@@ -461,8 +461,6 @@ void unrolled_contract_int3c2e(Int3c2eEnvVars envs, JKMatrix jk, BDiv3c2eBounds
                 double fac = ck/(aij*ak*sqrt(aij+ak));
                 if (pair_ij >= shl_pair1) {
                     fac = 0;
-                } else if (ish == jsh) {
-                    fac *= .5;
                 }
                 boys_fn(gamma_inc, theta, rr, jk.omega, fac, order, sp_id, nsp_per_block);
                 Rt[0] = gamma_inc[sp_id+order*nsp_per_block];
@@ -529,10 +527,11 @@ void unrolled_contract_int3c2e(Int3c2eEnvVars envs, JKMatrix jk, BDiv3c2eBounds
         __shared__ typename BlockReduceT::TempStorage temp_storage;
 #pragma unroll
         for (int k = 0; k < nfk; k++) {
-            vj_aux[k] = BlockReduceT(temp_storage).Sum(vj_aux[k]);
+            double sum_jaux = BlockReduceT(temp_storage).Sum(vj_aux[k]);
             if (thread_id == 0) {
-                atomicAdd(vj+k, vj_aux[k]);
+                atomicAdd(vj+k, sum_jaux);
             }
+            __syncthreads();
         }
     }
 }