For very auxiliary diffuse basis, small errors in pbc-rsdf_builder 2c2e integrals causes linear dependency issue

Author: sunqm

gpu4pyscf/lib/pbc/fill_int3c2e_v2.cu

@@ -26,9 +26,9 @@
 #define GOUT_WIDTH      54
 #define PAGE_SIZE       30
 #define REMOTE_THRESHOLD 50
-#define PAGES_PER_BLOCK  524288
-// approximately, 4000 images in each ijk shell triplet for 256 threads
-#define KTASKS_PER_BLOCK 16
+#define PAGES_PER_BLOCK  1048576
+// approximately, 15000 images in each ijk shell triplet for 256 threads
+#define KTASKS_PER_BLOCK 8
 
 typedef struct {
     int pair_ij;
@@ -44,8 +44,6 @@ __device__ __forceinline__ unsigned get_smid() {
     return smid;
 }
 
-#define allocate_page() (page_pool + atomicAdd(&num_pages, 1))
-
 __device__ __forceinline__
 void _filter_images(int& num_pages, // is stored in shm
                     ImgIdxPage *page_pool, PBCIntEnvVars &envs,
@@ -163,7 +161,12 @@ void _filter_images(int& num_pages, // is stored in shm
                         if (page != NULL) {
                             page->nimgs = PAGE_SIZE;
                         }
-                        page = allocate_page();
+                        int page_offset = atomicAdd(&num_pages, 1);
+                        if (page_offset >= PAGES_PER_BLOCK) {
+                            printf("Page overflow\n");
+                            __trap();
+                        }
+                        page = page_pool + page_offset;
                         page->pair_ij = pair_ij;
                         page->ksh = ksh;
                         counts = 0;
@@ -204,7 +207,6 @@ void pbc_int3c2e_latsum23_kernel(double *out, PBCIntEnvVars envs, PBCInt3c2eBoun
     int nksh_per_block = k_per_block * KTASKS_PER_BLOCK;
     int kcount0 = ksh_block_id * nksh_per_block;
     int kcount1 = min(ncells*bounds.nksh, kcount0 + nksh_per_block);
-    int nimgs = envs.nimgs;
 
     int li = bounds.li;
     int lj = bounds.lj;
@@ -248,7 +250,6 @@ void pbc_int3c2e_latsum23_kernel(double *out, PBCIntEnvVars envs, PBCInt3c2eBoun
                    kcount0, kcount1, log_cutoff);
     __syncthreads();
     if (num_pages >= PAGES_PER_BLOCK) {
-        printf("Page overflow\n");
         __trap();
     }
     __shared__ int img_max;

gpu4pyscf/pbc/df/int3c2e.py

@@ -56,7 +56,7 @@
 L_AUX_MAX = 6
 GOUT_WIDTH = 54
 THREADS = 256
-PAGES_PER_BLOCK = 524288
+PAGES_PER_BLOCK = 1048576
 PAGE_SIZE = 32 * 4 # Bytes
 
 def sr_aux_e2(cell, auxcell, omega, kpts=None, bvk_kmesh=None, j_only=False):

gpu4pyscf/pbc/df/rsdf_builder.py

@@ -332,7 +332,11 @@ def eigenvalue_decomposed_metric(j2c, linear_dep_threshold=LINEAR_DEP_THR):
     v1 = v[:,mask]
     v1 *= w[mask]**-.5
     j2c = v1
-    idx = cp.where(w < -linear_dep_threshold)[0]
+    # linear_dep_threshold for negative eigenvalues are too tight. Small errors
+    # in 2c2e metric would lead to small negative eigenvalues. They can be
+    # safely filtered.
+    #idx = cp.where(w < -linear_dep_threshold)[0]
+    idx = cp.where(w < -1e-4)[0]
     j2c_negative = None
     if len(idx) > 0:
         j2c_negative = (v[:,idx] * (-w[idx])**-.5)

gpu4pyscf/pbc/df/tests/test_rsdf_builder.py

@@ -638,3 +638,37 @@ def test_2c2e():
     dat = rsdf_builder._get_2c2e(cell, kpts, omega, with_long_range=True)
     ref = _get_2c2e_slow(cell, kpts, omega, with_long_range=True)
     assert abs(dat - cp.asarray(ref)).max() < 1e-10
+
+def test_kpts_compressed_linear_dep():
+    from pyscf.pbc.df import df as df_cpu
+    cell = pyscf.M(
+        atom='''
+        C 0.0 0.0 0.0
+        C 0.0 1.8 1.8
+        C 1.8 0.0 1.8
+        C 1.8 1.8 0.0''', a=np.eye(3) * 3.6,
+        basis=[[0, [4., 1.]],
+               [0, [.1, 1.]],
+               [0, [.035, 1.]]])
+    auxcell = df_cpu.make_auxcell(cell)
+    nao = cell.nao
+    kmesh = [2, 1, 1]
+    kpts = cell.make_kpts(kmesh)
+    dat, dat_neg, idx = rsdf_builder.compressed_cderi_kk(
+        cell, auxcell, kpts=kpts)
+    ref = build_cderi(cell, auxcell, kpts)[0]
+    kk_conserv = k2gamma.double_translation_indices(kmesh)
+    bvkmesh_Ls = k2gamma.translation_vectors_for_kmesh(cell, kmesh, True)
+    expLk = cp.exp(1j*cp.asarray(bvkmesh_Ls.dot(kpts.T)))
+    for kp in sorted(dat):
+        out = rsdf_builder.unpack_cderi(dat[kp], idx, kp, kk_conserv, expLk, nao)
+        ki_idx, kj_idx = np.where(kk_conserv == kp)
+        for ki, kj in zip(ki_idx, kj_idx):
+            if (ki, kj) in ref:
+                _ref = ref[ki, kj]
+            else:
+                _ref = ref[kj, ki].conj().transpose(0,2,1)
+            _ref = np.einsum('pij,plk->ijkl', _ref, _ref.conj(), optimize=True)
+            _dat = np.einsum('pij,plk->ijkl', out[ki], out[ki].conj(), optimize=True)
+            print(ki, kj)
+            assert abs(_ref - _dat).max() < 3e-7