cleanup cp vs np usage, early return with warp-divergence in ft_ao.cu

gpu4pyscf/df/tests/test_df_hessian.py

@@ -358,7 +358,7 @@ def test_unstable_j2c(self):
         mo_occ = mf.mo_occ
         test_hessian_round2 = hobj.partial_hess_elec(mo_energy, mo_coeff, mo_occ)
 
-        assert np.max(np.abs(test_hessian_round1 - test_hessian_round2)) < 2e-7
+        assert abs(test_hessian_round1 - test_hessian_round2).max().item() < 2e-7
 
 if __name__ == "__main__":
     print("Full Tests for DF Hessian")

gpu4pyscf/df/tests/test_df_rks_grad.py

@@ -94,8 +94,8 @@ def _check_grad(mol, grid_response=False, xc=xc0, disp=disp0, tol=1e-6):
     grad_fd = np.array(grad_fd).reshape(-1,3)
     print('finite difference gradient:')
     print(grad_fd)
-    print('difference between analytical and finite difference gradient:', cupy.linalg.norm(g_analy - grad_fd))
-    assert(cupy.linalg.norm(g_analy - grad_fd) < tol)
+    print('difference between analytical and finite difference gradient:', np.linalg.norm(g_analy - grad_fd))
+    assert(np.linalg.norm(g_analy - grad_fd) < tol)
 
 def _vs_cpu(mol, grid_response=False, xc=xc0, disp=disp0, tol=1e-9):
     mf = rks.RKS(mol, xc=xc).density_fit(auxbasis=auxbasis0)

gpu4pyscf/df/tests/test_df_uhf.py

@@ -81,8 +81,8 @@ def _check_grad(mol, tol=1e-5, disp=None):
     grad_fd = np.array(grad_fd).reshape(-1,3)
     print('finite difference gradient:')
     print(grad_fd)
-    print('difference between analytical and finite difference gradient:', cupy.linalg.norm(g_analy - grad_fd))
-    assert(cupy.linalg.norm(g_analy - grad_fd) < tol)
+    print('difference between analytical and finite difference gradient:', np.linalg.norm(g_analy - grad_fd))
+    assert(np.linalg.norm(g_analy - grad_fd) < tol)
 
 class KnownValues(unittest.TestCase):
     '''

gpu4pyscf/dft/tests/test_numint.py

@@ -272,7 +272,8 @@ def test_sparse_index(self):
                     i1 = min(i0+numint.MIN_BLK_SIZE, ngrids)
                     ref = numint._sparse_index(
                         opt._sorted_mol, grids.coords[i0:i1], opt.l_ctr_offsets, ao_loc, opt)
-                    assert all(np.array_equal(r, x) for r, x in zip(ref[1:], dat[i][1:]))
+                    assert all(r.shape == x.shape and bool((r == x).all())
+                               for r, x in zip(ref[1:], dat[i][1:]))
 
     def test_scale_ao(self):
         ao = cupy.random.rand(1, 3, 256)

gpu4pyscf/lib/cupy_helper.py

@@ -1265,7 +1265,7 @@ def batched_vec3_norm2(batched_vec3):
     '''
     einsum('gx,gx->g', vec3, vec3) for the (N,3)-array vec3
     '''
-    assert type(batched_vec3) is cupy.ndarray
+    assert isinstance(batched_vec3, cupy.ndarray)
     assert batched_vec3.dtype == cupy.float64
     assert batched_vec3.ndim == 2
     assert batched_vec3.shape[0] == 3 or batched_vec3.shape[1] == 3

gpu4pyscf/pbc/df/ft_ao.py

@@ -361,7 +361,7 @@ def ft_evaluator(self, batch_size=None, compressing=True, cart=None,
             dims, tmp = np.empty_like(dims), dims
             dims[cell.sorted_idx] = tmp
             ao_loc = cp.asarray(np.append(0, np.cumsum(dims.ravel())))
-            ao_loc = np.append(ao_loc[cell.sorted_idx], nao)
+            ao_loc = cp.append(ao_loc[cell.sorted_idx], nao)
         ao_loc = cp.asarray(ao_loc, dtype=np.int32)
 
         if batch_size is None:

gpu4pyscf/pbc/df/int3c2e.py

@@ -128,7 +128,7 @@ def sr_aux_e2(cell, auxcell, omega, kpts=None, bvk_kmesh=None, j_only=False):
         out = cp.empty((nkpts,nkpts,naux,nao,nao), dtype=np.complex128)
         kk_conserv = double_translation_indices(int3c2e_opt.bvk_kmesh)
         for k in range(nkpts):
-            ki_idx, kj_idx = np.where(kk_conserv == k)
+            ki_idx, kj_idx = cp.where(kk_conserv == k)
             out[k] = _unpack_cderi_v2(j3c[k], pair_address, kj_idx,
                                       conj_mapping, expLk, nao, axis)
         j3c = None

gpu4pyscf/pbc/df/rsdf_builder.py

@@ -404,8 +404,14 @@ def proc():
                 ctypes.c_int(naux), ctypes.c_int(nao_pairs),
                 ctypes.c_int(p0), ctypes.c_int(p1))
             j3c = None
-            if num_devices > 1:
-                stream.synchronize()
+            # store_col_segment enqueues its device->host write asynchronously
+            # and does not itself block. The next reader of `cderi` is a
+            # plain host-side read (mydf.loop()) with no ordering relationship
+            # to this stream, so this segment must be drained here every
+            # time -- not only when num_devices > 1. Skipping this on a
+            # single device happened to be safe on this stream/allocator
+            # combination but is not guaranteed in general.
+            stream.synchronize()
             t1 = log.timer_debug1(f'store int3c2e on Device {device_id}', *t1)
 
     multi_gpu.run(proc, non_blocking=True)
@@ -572,8 +578,9 @@ def proc():
                 if err != 0:
                     raise RuntimeError('store_col_segment kernel failed')
             j3c = None
-            if num_devices > 1:
-                stream.synchronize()
+            # See the matching comment in compressed_cderi_j_only: this must
+            # run unconditionally, not only when num_devices > 1.
+            stream.synchronize()
             t1 = log.timer_debug1(f'store int3c2e on Device {device_id}', *t1)
 
     multi_gpu.run(proc, non_blocking=True)

gpu4pyscf/pbc/df/tests/test_pbc_int3c2e.py

@@ -325,7 +325,7 @@ def test_contract_dm_kpts():
     np.random.seed(9)
     auxvec = np.random.rand(auxcell.nao)
     vj = opt.contract_auxvec(opt.auxcell.apply_C_dot(auxvec), kpts=kpts)
-    ref = cp.einsum('kpqr,r->kpq', j3c, auxvec)
+    ref = cp.einsum('kpqr,r->kpq', j3c, cp.asarray(auxvec))  # auxvec is host data
     assert abs(vj - ref).max() < 1e-10
 
 def test_int3c2e_batch_evaluation():

gpu4pyscf/pbc/dft/tests/test_pbc_numint.py

@@ -396,7 +396,7 @@ def test_uniform_grid_division_mode(self):
         test_coords = []
         for frag_grids in grids.loop_grids():
             test_coords.append(frag_grids.coords)
-            assert np.max(np.abs(frag_grids.weights - ref_weight)) < 1e-14
+            assert cp.max(cp.abs(frag_grids.weights - ref_weight)) < 1e-14
         test_coords = cp.vstack(test_coords).get()
 
         ref_coords = ref_coords[np.lexsort((ref_coords[:, 2], ref_coords[:, 1], ref_coords[:, 0])), :]

gpu4pyscf/pbc/scf/tests/test_pbc_scf_smearing.py

@@ -40,7 +40,7 @@ def test_krhf_smearing(self):
         mf = cell.KRHF(kpts=cell.make_kpts([2,1,1])).to_gpu()
         mf = mf.smearing(0.1, 'fermi')
         nkpts = len(mf.kpts)
-        mo_energy_kpts = cp.array([cp.arange(nao)*.2+cp.cos(i+.5)*.1 for i in range(nkpts)])
+        mo_energy_kpts = cp.array([cp.arange(nao)*.2+np.cos(i+.5)*.1 for i in range(nkpts)])
         mf.get_occ(mo_energy_kpts)
         self.assertAlmostEqual(mf.entropy, 6.1656394960533021/2, 9)
 
@@ -56,7 +56,7 @@ def test_kuhf_smearing(self):
         mf = cell.KUHF(kpts=cell.make_kpts([2,1,1])).to_gpu()
         mf = mf.smearing(0.1, 'fermi')
         nkpts = len(mf.kpts)
-        mo_energy_kpts = cp.array([cp.arange(nao)*.2+cp.cos(i+.5)*.1 for i in range(nkpts)])
+        mo_energy_kpts = cp.array([cp.arange(nao)*.2+np.cos(i+.5)*.1 for i in range(nkpts)])
         mo_energy_kpts = cp.array([mo_energy_kpts, mo_energy_kpts+cp.cos(mo_energy_kpts)*.02])
         mf.get_occ(mo_energy_kpts)
         self.assertAlmostEqual(mf.entropy, 6.1803390081500869/2, 9)

gpu4pyscf/pbc/tools/k2gamma.py

@@ -82,7 +82,8 @@ def double_translation_indices(kmesh):
     tz = cp.array(translation_map(kmesh[2]), dtype=np.int32)
     idx = cp.ravel_multi_index([tx[:,None,None,:,None,None],
                                 ty[None,:,None,None,:,None],
-                                tz[None,None,:,None,None,:]], kmesh)
+                                tz[None,None,:,None,None,:]],
+                                tuple(int(n) for n in kmesh))
     nk = np.prod(kmesh)
     return idx.reshape(nk, nk)
 

gpu4pyscf/qmmm/external_field.py

@@ -53,15 +53,15 @@ def __init__(self, method, electric_field=None, origin=None):
 
         if electric_field is not None:
             electric_field = cp.asarray(electric_field)
-            assert type(electric_field) is cp.ndarray
+            assert isinstance(electric_field, cp.ndarray)
             assert electric_field.shape == (3,)
         self.electric_field = electric_field
 
         if origin is None:
             origin = np.zeros(3)
         else:
             origin = cp.asarray(origin).get()
-            assert type(origin) is np.ndarray
+            assert isinstance(origin, np.ndarray)
             assert origin.shape == (3,)
         self.origin = origin