PBC multigrid optimization

.github/workflows/unittest.yml

@@ -24,6 +24,7 @@ jobs:
         docker pull pyscf/gpu4pyscf-devel:pyscf-2.14
 
     - name: Test with pytest
+      timeout-minutes: 240
       run: |
         docker run --gpus all \
           --rm \
@@ -46,6 +47,7 @@ jobs:
         docker pull pyscf/gpu4pyscf-devel:pyscf-2.14
 
     - name: Test with pytest
+      timeout-minutes: 240
       run: |
         docker run --gpus all \
           --rm \

benchmarks/pbc/run_dft_gdf.py

@@ -140,6 +140,7 @@ def xc_type(xc):
                     kstring = 'x'.join([str(x) for x in kmesh])
                     confstr = f'{xc}-k{kstring}'
                     if 'supercell' in method:
+                        cupy.fft.config.get_plan_cache().clear()
                         ncopy = method['supercell']
                         scstring = 'x'.join([str(x) for x in ncopy])
                         confstr = f'{scstring}-{confstr}'
@@ -158,6 +159,7 @@ def xc_type(xc):
                     else:
                         kpts = cell.make_kpts(kmesh)
                         mf = cell.KRKS(xc=xc, kpts=kpts).to_gpu()
+                    cupy.get_default_memory_pool().free_all_blocks()
                     mf = mf.multigrid_numint()
                     mf = mf.density_fit()
                     mf.max_cycle = 20
@@ -170,4 +172,3 @@ def xc_type(xc):
                         import traceback
                         traceback.print_stack()
                         traceback.print_exception(e)
-                    cupy.get_default_memory_pool().free_all_blocks()

benchmarks/pbc/run_dft_multigrid.py

@@ -94,24 +94,42 @@ def xc_type(xc):
     smearing:
       sigma: 0.005
 - filename:
-  - MgO.cif
+  - Mg4O4_cubic.cif
   method:
-  - xc: PBE
+  - xc:
+    - PBE
+    - R2SCAN
     supercell: [2,2,2]
-  - xc: PBE
+  - xc:
+    - PBE
+    - R2SCAN
     supercell: [3,3,3]
-  - xc: PBE
+  - xc:
+    - PBE
+    - R2SCAN
     supercell: [4,4,4]
+  - xc:
+    - PBE
+    - R2SCAN
+    supercell: [5,5,5]
 - filename:
   - Al2Mg3O12Si3_ICSD_80847.cif
   method:
-  - xc: PBE
+  - xc:
+    - PBE
+    - R2SCAN
     supercell: [1,1,1]
-  - xc: PBE
+  - xc:
+    - PBE
+    - R2SCAN
     supercell: [2,1,1]
-  - xc: PBE
+  - xc:
+    - PBE
+    - R2SCAN
     supercell: [2,2,1]
-  - xc: PBE
+  - xc:
+    - PBE
+    - R2SCAN
     supercell: [2,2,2]
 ''')
 
@@ -139,6 +157,7 @@ def xc_type(xc):
                     kstring = 'x'.join([str(x) for x in kmesh])
                     confstr = f'{xc}-k{kstring}'
                     if 'supercell' in method:
+                        cupy.fft.config.get_plan_cache().clear()
                         ncopy = method['supercell']
                         scstring = 'x'.join([str(x) for x in ncopy])
                         confstr = f'{scstring}-{confstr}'
@@ -157,6 +176,7 @@ def xc_type(xc):
                     else:
                         kpts = cell.make_kpts(kmesh)
                         mf = cell.KRKS(xc=xc, kpts=kpts).to_gpu()
+                    cupy.get_default_memory_pool().free_all_blocks()
                     mf = mf.multigrid_numint()
                     mf.max_cycle = 20
                     mf.conv_tol = 1e-6
@@ -168,4 +188,3 @@ def xc_type(xc):
                         import traceback
                         traceback.print_stack()
                         traceback.print_exception(e)
-                    cupy.get_default_memory_pool().free_all_blocks()

benchmarks/pbc/run_dft_rsjk.py

@@ -137,6 +137,7 @@ def xc_type(xc):
                     kstring = 'x'.join([str(x) for x in kmesh])
                     confstr = f'{xc}-k{kstring}'
                     if 'supercell' in method:
+                        cupy.fft.config.get_plan_cache().clear()
                         ncopy = method['supercell']
                         scstring = 'x'.join([str(x) for x in ncopy])
                         confstr = f'{scstring}-{confstr}'
@@ -155,6 +156,7 @@ def xc_type(xc):
                     else:
                         kpts = cell.make_kpts(kmesh)
                         mf = cell.KRKS(xc=xc, kpts=kpts).to_gpu()
+                    cupy.get_default_memory_pool().free_all_blocks()
                     mf = mf.multigrid_numint()
                     mf.rsjk = rsjk.PBCJKMatrixOpt(cell)
                     mf.max_cycle = 20
@@ -167,4 +169,3 @@ def xc_type(xc):
                         import traceback
                         traceback.print_stack()
                         traceback.print_exception(e)
-                    cupy.get_default_memory_pool().free_all_blocks()

examples/pbc/30-evaluate_density.py

@@ -74,7 +74,7 @@
 # by `pbc.dft.numint.NumInt` or `pbc.dft.numint.multigrid.MultiGridNumInt` directly.
 # For example, use the multigrid integrator to evaluate the density
 #
-from gpu4pyscf.pbc.dft.multigrid_v2 import MultiGridNumInt
+from gpu4pyscf.pbc.dft.multigrid_v3 import MultiGridNumInt
 ni = MultiGridNumInt(cell)
 ni.mesh = [10, 10, 10]
 dm = mf.make_rdm1()
@@ -98,7 +98,7 @@
 # Alternatively, the multigrid integrator provides a specialized routine to
 # compute the reciprocal-space density directly, which is more efficient.
 #
-from gpu4pyscf.pbc.dft.multigrid_v2 import MultiGridNumInt, _eval_rhoG
+from gpu4pyscf.pbc.dft.multigrid_v3 import MultiGridNumInt, _eval_rhoG
 ni = MultiGridNumInt(cell)
 dm = mf.make_rdm1()
 rhoG_direct = _eval_rhoG(ni, dm, kpts=kpts)

gpu4pyscf/dft/numint.py

@@ -2161,6 +2161,18 @@ def nlc_coeff(self, xc_code):
     def rsh_coeff(sef, xc_code):
         return libxc_cpu.rsh_coeff(xc_code)
 
+    def is_hybrid_xc(self, xc_code):
+        if xc_code is None:
+            return False
+        if self.rsh_coeff(xc_code) != (0, 0, 0):
+            return True
+        if self.hybrid_coeff(xc_code) != 0:
+            return True
+        return False
+
+    def is_nlc(self, xc_code):
+        return self.libxc.is_nlc(xc_code)
+
     def _xc_type(self, xc_code):
         return libxc_cpu.xc_type(xc_code)
 
@@ -2219,7 +2231,7 @@ def reset(self):
         return self
 
     def eval_xc_eff(self, xc_code, rho, deriv=1, *, omega=None, xctype=None,
-                    spin=None, work=None):
+                    spin=None, work=None, inplace=False):
         if spin is None:
             if rho.ndim >= 2 and rho.shape[0] == 2:
                 spin = 1
@@ -2264,9 +2276,12 @@ def eval_xc_eff(self, xc_code, rho, deriv=1, *, omega=None, xctype=None,
             nvar = 4
         else:
             nvar = 5
+
         out = [None] * 4
         for i in range(deriv+1):
-            if spin == 0:
+            if i == 1 and inplace:
+                out[i] = rho
+            elif spin == 0:
                 out[i] = cupy.empty([nvar] * i + [ngrids])
             else:
                 out[i] = cupy.empty([2, nvar] * i + [ngrids])

gpu4pyscf/lib/cupy_helper.py

@@ -28,6 +28,7 @@
 from gpu4pyscf.lib import multi_gpu
 from gpu4pyscf.lib.utils import load_library
 from gpu4pyscf.__config__ import num_devices, _p2p_access
+from gpu4pyscf.__config__ import props as gpu_specs
 from gpu4pyscf import __config__
 
 LMAX_ON_GPU = 7
@@ -1224,6 +1225,42 @@ def malloc(size):
         return default_mempool_malloc(size)
     cupy.cuda.set_allocator(malloc)
 
+def vec_dot(vec1, vec2):
+    '''
+    einsum('g,g->', vec1, vec2)
+    '''
+    vec1 = cupy.asarray(vec1)
+    vec2 = cupy.asarray(vec2)
+    assert vec1.dtype == vec2.dtype == cupy.float64
+    assert vec1.shape == vec2.shape
+    n = vec1.size
+
+    fn_name = 'vec_dot_kernel_order'
+    if fn_name not in _kernel_registery:
+        kernel_code = (r'''
+extern "C" __global__
+void ''' + fn_name + r'''(double *out, double *vec1, double *vec2, long long n) {
+    int tid = threadIdx.x;
+    int idx = blockDim.x * blockIdx.x + tid;
+    int stride = gridDim.x * blockDim.x;
+    double sum = 0;
+    for (long long g = idx; g < n; g += stride) {
+        sum += vec1[g] * vec2[g];
+    }
+    for (int offset = 16; offset > 0; offset >>= 1) {
+        sum += __shfl_down_sync(0xffffffff, sum, offset);
+    }
+    int lane = tid % 32;
+    if (lane == 0) atomicAdd(out, sum);
+}''')
+        _kernel_registery[fn_name] = cupy.RawKernel(kernel_code, fn_name)
+
+    kernel = _kernel_registery[fn_name]
+    out = cupy.zeros(1)
+    workers = gpu_specs['multiProcessorCount']
+    kernel((workers*2,), (1024,), (out, vec1, vec2, n))
+    return out[0]
+
 def batched_vec3_norm2(batched_vec3):
     '''
     einsum('gx,gx->g', vec3, vec3) for the (N,3)-array vec3

gpu4pyscf/lib/gvhf-rys/vhf.cuh

@@ -27,13 +27,6 @@
 #define LMAX1           (LMAX+1)
 #define NCART_MAX       ((LMAX+1)*(LMAX+2)/2)
 
-// performance drop when TILE>2, reason unclear
-#define TILE            2
-#define TILE2           (TILE*TILE)
-#define TILE4           (TILE2*TILE2)
-// when nroots > 5, GWIDTH=57 may be better
-#define GWIDTH          42
-// 2MB per block
 #define QUEUE_DEPTH     65536
 
 #define MIN(x, y)       ((x) < (y) ? (x) : (y))

gpu4pyscf/lib/multigrid/CMakeLists.txt

@@ -1,17 +1,15 @@
-set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --ptxas-options=-v")# -maxrregcount=128")
-
-add_library(mgrid SHARED
-  mg_driver.cu estimator.cu
-  eval_rho.cu
-  eval_mat_lda.cu eval_mat_gga.cu
-  eval_tau.cu eval_mat_tau.cu
-)
-
-set_target_properties(mgrid PROPERTIES
-  LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}
-  CUDA_SEPARABLE_COMPILATION ON)
-
-#target_link_libraries(ft_ao OpenMP::OpenMP_C)
+#set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --ptxas-options=-v")# -maxrregcount=128")
+#
+#add_library(mgrid SHARED
+#  mg_driver.cu estimator.cu
+#  eval_rho.cu
+#  eval_mat_lda.cu eval_mat_gga.cu
+#  eval_tau.cu eval_mat_tau.cu
+#)
+#
+#set_target_properties(mgrid PROPERTIES
+#  LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}
+#  CUDA_SEPARABLE_COMPILATION ON)
 
 add_library(mgrid_v2 SHARED
         multigrid_v2/drivers.cu
@@ -24,7 +22,28 @@ set_target_properties(mgrid_v2 PROPERTIES
         LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}
         CUDA_SEPARABLE_COMPILATION ON
 )
-
 if(ENABLE_FP32_MULTIGRID)
   target_compile_definitions(mgrid_v2 PRIVATE GPU4PYSCF_ENABLE_FP32_MULTIGRID)
 endif()
+
+add_library(mgrid_v3 SHARED
+  multigrid_v3/screen.cu
+  multigrid_v3/drivers.cu
+  multigrid_v3/eval_density.cu
+  multigrid_v3/eval_tau.cu
+  #multigrid_v3/eval_lda_mat.cu
+  #multigrid_v3/eval_mgga_mat.cu
+  multigrid_v3/eval_lda_mat_v2.cu
+  multigrid_v3/eval_mgga_mat_v2.cu
+  multigrid_v3/eval_lda_strain_grad.cu
+  multigrid_v3/eval_mgga_strain_grad.cu
+  multigrid_v3/aft_eval_density.cu
+  multigrid_v3/aft_eval_tau.cu
+  multigrid_v3/aft_eval_lda_mat.cu
+  multigrid_v3/aft_eval_mgga_mat.cu
+)
+
+set_target_properties(mgrid_v3 PROPERTIES
+  LIBRARY_OUTPUT_DIRECTORY ${PROJECT_SOURCE_DIR}
+  CUDA_SEPARABLE_COMPILATION ON
+)