diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml index 7fd51c3b5..fdcd227df 100644 --- a/.github/workflows/unittest.yml +++ b/.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 \ diff --git a/benchmarks/pbc/run_dft_gdf.py b/benchmarks/pbc/run_dft_gdf.py index fba05db27..5d7ae2a81 100644 --- a/benchmarks/pbc/run_dft_gdf.py +++ b/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() diff --git a/benchmarks/pbc/run_dft_multigrid.py b/benchmarks/pbc/run_dft_multigrid.py index 0899417bb..45681770c 100644 --- a/benchmarks/pbc/run_dft_multigrid.py +++ b/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() diff --git a/benchmarks/pbc/run_dft_rsjk.py b/benchmarks/pbc/run_dft_rsjk.py index 18674c1e3..a21b72bef 100644 --- a/benchmarks/pbc/run_dft_rsjk.py +++ b/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() diff --git a/examples/pbc/30-evaluate_density.py b/examples/pbc/30-evaluate_density.py index 94d2d6685..b8b21f430 100644 --- a/examples/pbc/30-evaluate_density.py +++ b/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) diff --git a/gpu4pyscf/dft/numint.py b/gpu4pyscf/dft/numint.py index 845ee6671..e9c39aa34 100644 --- a/gpu4pyscf/dft/numint.py +++ b/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]) diff --git a/gpu4pyscf/lib/cupy_helper.py b/gpu4pyscf/lib/cupy_helper.py index 40bafb6a6..7239603f5 100644 --- a/gpu4pyscf/lib/cupy_helper.py +++ b/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 diff --git a/gpu4pyscf/lib/gvhf-rys/vhf.cuh b/gpu4pyscf/lib/gvhf-rys/vhf.cuh index 2e53c02da..208a3bc8e 100644 --- a/gpu4pyscf/lib/gvhf-rys/vhf.cuh +++ b/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)) diff --git a/gpu4pyscf/lib/multigrid/CMakeLists.txt b/gpu4pyscf/lib/multigrid/CMakeLists.txt index b0d86c0d4..7cd6eb530 100644 --- a/gpu4pyscf/lib/multigrid/CMakeLists.txt +++ b/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 +) \ No newline at end of file diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_density.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_density.cu new file mode 100644 index 000000000..653767db4 --- /dev/null +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_density.cu @@ -0,0 +1,245 @@ +/* + * Copyright 2026 The PySCF Developers. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include +#include "gvhf-rys/vhf.cuh" +#include "gvhf-rys/rys_contract_k.cuh" +#include "constant_objects.cuh" +#include "utils.cuh" +#include "aft_recursion.cuh" + +#define THREADS 256 +#define NGV_PER_BLOCK 16 +#define DENSITY_WIDTH 16 +#define REMOTE_THRESHOLD 50 +// pi^1.5 +#define OVERLAP_FAC 5.56832799683170787 + +#if CUDA_VERSION >= 12040 +__global__ __maxnreg__(128) static +#else +__global__ static +#endif +void orth_aopair_dm_kernel(double *outR, double *outI, double *dm, + PBCIntEnvVars envs, int *shl_pair_offsets, + int64_t *bas_ij_idx, double *G_bases, double *L_bases, + int *mesh_cum, int *nimgs_cum, int ntiles, + double factor) +{ + int thread_id = threadIdx.x; + int x_id = thread_id / NGV_PER_BLOCK; + int Gv_id = thread_id % NGV_PER_BLOCK; + int sp_block_id = blockIdx.x / ntiles; + int tile_id = blockIdx.x % ntiles; + __shared__ double gx[NGV_PER_BLOCK*3*2*LMAX1*LMAX1]; + __shared__ double swap[NGV_PER_BLOCK*3*2*(LMAX+LMAX+1)]; + __shared__ int mesh_start[3]; + __shared__ int ri, rj; + __shared__ size_t ij_offset; + __shared__ double fac, ai, aj; + + int *bas = envs.bas; + int nbas = envs.nbas; + double *env = envs.env; + + int mesh_x = mesh_cum[1] - mesh_cum[0]; + int mesh_y = mesh_cum[2] - mesh_cum[1]; + int mesh_z = mesh_cum[3] - mesh_cum[2]; + + if (thread_id == 0) { + int ntiles_y = (mesh_y + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; + int ntiles_z = (mesh_z + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; + int tile_z = tile_id % ntiles_z; + int tile_xy = tile_id / ntiles_z; + int tile_y = tile_xy % ntiles_y; + int tile_x = tile_xy / ntiles_y; + mesh_start[0] = tile_x * NGV_PER_BLOCK; + mesh_start[1] = tile_y * NGV_PER_BLOCK; + mesh_start[2] = tile_z * NGV_PER_BLOCK; + } + + double density_R[DENSITY_WIDTH]; + double density_I[DENSITY_WIDTH]; +#pragma unroll + for (int n = 0; n < DENSITY_WIDTH; ++n) { + density_R[n] = 0.; + density_I[n] = 0.; + } + + int shl_pair0 = shl_pair_offsets[sp_block_id]; + int shl_pair1 = shl_pair_offsets[sp_block_id+1]; + for (int pair_idx = shl_pair0; pair_idx < shl_pair1; pair_idx++) { + __syncthreads(); + int64_t bas_ij = bas_ij_idx[pair_idx]; + int ish = bas_ij / NBAS_MAX; + int jsh = bas_ij % NBAS_MAX; + int li = bas[ish*BAS_SLOTS+ANG_OF]; + int lj = bas[jsh*BAS_SLOTS+ANG_OF]; + if (thread_id == 0) { + int ci = bas[ish*BAS_SLOTS+PTR_COEFF]; + int cj = bas[jsh*BAS_SLOTS+PTR_COEFF]; + ai = env[bas[ish*BAS_SLOTS+PTR_EXP]]; + aj = env[bas[jsh*BAS_SLOTS+PTR_EXP]]; + double aij = ai + aj; + ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; + rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; + fac = OVERLAP_FAC * env[ci] * env[cj] / (aij * sqrt(aij)) * factor; + int ish_cell0 = ish; + int bvk_cell_id = jsh / nbas; + int jsh_cell0 = jsh - nbas * bvk_cell_id; + if (ish_cell0 == jsh_cell0) { + fac *= .5; + } + int i0 = envs.ao_loc[ish_cell0]; + int j0 = envs.ao_loc[jsh_cell0]; + size_t nao = envs.ao_loc[nbas]; + ij_offset = bvk_cell_id * nao * nao + i0 * nao + j0; + } + + constexpr int stride_i = NGV_PER_BLOCK * 6; + int stride_j = stride_i * LMAX1; + for (int n = thread_id; n < stride_j*(lj+1); n += THREADS) { + gx[n] = 0; + } + __syncthreads(); + + if (x_id < 3) { + double aij = ai + aj; + double a2 = .5 / aij; + double aj_aij = aj * 2 * a2; + double theta_ij = ai * aj_aij; + double kx = 0; + int _Gv_id = mesh_cum[x_id] + mesh_start[x_id] + Gv_id; + if (_Gv_id < mesh_cum[x_id+1]) { + kx = G_bases[_Gv_id]; + } + int addrR = x_id * NGV_PER_BLOCK*2 + Gv_id; + for (int img = nimgs_cum[x_id]; img < nimgs_cum[x_id+1]; ++img) { + double Lx = L_bases[img]; + double xi = env[ri+x_id]; + double xjxi = env[rj+x_id] + Lx - xi; + double theta_rr = theta_ij * xjxi * xjxi + .5*a2 * kx * kx; + if (theta_rr > REMOTE_THRESHOLD) continue; + vrr_hrr(gx, swap, addrR, li, lj, stride_j, a2, xjxi, aj_aij, xi, + kx, theta_rr); + } + } + + __syncthreads(); + int y_in_tile = thread_id / NGV_PER_BLOCK; + int z_in_tile = Gv_id; + if (mesh_start[1] + y_in_tile < mesh_y && mesh_start[2] + z_in_tile < mesh_z) { + int nfi = c_nf[li]; + int nfj = c_nf[lj]; + int idx_i = lex_xyz_offset(li); + int idx_j = lex_xyz_offset(lj); + size_t nao = envs.ao_loc[nbas]; + for (int i = 0; i < nfi; ++i) { + for (int j = 0; j < nfj; ++j) { + int ix = _c_cartesian_lexical_xyz[idx_i+i*3+0]; + int iy = _c_cartesian_lexical_xyz[idx_i+i*3+1]; + int iz = _c_cartesian_lexical_xyz[idx_i+i*3+2]; + int jx = _c_cartesian_lexical_xyz[idx_j+j*3+0]; + int jy = _c_cartesian_lexical_xyz[idx_j+j*3+1]; + int jz = _c_cartesian_lexical_xyz[idx_j+j*3+2]; + int addrx = ix*stride_i + jx*stride_j; + int addry = iy*stride_i + jy*stride_j + NGV_PER_BLOCK*2 + y_in_tile; + int addrz = iz*stride_i + jz*stride_j + NGV_PER_BLOCK*4 + z_in_tile; + double dm_fac = dm[ij_offset + i*nao+j] * fac; + double *gxR = gx; + double *gxI = gxR + NGV_PER_BLOCK; + double yR = gxR[addry]; + double yI = gxI[addry]; + double zR = gxR[addrz]; + double zI = gxI[addrz]; + double yzR, yzI; + multiply(yR, yI, zR, zI, yzR, yzI); + yzR *= dm_fac; + yzI *= dm_fac; +#pragma unroll + for (int n = 0; n < DENSITY_WIDTH; ++n) { + int x = n; + if (mesh_start[0] + x >= mesh_x) break; + double xR = gxR[addrx+x]; + double xI = gxI[addrx+x]; + double xyzR, xyzI; + multiply(xR, xI, yzR, yzI, xyzR, xyzI); + density_R[n] += xyzR; + density_I[n] += xyzI; + } + } } + } + } + + int Gx0 = mesh_start[0]; + int Gy0 = mesh_start[1]; + int Gz0 = mesh_start[2]; + int y_in_tile = thread_id / NGV_PER_BLOCK; + int z_in_tile = Gv_id; + int y = Gy0 + y_in_tile; + int z = Gz0 + z_in_tile; + if (y < mesh_y && z < mesh_z) { +#pragma unroll + for (int n = 0; n < DENSITY_WIDTH; ++n) { + int x = Gx0 + n; + if (x >= mesh_x) break; + int abc_idx = (x*mesh_y+y)*mesh_z+z; + atomicAdd(outR+abc_idx, density_R[n]); + atomicAdd(outI+abc_idx, density_I[n]); + } + } +} + +//__global__ static +//void monoclinic_aopair_dm_kernel(double *outR, double *outI, double *dm, +// PBCIntEnvVars *envs, +// int *shl_pair_offsets, int64_t *bas_ij_idx, +// double *G_bases, int *mesh_cum, int *nimgs_cum, +// int *mesh, int nbatches_shl_pair) +//{ +//} + +extern "C" { +int orth_contract_aopair_dm(double *outR, double *outI, + double *placeholderR, double *placeholderI, double *dm, + PBCIntEnvVars *envs, int *shl_pair_offsets, + int64_t *bas_ij_idx, double *G_bases, double *L_bases, + int *mesh_cum, int *nimgs_cum, int *mesh, + int nbatches_shl_pair, double factor) +{ + int mesh_x = mesh[0]; + int mesh_y = mesh[1]; + int mesh_z = mesh[2]; + int ntiles_x = (mesh_x + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; + int ntiles_y = (mesh_y + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; + int ntiles_z = (mesh_z + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; + int ntiles = ntiles_x * ntiles_y * ntiles_z; + orth_aopair_dm_kernel<<>>( + outR, outI, dm, *envs, shl_pair_offsets, bas_ij_idx, G_bases, L_bases, + mesh_cum, nimgs_cum, ntiles, factor); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in orth_aopair_dm_kernel: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} +} diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_lda_mat.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_lda_mat.cu new file mode 100644 index 000000000..294e1c9a1 --- /dev/null +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_lda_mat.cu @@ -0,0 +1,258 @@ +/* + * Copyright 2026 The PySCF Developers. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include +#include "gvhf-rys/vhf.cuh" +#include "gvhf-rys/rys_contract_k.cuh" +#include "constant_objects.cuh" +#include "utils.cuh" +#include "aft_recursion.cuh" + +#define WARP_SIZE 32 +#define WARPS 8 +#define THREADS 256 +#define NGV_PER_BLOCK 16 +#define DENSITY_WIDTH 16 +#define TILES_PER_BATCH 64 +#define REMOTE_THRESHOLD 50 +// pi^1.5 +#define OVERLAP_FAC 5.56832799683170787 + +#if CUDA_VERSION >= 12040 +__global__ __maxnreg__(128) static +#else +__global__ static +#endif +void orth_lda_mat_kernel(double *out, cuDoubleComplex *vxcG, + PBCIntEnvVars envs, int64_t *bas_ij_idx, + double *G_bases, double *L_bases, + int *mesh_cum, int *nimgs_cum, + int npair, int ntiles_x, int ntiles_y, int ntiles_z) +{ + int thread_id = threadIdx.x; + int x_id = thread_id / NGV_PER_BLOCK; + int Gv_id = thread_id % NGV_PER_BLOCK; + __shared__ int tile_batch; + int pair_id = blockIdx.x % npair; + if (thread_id == 0) { + tile_batch = blockIdx.x / npair; + } + __shared__ double gx[NGV_PER_BLOCK*3*2*LMAX1*LMAX1]; + __shared__ double swap[NGV_PER_BLOCK*3*2*(LMAX+LMAX+1)]; + __shared__ int mesh_start[3]; + __shared__ double vjR[NCART_MAX*NCART_MAX * WARPS]; + __shared__ int ri, rj, li, lj; + __shared__ double ai, aj; + + int mesh_x = mesh_cum[1] - mesh_cum[0]; + int mesh_y = mesh_cum[2] - mesh_cum[1]; + int mesh_z = mesh_cum[3] - mesh_cum[2]; + int nbas = envs.nbas; + int *bas = envs.bas; + double *env = envs.env; + int64_t bas_ij = bas_ij_idx[pair_id]; + int ish = bas_ij / NBAS_MAX; + int jsh = bas_ij % NBAS_MAX; + if (thread_id == 0) { + li = bas[ish*BAS_SLOTS+ANG_OF]; + lj = bas[jsh*BAS_SLOTS+ANG_OF]; + ai = env[bas[ish*BAS_SLOTS+PTR_EXP]]; + aj = env[bas[jsh*BAS_SLOTS+PTR_EXP]]; + ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; + rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; + } + __syncthreads(); + + int nfi = c_nf[li]; + int nfj = c_nf[lj]; + for (int n = thread_id; n < nfi * nfj * WARPS; n += THREADS) { + vjR[n] = 0.; + } + + int ntiles = ntiles_x * ntiles_y * ntiles_z; + int tile0 = tile_batch * TILES_PER_BATCH; + int tile1 = min(tile0 + TILES_PER_BATCH, ntiles); + for (int tile_id = tile0; tile_id < tile1; tile_id++) { + __syncthreads(); + int tile_z = tile_id % ntiles_z; + int tile_xy = tile_id / ntiles_z; + int tile_y = tile_xy % ntiles_y; + int tile_x = tile_xy / ntiles_y; + if (thread_id == 0) { + mesh_start[0] = tile_x * NGV_PER_BLOCK; + mesh_start[1] = tile_y * NGV_PER_BLOCK; + mesh_start[2] = tile_z * NGV_PER_BLOCK; + } + + constexpr int stride_i = NGV_PER_BLOCK * 6; + constexpr int stride_j = stride_i * LMAX1; + for (int n = thread_id; n < stride_j*(lj+1); n += THREADS) { + gx[n] = 0; + } + __syncthreads(); + + if (x_id < 3) { + double aij = ai + aj; + double a2 = .5 / aij; + double aj_aij = aj * 2 * a2; + double theta_ij = ai * aj_aij; + double kx = 0; + int _Gv_id = mesh_cum[x_id] + mesh_start[x_id] + Gv_id; + if (_Gv_id < mesh_cum[x_id+1]) { + kx = G_bases[_Gv_id]; + } + int addrR = x_id * NGV_PER_BLOCK*2 + Gv_id; + for (int img = nimgs_cum[x_id]; img < nimgs_cum[x_id+1]; ++img) { + double Lx = L_bases[img]; + double xi = env[ri+x_id]; + double xjxi = env[rj+x_id] + Lx - xi; + double theta_rr = theta_ij * xjxi * xjxi + .5*a2 * kx * kx; + if (theta_rr > REMOTE_THRESHOLD) continue; + vrr_hrr(gx, swap, addrR, li, lj, stride_j, a2, xjxi, aj_aij, xi, + kx, theta_rr); + } + } + __syncthreads(); + + int y_in_tile = thread_id / NGV_PER_BLOCK; + int z_in_tile = Gv_id; + int y = mesh_start[1] + y_in_tile; + int z = mesh_start[2] + z_in_tile; + double vG_R[DENSITY_WIDTH]; + double vG_I[DENSITY_WIDTH]; + if (y < mesh_y && z < mesh_z) { +#pragma unroll + for (int n = 0; n < DENSITY_WIDTH; ++n) { + int x = mesh_start[0] + n; + if (x >= mesh_x) break; + size_t addr = (x * mesh_y + y) * (size_t)mesh_z + z; + cuDoubleComplex val = vxcG[addr]; + vG_R[n] = val.x; + vG_I[n] = val.y; + } + } + + int nfi = c_nf[li]; + int nfj = c_nf[lj]; + int idx_i = lex_xyz_offset(li); + int idx_j = lex_xyz_offset(lj); + for (int i = 0; i < nfi; ++i) { + for (int j = 0; j < nfj; ++j) { + double s = 0; + if (y < mesh_y && z < mesh_z) { + int ix = _c_cartesian_lexical_xyz[idx_i+i*3+0]; + int iy = _c_cartesian_lexical_xyz[idx_i+i*3+1]; + int iz = _c_cartesian_lexical_xyz[idx_i+i*3+2]; + int jx = _c_cartesian_lexical_xyz[idx_j+j*3+0]; + int jy = _c_cartesian_lexical_xyz[idx_j+j*3+1]; + int jz = _c_cartesian_lexical_xyz[idx_j+j*3+2]; + int addrx = ix*stride_i + jx*stride_j; + int addry = iy*stride_i + jy*stride_j + NGV_PER_BLOCK*2 + y_in_tile; + int addrz = iz*stride_i + jz*stride_j + NGV_PER_BLOCK*4 + z_in_tile; + double *gxR = gx; + double *gxI = gxR + NGV_PER_BLOCK; + double yR = gxR[addry]; + double yI = gxI[addry]; + double zR = gxR[addrz]; + double zI = gxI[addrz]; + double yzR, yzI; + multiply(yR, yI, zR, zI, yzR, yzI); +#pragma unroll + for (int n = 0; n < DENSITY_WIDTH; ++n) { + int x = n; + if (mesh_start[0] + x >= mesh_x) break; + double xR = gxR[addrx+x]; + double xI = gxI[addrx+x]; + double xyzR, xyzI; + multiply(xR, xI, yzR, yzI, xyzR, xyzI); + s += xyzR * vG_R[n] - xyzI * vG_I[n]; + } + } + for (int offset = 16; offset > 0; offset >>= 1) { + s += __shfl_down_sync(0xffffffff, s, offset); + } + int lane = thread_id % WARP_SIZE; + int warp = thread_id / WARP_SIZE; + if (lane == 0) { + vjR[warp + WARPS*(i*nfj+j)] += s; + } + } } + } + + __syncthreads(); + for (int n = thread_id; n < nfi * nfj; n += THREADS) { + int ci = bas[ish*BAS_SLOTS+PTR_COEFF]; + int cj = bas[jsh*BAS_SLOTS+PTR_COEFF]; + double aij = ai + aj; + double fac = OVERLAP_FAC * env[ci] * env[cj] / (aij * sqrt(aij)); + int ish_cell0 = ish; + int bvk_cell_id = jsh / nbas; + int jsh_cell0 = jsh - nbas * bvk_cell_id; + if (ish_cell0 == jsh_cell0) { + fac *= .5; + } + size_t nao = envs.ao_loc[nbas]; + int i0 = envs.ao_loc[ish_cell0]; + int j0 = envs.ao_loc[jsh_cell0]; + int i = n * c_div_nf[lj]; + int j = n - nfj * i; + double s = 0; + for (int m = 0; m < WARPS; m++) { + s += vjR[n*WARPS+m]; + } + atomicAdd(out + bvk_cell_id*nao*nao + (i0+i)*nao + j0+j, s * fac); + } +} + +//__global__ static +//void monoclinic_aopair_coulG_kernel(double *out, double *coulG_R, double *coulG_I, +// PBCIntEnvVars *envs, +// int *shl_pair_offsets, int64_t *bas_ij_idx, +// double *G_bases, int *mesh_cum, int *nimgs_cum, +// int *mesh, int nbatches_shl_pair) +//{ +//} + +extern "C" { +int orth_aft_lda_mat(double *out, cuDoubleComplex *vxcG, cuDoubleComplex *placeholder, + PBCIntEnvVars *envs, int64_t *bas_ij_idx, + double *G_bases, double *L_bases, int *mesh_cum, + int *nimgs_cum, int *mesh, int npair) +{ + int mesh_x = mesh[0]; + int mesh_y = mesh[1]; + int mesh_z = mesh[2]; + int ntiles_x = (mesh_x + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; + int ntiles_y = (mesh_y + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; + int ntiles_z = (mesh_z + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; + int ntiles = ntiles_x * ntiles_y * ntiles_z; + int ntile_batch = (ntiles + TILES_PER_BATCH-1) / TILES_PER_BATCH; + orth_lda_mat_kernel<<>>( + out, vxcG, *envs, bas_ij_idx, G_bases, L_bases, + mesh_cum, nimgs_cum, npair, ntiles_x, ntiles_y, ntiles_z); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in orth_lda_mat_kernel: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} +} diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_mgga_mat.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_mgga_mat.cu new file mode 100644 index 000000000..9d415a738 --- /dev/null +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_mgga_mat.cu @@ -0,0 +1,310 @@ +/* + * Copyright 2026 The PySCF Developers. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include +#include "gvhf-rys/vhf.cuh" +#include "gvhf-rys/rys_contract_k.cuh" +#include "constant_objects.cuh" +#include "utils.cuh" +#include "aft_recursion.cuh" + +#define WARP_SIZE 32 +#define WARPS 8 +#define THREADS 256 +#define NGV_PER_BLOCK 16 +#define DENSITY_WIDTH 16 +#define TILES_PER_BATCH 64 +#define REMOTE_THRESHOLD 50 +// pi^1.5 +#define OVERLAP_FAC 5.56832799683170787 + +__global__ static +void orth_mgga_mat_kernel(double *out, cuDoubleComplex *vrhoG, + cuDoubleComplex *vtauG, + PBCIntEnvVars envs, int64_t *bas_ij_idx, + double *G_bases, double *L_bases, + int *mesh_cum, int *nimgs_cum, + int npair, int ntiles_x, int ntiles_y, int ntiles_z) +{ + int thread_id = threadIdx.x; + int x_id = thread_id / NGV_PER_BLOCK; + int Gv_id = thread_id % NGV_PER_BLOCK; + __shared__ int tile_batch; + int pair_id = blockIdx.x % npair; + if (thread_id == 0) { + tile_batch = blockIdx.x / npair; + } + __shared__ double gx[NGV_PER_BLOCK*3*2*(LMAX1+1)*(LMAX1+1)]; + __shared__ int mesh_start[3]; + __shared__ double vjR[NCART_MAX*NCART_MAX * WARPS]; + __shared__ int ri, rj, li, lj; + __shared__ double ai, aj; + + int mesh_x = mesh_cum[1] - mesh_cum[0]; + int mesh_y = mesh_cum[2] - mesh_cum[1]; + int mesh_z = mesh_cum[3] - mesh_cum[2]; + int nbas = envs.nbas; + int *bas = envs.bas; + double *env = envs.env; + int64_t bas_ij = bas_ij_idx[pair_id]; + int ish = bas_ij / NBAS_MAX; + int jsh = bas_ij % NBAS_MAX; + if (thread_id == 0) { + li = bas[ish*BAS_SLOTS+ANG_OF]; + lj = bas[jsh*BAS_SLOTS+ANG_OF]; + ai = env[bas[ish*BAS_SLOTS+PTR_EXP]]; + aj = env[bas[jsh*BAS_SLOTS+PTR_EXP]]; + ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; + rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; + } + __syncthreads(); + + int nfi = c_nf[li]; + int nfj = c_nf[lj]; + for (int n = thread_id; n < nfi * nfj * WARPS; n += THREADS) { + vjR[n] = 0.; + } + + int ntiles = ntiles_x * ntiles_y * ntiles_z; + int tile0 = tile_batch * TILES_PER_BATCH; + int tile1 = min(tile0 + TILES_PER_BATCH, ntiles); + for (int tile_id = tile0; tile_id < tile1; tile_id++) { + __syncthreads(); + int tile_z = tile_id % ntiles_z; + int tile_xy = tile_id / ntiles_z; + int tile_y = tile_xy % ntiles_y; + int tile_x = tile_xy / ntiles_y; + if (thread_id == 0) { + mesh_start[0] = tile_x * NGV_PER_BLOCK; + mesh_start[1] = tile_y * NGV_PER_BLOCK; + mesh_start[2] = tile_z * NGV_PER_BLOCK; + } + + constexpr int stride_i = NGV_PER_BLOCK * 6; + constexpr int stride_j = stride_i * (LMAX+2); + for (int n = thread_id; n < stride_j*(lj+2); n += THREADS) { + gx[n] = 0; + } + __syncthreads(); + + if (x_id < 3) { + double aij = ai + aj; + double a2 = .5 / aij; + double aj_aij = aj * 2 * a2; + double theta_ij = ai * aj_aij; + double kx = 0; + int _Gv_id = mesh_cum[x_id] + mesh_start[x_id] + Gv_id; + if (_Gv_id < mesh_cum[x_id+1]) { + kx = G_bases[_Gv_id]; + } + int addrR = x_id * NGV_PER_BLOCK*2 + Gv_id; + for (int img = nimgs_cum[x_id]; img < nimgs_cum[x_id+1]; ++img) { + double Lx = L_bases[img]; + double xi = env[ri+x_id]; + double xjxi = env[rj+x_id] + Lx - xi; + double theta_rr = theta_ij * xjxi * xjxi + .5*a2 * kx * kx; + if (theta_rr > REMOTE_THRESHOLD) continue; + double xpa = xjxi * aj_aij; + double xij = xpa + xi; + double kR = kx * xij; + double g00R; + double g00I; + sincos(-kR, &g00I, &g00R); + double Kab = exp(-theta_rr); + g00R *= Kab; + g00I *= Kab; + double RpaR = xpa; + double RpaI = -a2 * kx; + switch (li*LMAX1+lj) { + case (0*LMAX1+0): vrr_hrr<1,1>(gx, addrR, stride_j, a2, xjxi, RpaR, RpaI, g00R, g00I); break; + case (1*LMAX1+0): vrr_hrr<2,1>(gx, addrR, stride_j, a2, xjxi, RpaR, RpaI, g00R, g00I); break; + case (1*LMAX1+1): vrr_hrr<2,2>(gx, addrR, stride_j, a2, xjxi, RpaR, RpaI, g00R, g00I); break; + case (2*LMAX1+0): vrr_hrr<3,1>(gx, addrR, stride_j, a2, xjxi, RpaR, RpaI, g00R, g00I); break; + case (2*LMAX1+1): vrr_hrr<3,2>(gx, addrR, stride_j, a2, xjxi, RpaR, RpaI, g00R, g00I); break; + case (2*LMAX1+2): vrr_hrr<3,3>(gx, addrR, stride_j, a2, xjxi, RpaR, RpaI, g00R, g00I); break; + case (3*LMAX1+0): vrr_hrr<4,1>(gx, addrR, stride_j, a2, xjxi, RpaR, RpaI, g00R, g00I); break; + case (3*LMAX1+1): vrr_hrr<4,2>(gx, addrR, stride_j, a2, xjxi, RpaR, RpaI, g00R, g00I); break; + case (3*LMAX1+2): vrr_hrr<4,3>(gx, addrR, stride_j, a2, xjxi, RpaR, RpaI, g00R, g00I); break; + case (3*LMAX1+3): vrr_hrr<4,4>(gx, addrR, stride_j, a2, xjxi, RpaR, RpaI, g00R, g00I); break; + case (4*LMAX1+0): vrr_hrr<5,1>(gx, addrR, stride_j, a2, xjxi, RpaR, RpaI, g00R, g00I); break; + case (4*LMAX1+1): vrr_hrr<5,2>(gx, addrR, stride_j, a2, xjxi, RpaR, RpaI, g00R, g00I); break; + case (4*LMAX1+2): vrr_hrr<5,3>(gx, addrR, stride_j, a2, xjxi, RpaR, RpaI, g00R, g00I); break; + case (4*LMAX1+3): vrr_hrr<5,4>(gx, addrR, stride_j, a2, xjxi, RpaR, RpaI, g00R, g00I); break; + case (4*LMAX1+4): vrr_hrr<5,5>(gx, addrR, stride_j, a2, xjxi, RpaR, RpaI, g00R, g00I); break; + } + } + } + __syncthreads(); + + int y_in_tile = thread_id / NGV_PER_BLOCK; + int z_in_tile = Gv_id; + int y = mesh_start[1] + y_in_tile; + int z = mesh_start[2] + z_in_tile; + double vrho_R[DENSITY_WIDTH]; + double vrho_I[DENSITY_WIDTH]; + double vtau_R[DENSITY_WIDTH]; + double vtau_I[DENSITY_WIDTH]; + if (y < mesh_y && z < mesh_z) { +#pragma unroll + for (int n = 0; n < DENSITY_WIDTH; ++n) { + int x = mesh_start[0] + n; + if (x >= mesh_x) break; + size_t addr = (x * mesh_y + y) * (size_t)mesh_z + z; + cuDoubleComplex val = vrhoG[addr]; + vrho_R[n] = val.x; + vrho_I[n] = val.y; + val = vtauG[addr]; + vtau_R[n] = val.x / 2; + vtau_I[n] = val.y / 2; + } + } + + int nfi = c_nf[li]; + int nfj = c_nf[lj]; + int idx_i = lex_xyz_offset(li); + int idx_j = lex_xyz_offset(lj); + for (int i = 0; i < nfi; ++i) { + for (int j = 0; j < nfj; ++j) { + double s = 0; + if (y < mesh_y && z < mesh_z) { + int ix = _c_cartesian_lexical_xyz[idx_i+i*3+0]; + int iy = _c_cartesian_lexical_xyz[idx_i+i*3+1]; + int iz = _c_cartesian_lexical_xyz[idx_i+i*3+2]; + int jx = _c_cartesian_lexical_xyz[idx_j+j*3+0]; + int jy = _c_cartesian_lexical_xyz[idx_j+j*3+1]; + int jz = _c_cartesian_lexical_xyz[idx_j+j*3+2]; + int addrx = ix*stride_i + jx*stride_j; + int addry = iy*stride_i + jy*stride_j + NGV_PER_BLOCK*2 + y_in_tile; + int addrz = iz*stride_i + jz*stride_j + NGV_PER_BLOCK*4 + z_in_tile; + double *gxR = gx; + double *gxI = gxR + NGV_PER_BLOCK; + double yR0 = gxR[addry]; + double yI0 = gxI[addry]; + double zR0 = gxR[addrz]; + double zI0 = gxI[addrz]; + double yzR00, yzI00; + multiply(yR0, yI0, zR0, zI0, yzR00, yzI00); + + double ai2 = ai * -2; + double aj2 = aj * -2; + double yR3, yI3; + dIdJ_gx(gxR, addry, stride_i, stride_j, iy, jy, ai2, aj2, yR3, yI3); + double yzR33, yzI33; + multiply(yR3, yI3, zR0, zI0, yzR33, yzI33); + + double zR3, zI3; + dIdJ_gx(gxR, addrz, stride_i, stride_j, iz, jz, ai2, aj2, zR3, zI3); + double tmpR, tmpI; + multiply(yR0, yI0, zR3, zI3, tmpR, tmpI); + yzR33 += tmpR; + yzI33 += tmpI; +#pragma unroll + for (int n = 0; n < DENSITY_WIDTH; ++n) { + int x = n; + if (mesh_start[0] + x >= mesh_x) break; + int addr = addrx + x; + double xR3, xI3; + dIdJ_gx(gxR, addr, stride_i, stride_j, ix, jx, ai2, aj2, xR3, xI3); + double xyzR, xyzI; + multiply(xR3, xI3, yzR00, yzI00, xyzR, xyzI); + + double xR0 = gxR[addr]; + double xI0 = gxI[addr]; + double tmpR, tmpI; + multiply(xR0, xI0, yzR33, yzI33, tmpR, tmpI); + xyzR += tmpR; + xyzI += tmpI; + s += xyzR * vtau_R[n] - xyzI * vtau_I[n]; + + multiply(xR0, xI0, yzR00, yzI00, xyzR, xyzI); + s += xyzR * vrho_R[n] - xyzI * vrho_I[n]; + } + } + for (int offset = 16; offset > 0; offset >>= 1) { + s += __shfl_down_sync(0xffffffff, s, offset); + } + int lane = thread_id % WARP_SIZE; + int warp = thread_id / WARP_SIZE; + if (lane == 0) { + vjR[warp + WARPS*(i*nfj+j)] += s; + } + } } + } + + __syncthreads(); + for (int n = thread_id; n < nfi * nfj; n += THREADS) { + int ci = bas[ish*BAS_SLOTS+PTR_COEFF]; + int cj = bas[jsh*BAS_SLOTS+PTR_COEFF]; + double aij = ai + aj; + double fac = OVERLAP_FAC * env[ci] * env[cj] / (aij * sqrt(aij)); + int ish_cell0 = ish; + int bvk_cell_id = jsh / nbas; + int jsh_cell0 = jsh - nbas * bvk_cell_id; + if (ish_cell0 == jsh_cell0) { + fac *= .5; + } + size_t nao = envs.ao_loc[nbas]; + int i0 = envs.ao_loc[ish_cell0]; + int j0 = envs.ao_loc[jsh_cell0]; + int i = n * c_div_nf[lj]; + int j = n - nfj * i; + double s = 0; + for (int m = 0; m < WARPS; m++) { + s += vjR[n*WARPS+m]; + } + atomicAdd(out + bvk_cell_id*nao*nao + (i0+i)*nao + j0+j, s * fac); + } +} + +//__global__ static +//void monoclinic_aopair_coulG_kernel(double *out, double *coulG_R, double *coulG_I, +// PBCIntEnvVars *envs, +// int *shl_pair_offsets, int64_t *bas_ij_idx, +// double *G_bases, int *mesh_cum, int *nimgs_cum, +// int *mesh, int nbatches_shl_pair) +//{ +//} + +extern "C" { +int orth_aft_mgga_mat(double *out, cuDoubleComplex *vrhoG, cuDoubleComplex *vtauG, + PBCIntEnvVars *envs, int64_t *bas_ij_idx, + double *G_bases, double *L_bases, int *mesh_cum, + int *nimgs_cum, int *mesh, int npair) +{ + int mesh_x = mesh[0]; + int mesh_y = mesh[1]; + int mesh_z = mesh[2]; + int ntiles_x = (mesh_x + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; + int ntiles_y = (mesh_y + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; + int ntiles_z = (mesh_z + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; + int ntiles = ntiles_x * ntiles_y * ntiles_z; + int ntile_batch = (ntiles + TILES_PER_BATCH-1) / TILES_PER_BATCH; + orth_mgga_mat_kernel<<>>( + out, vrhoG, vtauG, *envs, bas_ij_idx, G_bases, L_bases, + mesh_cum, nimgs_cum, npair, ntiles_x, ntiles_y, ntiles_z); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in orth_mgga_mat_kernel: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} +} diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_tau.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_tau.cu new file mode 100644 index 000000000..6084ee630 --- /dev/null +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/aft_eval_tau.cu @@ -0,0 +1,263 @@ +/* + * Copyright 2026 The PySCF Developers. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include +#include "gvhf-rys/vhf.cuh" +#include "gvhf-rys/rys_contract_k.cuh" +#include "constant_objects.cuh" +#include "utils.cuh" +#include "aft_recursion.cuh" + +#define THREADS 256 +#define NGV_PER_BLOCK 16 +#define DENSITY_WIDTH 16 +#define REMOTE_THRESHOLD 50 +// pi^1.5 +#define OVERLAP_FAC 5.56832799683170787 + +__global__ static +void orth_ft_tau_dm_kernel(double *densityR, double *densityI, double *tauR, double *tauI, + double *dm, PBCIntEnvVars envs, int *shl_pair_offsets, + int64_t *bas_ij_idx, double *G_bases, double *L_bases, + int *mesh_cum, int *nimgs_cum, int ntiles, double factor) +{ + int thread_id = threadIdx.x; + int x_id = thread_id / NGV_PER_BLOCK; + int Gv_id = thread_id % NGV_PER_BLOCK; + int sp_block_id = blockIdx.x / ntiles; + int tile_id = blockIdx.x % ntiles; + __shared__ double gx[NGV_PER_BLOCK*3*2*(LMAX1+1)*(LMAX1+1)]; + __shared__ double swap[NGV_PER_BLOCK*3*2*(LMAX+LMAX+3)]; + __shared__ int mesh_start[3]; + __shared__ int ri, rj; + __shared__ size_t ij_offset; + __shared__ double fac, ai, aj; + + int *bas = envs.bas; + int nbas = envs.nbas; + double *env = envs.env; + + int mesh_x = mesh_cum[1] - mesh_cum[0]; + int mesh_y = mesh_cum[2] - mesh_cum[1]; + int mesh_z = mesh_cum[3] - mesh_cum[2]; + + if (thread_id == 0) { + int ntiles_y = (mesh_y + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; + int ntiles_z = (mesh_z + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; + int tile_z = tile_id % ntiles_z; + int tile_xy = tile_id / ntiles_z; + int tile_y = tile_xy % ntiles_y; + int tile_x = tile_xy / ntiles_y; + mesh_start[0] = tile_x * NGV_PER_BLOCK; + mesh_start[1] = tile_y * NGV_PER_BLOCK; + mesh_start[2] = tile_z * NGV_PER_BLOCK; + } + + double rho_R[DENSITY_WIDTH]; + double rho_I[DENSITY_WIDTH]; + double tau_R[DENSITY_WIDTH]; + double tau_I[DENSITY_WIDTH]; +#pragma unroll + for (int n = 0; n < DENSITY_WIDTH; ++n) { + rho_R[n] = 0.; + rho_I[n] = 0.; + tau_R[n] = 0.; + tau_I[n] = 0.; + } + + int shl_pair0 = shl_pair_offsets[sp_block_id]; + int shl_pair1 = shl_pair_offsets[sp_block_id+1]; + for (int pair_idx = shl_pair0; pair_idx < shl_pair1; pair_idx++) { + __syncthreads(); + int64_t bas_ij = bas_ij_idx[pair_idx]; + int ish = bas_ij / NBAS_MAX; + int jsh = bas_ij % NBAS_MAX; + int li = bas[ish*BAS_SLOTS+ANG_OF]; + int lj = bas[jsh*BAS_SLOTS+ANG_OF]; + if (thread_id == 0) { + int ci = bas[ish*BAS_SLOTS+PTR_COEFF]; + int cj = bas[jsh*BAS_SLOTS+PTR_COEFF]; + ai = env[bas[ish*BAS_SLOTS+PTR_EXP]]; + aj = env[bas[jsh*BAS_SLOTS+PTR_EXP]]; + double aij = ai + aj; + ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; + rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; + fac = OVERLAP_FAC * env[ci] * env[cj] / (aij * sqrt(aij)) * factor; + int ish_cell0 = ish; + int bvk_cell_id = jsh / nbas; + int jsh_cell0 = jsh - nbas * bvk_cell_id; + if (ish_cell0 == jsh_cell0) { + fac *= .5; + } + int i0 = envs.ao_loc[ish_cell0]; + int j0 = envs.ao_loc[jsh_cell0]; + size_t nao = envs.ao_loc[nbas]; + ij_offset = bvk_cell_id * nao * nao + i0 * nao + j0; + } + + constexpr int stride_i = NGV_PER_BLOCK * 6; + constexpr int stride_j = stride_i * (LMAX+2); + for (int n = thread_id; n < stride_j*(lj+2); n += THREADS) { + gx[n] = 0; + } + __syncthreads(); + + if (x_id < 3) { + double aij = ai + aj; + double a2 = .5 / aij; + double aj_aij = aj * 2 * a2; + double theta_ij = ai * aj_aij; + double kx = 0; + int _Gv_id = mesh_cum[x_id] + mesh_start[x_id] + Gv_id; + if (_Gv_id < mesh_cum[x_id+1]) { + kx = G_bases[_Gv_id]; + } + int addrR = x_id * NGV_PER_BLOCK*2 + Gv_id; + for (int img = nimgs_cum[x_id]; img < nimgs_cum[x_id+1]; ++img) { + double Lx = L_bases[img]; + double xi = env[ri+x_id]; + double xjxi = env[rj+x_id] + Lx - xi; + double theta_rr = theta_ij * xjxi * xjxi + .5*a2 * kx * kx; + if (theta_rr > REMOTE_THRESHOLD) continue; + vrr_hrr(gx, swap, addrR, li+1, lj+1, stride_j, a2, xjxi, aj_aij, xi, + kx, theta_rr); + } + } + + __syncthreads(); + int y_in_tile = thread_id / NGV_PER_BLOCK; + int z_in_tile = Gv_id; + if (mesh_start[1] + y_in_tile < mesh_y && mesh_start[2] + z_in_tile < mesh_z) { + int nfi = c_nf[li]; + int nfj = c_nf[lj]; + int idx_i = lex_xyz_offset(li); + int idx_j = lex_xyz_offset(lj); + size_t nao = envs.ao_loc[nbas]; + for (int i = 0; i < nfi; ++i) { + for (int j = 0; j < nfj; ++j) { + int ix = _c_cartesian_lexical_xyz[idx_i+i*3+0]; + int iy = _c_cartesian_lexical_xyz[idx_i+i*3+1]; + int iz = _c_cartesian_lexical_xyz[idx_i+i*3+2]; + int jx = _c_cartesian_lexical_xyz[idx_j+j*3+0]; + int jy = _c_cartesian_lexical_xyz[idx_j+j*3+1]; + int jz = _c_cartesian_lexical_xyz[idx_j+j*3+2]; + int addrx = ix*stride_i + jx*stride_j; + int addry = iy*stride_i + jy*stride_j + NGV_PER_BLOCK*2 + y_in_tile; + int addrz = iz*stride_i + jz*stride_j + NGV_PER_BLOCK*4 + z_in_tile; + double dm_fac = dm[ij_offset + i*nao+j] * fac; + double *gxR = gx; + double *gxI = gxR + NGV_PER_BLOCK; + double yR0 = gxR[addry]; + double yI0 = gxI[addry]; + double zR0 = gxR[addrz]; + double zI0 = gxI[addrz]; + double yzR00, yzI00; + multiply(yR0, yI0, zR0, zI0, yzR00, yzI00); + yzR00 *= dm_fac; + yzI00 *= dm_fac; + + double ai2 = ai * -2; + double aj2 = aj * -2; + double yR3, yI3; + dIdJ_gx(gxR, addry, stride_i, stride_j, iy, jy, ai2, aj2, yR3, yI3); + double yzR33, yzI33; + multiply(yR3, yI3, zR0, zI0, yzR33, yzI33); + + double zR3, zI3; + dIdJ_gx(gxR, addrz, stride_i, stride_j, iz, jz, ai2, aj2, zR3, zI3); + double tmpR, tmpI; + multiply(yR0, yI0, zR3, zI3, tmpR, tmpI); + yzR33 = (yzR33 + tmpR) * dm_fac; + yzI33 = (yzI33 + tmpI) * dm_fac; +#pragma unroll + for (int n = 0; n < DENSITY_WIDTH; ++n) { + int x = n; + if (mesh_start[0] + x >= mesh_x) break; + int addr = addrx + x; + double xR3, xI3; + dIdJ_gx(gxR, addr, stride_i, stride_j, ix, jx, ai2, aj2, xR3, xI3); + double xyzR, xyzI; + multiply(xR3, xI3, yzR00, yzI00, xyzR, xyzI); + tau_R[n] += xyzR; + tau_I[n] += xyzI; + + double xR0 = gxR[addr]; + double xI0 = gxI[addr]; + multiply(xR0, xI0, yzR00, yzI00, xyzR, xyzI); + rho_R[n] += xyzR; + rho_I[n] += xyzI; + + multiply(xR0, xI0, yzR33, yzI33, xyzR, xyzI); + tau_R[n] += xyzR; + tau_I[n] += xyzI; + + } + } } + } + } + + int Gx0 = mesh_start[0]; + int Gy0 = mesh_start[1]; + int Gz0 = mesh_start[2]; + int y_in_tile = thread_id / NGV_PER_BLOCK; + int z_in_tile = Gv_id; + int y = Gy0 + y_in_tile; + int z = Gz0 + z_in_tile; + if (y < mesh_y && z < mesh_z) { +#pragma unroll + for (int n = 0; n < DENSITY_WIDTH; ++n) { + int x = Gx0 + n; + if (x >= mesh_x) break; + int abc_idx = (x * mesh_y + y) * mesh_z + z; + atomicAdd(densityR+abc_idx, rho_R[n]); + atomicAdd(densityI+abc_idx, rho_I[n]); + atomicAdd(tauR+abc_idx, tau_R[n]/2); + atomicAdd(tauI+abc_idx, tau_I[n]/2); + } + } +} + +extern "C" { +int orth_contract_ft_tau_dm(double *densityR, double *densityI, + double *tauR, double *tauI, double *dm, + PBCIntEnvVars *envs, int *shl_pair_offsets, + int64_t *bas_ij_idx, double *G_bases, double *L_bases, + int *mesh_cum, int *nimgs_cum, int *mesh, + int nbatches_shl_pair, double factor) +{ + int mesh_x = mesh[0]; + int mesh_y = mesh[1]; + int mesh_z = mesh[2]; + int ntiles_x = (mesh_x + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; + int ntiles_y = (mesh_y + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; + int ntiles_z = (mesh_z + NGV_PER_BLOCK - 1) / NGV_PER_BLOCK; + int ntiles = ntiles_x * ntiles_y * ntiles_z; + orth_ft_tau_dm_kernel<<>>( + densityR, densityI, tauR, tauI, dm, *envs, shl_pair_offsets, bas_ij_idx, G_bases, L_bases, + mesh_cum, nimgs_cum, ntiles, factor); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in orth_ft_tau_dm_kernel: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} +} diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/aft_recursion.cuh b/gpu4pyscf/lib/multigrid/multigrid_v3/aft_recursion.cuh new file mode 100644 index 000000000..7d8ff37b7 --- /dev/null +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/aft_recursion.cuh @@ -0,0 +1,252 @@ +/* + * Copyright 2026 The PySCF Developers. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#define NGV_PER_BLOCK 16 + +__forceinline__ __device__ +void vrr_hrr(double *gx, double *swap, int addrR, int li, int lj, + int stride_j, double a2, double xjxi, double aj_aij, + double xi, double kx, double theta_rr) +{ + constexpr int stride_i = NGV_PER_BLOCK * 6; + int addrI = addrR + NGV_PER_BLOCK; + int lij = li + lj; + double xpa = xjxi * aj_aij; + double xij = xpa + xi; + double kR = kx * xij; + double s0xR, s1xR, s2xR; + double s0xI, s1xI, s2xI; + sincos(-kR, &s0xI, &s0xR); + double Kab = exp(-theta_rr); + s0xR *= Kab; + s0xI *= Kab; + swap[addrR] = s0xR; + swap[addrI] = s0xI; + gx[addrR] += s0xR; + gx[addrI] += s0xI; + if (lij > 0) { + double RpaR = xpa; + double RpaI = -a2 * kx; + s1xR = RpaR * s0xR - RpaI * s0xI; + s1xI = RpaR * s0xI + RpaI * s0xR; + swap[addrR+stride_i] = s1xR; + swap[addrI+stride_i] = s1xI; + if (0 < li) { + gx[addrR+stride_i] += s1xR; + gx[addrI+stride_i] += s1xI; + } + for (int i = 2; i <= lij; i++) { + double ia2 = (i-1) * a2; + s2xR = ia2 * s0xR + RpaR * s1xR - RpaI * s1xI; + s2xI = ia2 * s0xI + RpaR * s1xI + RpaI * s1xR; + swap[addrR+i*stride_i] = s2xR; + swap[addrI+i*stride_i] = s2xI; + if (i <= li) { + gx[addrR+i*stride_i] += s2xR; + gx[addrI+i*stride_i] += s2xI; + } + s0xR = s1xR; + s0xI = s1xI; + s1xR = s2xR; + s1xI = s2xI; + } + } + for (int j = 1; j <= lj; ++j) { + int i = lij - j; + s1xR = swap[addrR+(i+1)*stride_i]; + s1xI = swap[addrI+(i+1)*stride_i]; + for (; i >= 0; --i) { + s0xR = swap[addrR+i*stride_i]; + s0xI = swap[addrI+i*stride_i]; + s2xR = s1xR - xjxi * s0xR; + s2xI = s1xI - xjxi * s0xI; + swap[addrR+i*stride_i] = s2xR; + swap[addrI+i*stride_i] = s2xI; + if (i <= li) { + int ij = i * stride_i + j * stride_j; + gx[addrR+ij] += s2xR; + gx[addrI+ij] += s2xI; + } + s1xR = s0xR; + s1xI = s0xI; + } + } +} + +template __forceinline__ __device__ +void vrr_hrr(double *gx, int addrR, int stride_j, double a2, double xjxi, + double RpaR, double RpaI, double g00R, double g00I) +{ + constexpr int stride_i = NGV_PER_BLOCK * 6; + int addrI = addrR + NGV_PER_BLOCK; + double swapR[LI+LJ+1]; + double swapI[LI+LJ+1]; + double s0xR, s1xR, s2xR; + double s0xI, s1xI, s2xI; + s0xR = g00R; + s0xI = g00I; + swapR[0] = s0xR; + swapI[0] = s0xI; + gx[addrR] += s0xR; + gx[addrI] += s0xI; + constexpr int lij = LI + LJ; + if (lij > 0) { + s1xR = RpaR * s0xR - RpaI * s0xI; + s1xI = RpaR * s0xI + RpaI * s0xR; + swapR[1] = s1xR; + swapI[1] = s1xI; + if (0 < LI) { + gx[addrR+stride_i] += s1xR; + gx[addrI+stride_i] += s1xI; + } +#pragma unroll + for (int i = 2; i <= lij; i++) { + double ia2 = (i-1) * a2; + s2xR = ia2 * s0xR + RpaR * s1xR - RpaI * s1xI; + s2xI = ia2 * s0xI + RpaR * s1xI + RpaI * s1xR; + swapR[i] = s2xR; + swapI[i] = s2xI; + if (i <= LI) { + gx[addrR+i*stride_i] += s2xR; + gx[addrI+i*stride_i] += s2xI; + } + s0xR = s1xR; + s0xI = s1xI; + s1xR = s2xR; + s1xI = s2xI; + } + } +#pragma unroll + for (int j = 1; j <= LJ; ++j) { + int i = lij - j; + s1xR = swapR[i+1]; + s1xI = swapI[i+1]; +#pragma unroll + for (; i >= 0; --i) { + s0xR = swapR[i]; + s0xI = swapI[i]; + s2xR = s1xR - xjxi * s0xR; + s2xI = s1xI - xjxi * s0xI; + swapR[i] = s2xR; + swapI[i] = s2xI; + if (i <= LI) { + int ij = i * stride_i + j * stride_j; + gx[addrR+ij] += s2xR; + gx[addrI+ij] += s2xI; + } + s1xR = s0xR; + s1xI = s0xI; + } + } +} + +__forceinline__ __device__ +void dI_gx(double *gx, int addr, int stride_i, int li, + double ai2, double &outR, double &outI) +{ + outR = ai2 * gx[addr+stride_i]; + outI = ai2 * gx[addr+stride_i+NGV_PER_BLOCK]; + if (li > 0) { + outR += li * gx[addr-stride_i]; + outI += li * gx[addr-stride_i+NGV_PER_BLOCK]; + } +} + +__forceinline__ __device__ +void dIdJ_gx(double *gx, int addr, int stride_i, int stride_j, int li, int lj, + double ai2, double aj2, double &outR, double &outI) +{ + outR = ai2 * gx[addr+stride_i+stride_j]; + outI = ai2 * gx[addr+stride_i+stride_j+NGV_PER_BLOCK]; + if (li > 0) { + outR += li * gx[addr-stride_i+stride_j]; + outI += li * gx[addr-stride_i+stride_j+NGV_PER_BLOCK]; + } + outR *= aj2; + outI *= aj2; + if (lj > 0) { + double f1R = ai2 * gx[addr+stride_i-stride_j]; + double f1I = ai2 * gx[addr+stride_i-stride_j+NGV_PER_BLOCK]; + if (li > 0) { + f1R += li * gx[addr-stride_i-stride_j]; + f1I += li * gx[addr-stride_i-stride_j+NGV_PER_BLOCK]; + } + outR += lj * f1R; + outI += lj * f1I; + } +} + +__forceinline__ __device__ +void dIdJ_gx(double *gx, int addr, int stride_i, int li, + double ai2, double kx, double &outR, double &outI) +{ + int li2 = li * 2 + 1; + outR = li2 * gx[addr]; + outI = li2 * gx[addr+NGV_PER_BLOCK]; + outR += ai2 * gx[addr+stride_i*2]; + outI += ai2 * gx[addr+stride_i*2+NGV_PER_BLOCK]; + outR *= ai2; + outI *= ai2; + if (li > 1) { + int lili = li * (li-1); + outR += lili * gx[addr-stride_i*2]; + outI += lili * gx[addr-stride_i*2+NGV_PER_BLOCK]; + } + + // (d i|d j) + (d2 i| j) + -iG*(d i|j) = 0 + double f1R = ai2 * gx[addr+stride_i]; + double f1I = ai2 * gx[addr+stride_i+NGV_PER_BLOCK]; + if (li > 0) { + f1R += li * gx[addr-stride_i]; + f1I += li * gx[addr-stride_i+NGV_PER_BLOCK]; + } + outR = f1I * -kx - outR; + outI = f1R * kx - outI; +} + +__forceinline__ __device__ +void d2IdJ_gx(double *gx, int addr, int stride_i, int stride_j, int li, int lj, + double ai2, double aj2, double &outR, double &outI) +{ + int li2 = li * 2 + 1; + int lili = li * (li-1); + outR = li2 * gx[addr+stride_j]; + outI = li2 * gx[addr+stride_j+NGV_PER_BLOCK]; + outR += ai2 * gx[addr+stride_i*2+stride_j]; + outI += ai2 * gx[addr+stride_i*2+stride_j+NGV_PER_BLOCK]; + outR *= ai2; + outI *= ai2; + if (li > 1) { + outR += lili * gx[addr-stride_i*2+stride_j]; + outI += lili * gx[addr-stride_i*2+stride_j+NGV_PER_BLOCK]; + } + outR *= aj2; + outI *= aj2; + if (lj > 0) { + double f1R = li2 * gx[addr-stride_j]; + double f1I = li2 * gx[addr-stride_j+NGV_PER_BLOCK]; + f1R += ai2 * gx[addr+stride_i*2-stride_j]; + f1I += ai2 * gx[addr+stride_i*2-stride_j+NGV_PER_BLOCK]; + f1R *= ai2; + f1I *= ai2; + if (li > 1) { + f1R += lili * gx[addr-stride_i*2-stride_j]; + f1I += lili * gx[addr-stride_i*2-stride_j+NGV_PER_BLOCK]; + } + outR += lj * f1R; + outI += lj * f1I; + } +} diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/cartesian.cuh b/gpu4pyscf/lib/multigrid/multigrid_v3/cartesian.cuh new file mode 100644 index 000000000..29ff2009f --- /dev/null +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/cartesian.cuh @@ -0,0 +1,613 @@ +/* + * Copyright 2025-2026 The PySCF Developers. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +__forceinline__ __device__ +void rename_registers(double i_cartesian[], int start, int stop, int count) +{ + if (start > 0) { +#pragma unroll + for (int i = 0; i < min(stop-start, count); ++i) { + i_cartesian[i] = i_cartesian[i+start]; + } + } +} + +template __forceinline__ __device__ +void gto_cartesian(double values[], double fx, double fy, double fz) +{ + if constexpr (ANG == 0) { + values[0] = 1; + } else if constexpr (ANG == 1) { + values[0] = fx; + values[1] = fy; + values[2] = fz; + } else if constexpr (ANG == 2) { + values[0] = fx * fx; + values[1] = fx * fy; + values[2] = fx * fz; + values[3] = fy * fy; + values[4] = fy * fz; + values[5] = fz * fz; + } else if constexpr (ANG == 3) { + double xx = fx * fx; + double yy = fy * fy; + double zz = fz * fz; + values[0] = xx * fx; + values[1] = xx * fy; + values[2] = xx * fz; + values[3] = fx * yy; + values[4] = fx * fy * fz; + values[5] = fx * zz; + values[6] = yy * fy; + values[7] = yy * fz; + values[8] = fy * zz; + values[9] = fz * zz; + } else if constexpr (ANG == 4) { + double xx = fx * fx; + double yy = fy * fy; + double zz = fz * fz; + double xxx = xx * fx; + double yyy = yy * fy; + double zzz = zz * fz; + values[0 ] = xxx * fx; + values[1 ] = xxx * fy; + values[2 ] = xxx * fz; + values[3 ] = xx * yy; + values[4 ] = xx * fy * fz; + values[5 ] = xx * zz; + values[6 ] = fx * yyy; + values[7 ] = fx * yy * fz; + values[8 ] = fx * fy * zz; + values[9 ] = fx * zzz; + values[10] = yyy * fy; + values[11] = yyy * fz; + values[12] = yy * zz; + values[13] = fy * zzz; + values[14] = fz * zzz; + } +} + +template __forceinline__ __device__ +void gto_deriv1(double gradient_values[], double original_values[], + double fx, double fy, double fz, double exponent) +{ + double a2 = -2 * exponent; + double minus_2afx = a2 * fx; + double minus_2afy = a2 * fy; + double minus_2afz = a2 * fz; + if constexpr (ANG == 0) { + // For s orbital (ANG=0), f(x,y,z) = 1 + // f'_x = 0, so g_x = -2 * exponent * fx + gradient_values[0] = minus_2afx; // x gradient + gradient_values[1] = minus_2afy; // y gradient + gradient_values[2] = minus_2afz; // z gradient + } else if constexpr (ANG == 1) { + // For p orbitals (ANG=1), f(x,y,z) = {x, y, z} + // First row: x gradient + gradient_values[0] = 1 + minus_2afx * fx; // d/dx(x) - 2*exponent*x*fx + gradient_values[1] = minus_2afx * fy; // d/dx(y) - 2*exponent*y*fx + gradient_values[2] = minus_2afx * fz; // d/dx(z) - 2*exponent*z*fx + // Second row: y gradient + gradient_values[3] = minus_2afy * fx; // d/dy(x) - 2*exponent*x*fy + gradient_values[4] = 1 + minus_2afy * fy; // d/dy(y) - 2*exponent*y*fy + gradient_values[5] = minus_2afy * fz; // d/dy(z) - 2*exponent*z*fy + // Third row: z gradient + gradient_values[6] = minus_2afz * fx; // d/dz(x) - 2*exponent*x*fz + gradient_values[7] = minus_2afz * fy; // d/dz(y) - 2*exponent*y*fz + gradient_values[8] = 1 + minus_2afz * fz; // d/dz(z) - 2*exponent*z*fz + } else if constexpr (ANG == 2) { + // For d orbitals (ANG=2), f(x,y,z) = {xx, xy, xz, yy, yz, zz} + // First row: x gradient + gradient_values[0] = + 2 * fx + minus_2afx * original_values[0]; // d/dx(xx) - 2*exponent*xx*fx + gradient_values[1] = + fy + minus_2afx * original_values[1]; // d/dx(xy) - 2*exponent*xy*fx + gradient_values[2] = + fz + minus_2afx * original_values[2]; // d/dx(xz) - 2*exponent*xz*fx + gradient_values[3] = + minus_2afx * original_values[3]; // d/dx(yy) - 2*exponent*yy*fx + gradient_values[4] = + minus_2afx * original_values[4]; // d/dx(yz) - 2*exponent*yz*fx + gradient_values[5] = + minus_2afx * original_values[5]; // d/dx(zz) - 2*exponent*zz*fx + // Second row: y gradient + gradient_values[6] = + minus_2afy * original_values[0]; // d/dy(xx) - 2*exponent*xx*fy + gradient_values[7] = + fx + minus_2afy * original_values[1]; // d/dy(xy) - 2*exponent*xy*fy + gradient_values[8] = + minus_2afy * original_values[2]; // d/dy(xz) - 2*exponent*xz*fy + gradient_values[9] = + 2 * fy + minus_2afy * original_values[3]; // d/dy(yy) - 2*exponent*yy*fy + gradient_values[10] = + fz + minus_2afy * original_values[4]; // d/dy(yz) - 2*exponent*yz*fy + gradient_values[11] = + minus_2afy * original_values[5]; // d/dy(zz) - 2*exponent*zz*fy + // Third row: z gradient + gradient_values[12] = + minus_2afz * original_values[0]; // d/dz(xx) - 2*exponent*xx*fz + gradient_values[13] = + minus_2afz * original_values[1]; // d/dz(xy) - 2*exponent*xy*fz + gradient_values[14] = + fx + minus_2afz * original_values[2]; // d/dz(xz) - 2*exponent*xz*fz + gradient_values[15] = + minus_2afz * original_values[3]; // d/dz(yy) - 2*exponent*yy*fz + gradient_values[16] = + fy + minus_2afz * original_values[4]; // d/dz(yz) - 2*exponent*yz*fz + gradient_values[17] = + 2 * fz + minus_2afz * original_values[5]; // d/dz(zz) - 2*exponent*zz*fz + } else if constexpr (ANG == 3) { + // For f orbitals (ANG=3), f(x,y,z) = {xxx, xxy, xxz, xyy, xyz, xzz, yyy, + // yyz, yzz, zzz} + // First row: x gradient + gradient_values[0] = + 3 * fx * fx + + minus_2afx * original_values[0]; // d/dx(xxx) - 2*exponent*xxx*fx + gradient_values[1] = + 2 * fx * fy + + minus_2afx * original_values[1]; // d/dx(xxy) - 2*exponent*xxy*fx + gradient_values[2] = + 2 * fx * fz + + minus_2afx * original_values[2]; // d/dx(xxz) - 2*exponent*xxz*fx + gradient_values[3] = + fy * fy + + minus_2afx * original_values[3]; // d/dx(xyy) - 2*exponent*xyy*fx + gradient_values[4] = + fy * fz + + minus_2afx * original_values[4]; // d/dx(xyz) - 2*exponent*xyz*fx + gradient_values[5] = + fz * fz + + minus_2afx * original_values[5]; // d/dx(xzz) - 2*exponent*xzz*fx + gradient_values[6] = + minus_2afx * original_values[6]; // d/dx(yyy) - 2*exponent*yyy*fx + gradient_values[7] = + minus_2afx * original_values[7]; // d/dx(yyz) - 2*exponent*yyz*fx + gradient_values[8] = + minus_2afx * original_values[8]; // d/dx(yzz) - 2*exponent*yzz*fx + gradient_values[9] = + minus_2afx * original_values[9]; // d/dx(zzz) - 2*exponent*zzz*fx + // Second row: y gradient + gradient_values[10] = + minus_2afy * original_values[0]; // d/dy(xxx) - 2*exponent*xxx*fy + gradient_values[11] = + fx * fx + + minus_2afy * original_values[1]; // d/dy(xxy) - 2*exponent*xxy*fy + gradient_values[12] = + minus_2afy * original_values[2]; // d/dy(xxz) - 2*exponent*xxz*fy + gradient_values[13] = + 2 * fx * fy + + minus_2afy * original_values[3]; // d/dy(xyy) - 2*exponent*xyy*fy + gradient_values[14] = + fx * fz + + minus_2afy * original_values[4]; // d/dy(xyz) - 2*exponent*xyz*fy + gradient_values[15] = + minus_2afy * original_values[5]; // d/dy(xzz) - 2*exponent*xzz*fy + gradient_values[16] = + 3 * fy * fy + + minus_2afy * original_values[6]; // d/dy(yyy) - 2*exponent*yyy*fy + gradient_values[17] = + 2 * fy * fz + + minus_2afy * original_values[7]; // d/dy(yyz) - 2*exponent*yyz*fy + gradient_values[18] = + fz * fz + + minus_2afy * original_values[8]; // d/dy(yzz) - 2*exponent*yzz*fy + gradient_values[19] = + minus_2afy * original_values[9]; // d/dy(zzz) - 2*exponent*zzz*fy + // Third row: z gradient + gradient_values[20] = + minus_2afz * original_values[0]; // d/dz(xxx) - 2*exponent*xxx*fz + gradient_values[21] = + minus_2afz * original_values[1]; // d/dz(xxy) - 2*exponent*xxy*fz + gradient_values[22] = + fx * fx + + minus_2afz * original_values[2]; // d/dz(xxz) - 2*exponent*xxz*fz + gradient_values[23] = + minus_2afz * original_values[3]; // d/dz(xyy) - 2*exponent*xyy*fz + gradient_values[24] = + fx * fy + + minus_2afz * original_values[4]; // d/dz(xyz) - 2*exponent*xyz*fz + gradient_values[25] = + 2 * fx * fz + + minus_2afz * original_values[5]; // d/dz(xzz) - 2*exponent*xzz*fz + gradient_values[26] = + minus_2afz * original_values[6]; // d/dz(yyy) - 2*exponent*yyy*fz + gradient_values[27] = + fy * fy + + minus_2afz * original_values[7]; // d/dz(yyz) - 2*exponent*yyz*fz + gradient_values[28] = + 2 * fy * fz + + minus_2afz * original_values[8]; // d/dz(yzz) - 2*exponent*yzz*fz + gradient_values[29] = + 3 * fz * fz + + minus_2afz * original_values[9]; // d/dz(zzz) - 2*exponent*zzz*fz + } else if constexpr (ANG == 4) { + // For g orbitals (ANG=4), f(x,y,z) = {xxxx, xxxy, xxxz, xxyy, xxyz, xxzz, + // xyyy, xyyz, xyzz, xzzz, yyyy, yyyz, yyzz, yzzz, zzzz} + // First row: x gradient + gradient_values[0] = + 4 * fx * fx * fx + + minus_2afx * original_values[0]; // d/dx(xxxx) - 2*exponent*xxxx*fx + gradient_values[1] = + 3 * fx * fx * fy + + minus_2afx * original_values[1]; // d/dx(xxxy) - 2*exponent*xxxy*fx + gradient_values[2] = + 3 * fx * fx * fz + + minus_2afx * original_values[2]; // d/dx(xxxz) - 2*exponent*xxxz*fx + gradient_values[3] = + 2 * fx * fy * fy + + minus_2afx * original_values[3]; // d/dx(xxyy) - 2*exponent*xxyy*fx + gradient_values[4] = + 2 * fx * fy * fz + + minus_2afx * original_values[4]; // d/dx(xxyz) - 2*exponent*xxyz*fx + gradient_values[5] = + 2 * fx * fz * fz + + minus_2afx * original_values[5]; // d/dx(xxzz) - 2*exponent*xxzz*fx + gradient_values[6] = + fy * fy * fy + + minus_2afx * original_values[6]; // d/dx(xyyy) - 2*exponent*xyyy*fx + gradient_values[7] = + fy * fy * fz + + minus_2afx * original_values[7]; // d/dx(xyyz) - 2*exponent*xyyz*fx + gradient_values[8] = + fy * fz * fz + + minus_2afx * original_values[8]; // d/dx(xyzz) - 2*exponent*xyzz*fx + gradient_values[9] = + fz * fz * fz + + minus_2afx * original_values[9]; // d/dx(xzzz) - 2*exponent*xzzz*fx + gradient_values[10] = + minus_2afx * original_values[10]; // d/dx(yyyy) - 2*exponent*yyyy*fx + gradient_values[11] = + minus_2afx * original_values[11]; // d/dx(yyyz) - 2*exponent*yyyz*fx + gradient_values[12] = + minus_2afx * original_values[12]; // d/dx(yyzz) - 2*exponent*yyzz*fx + gradient_values[13] = + minus_2afx * original_values[13]; // d/dx(yzzz) - 2*exponent*yzzz*fx + gradient_values[14] = + minus_2afx * original_values[14]; // d/dx(zzzz) - 2*exponent*zzzz*fx + // Second row: y gradient + gradient_values[15] = + minus_2afy * original_values[0]; // d/dy(xxxx) - 2*exponent*xxxx*fy + gradient_values[16] = + fx * fx * fx + + minus_2afy * original_values[1]; // d/dy(xxxy) - 2*exponent*xxxy*fy + gradient_values[17] = + minus_2afy * original_values[2]; // d/dy(xxxz) - 2*exponent*xxxz*fy + gradient_values[18] = + 2 * fx * fx * fy + + minus_2afy * original_values[3]; // d/dy(xxyy) - 2*exponent*xxyy*fy + gradient_values[19] = + fx * fx * fz + + minus_2afy * original_values[4]; // d/dy(xxyz) - 2*exponent*xxyz*fy + gradient_values[20] = + minus_2afy * original_values[5]; // d/dy(xxzz) - 2*exponent*xxzz*fy + gradient_values[21] = + 3 * fx * fy * fy + + minus_2afy * original_values[6]; // d/dy(xyyy) - 2*exponent*xyyy*fy + gradient_values[22] = + 2 * fx * fy * fz + + minus_2afy * original_values[7]; // d/dy(xyyz) - 2*exponent*xyyz*fy + gradient_values[23] = + fx * fz * fz + + minus_2afy * original_values[8]; // d/dy(xyzz) - 2*exponent*xyzz*fy + gradient_values[24] = + minus_2afy * original_values[9]; // d/dy(xzzz) - 2*exponent*xzzz*fy + gradient_values[25] = + 4 * fy * fy * fy + + minus_2afy * original_values[10]; // d/dy(yyyy) - 2*exponent*yyyy*fy + gradient_values[26] = + 3 * fy * fy * fz + + minus_2afy * original_values[11]; // d/dy(yyyz) - 2*exponent*yyyz*fy + gradient_values[27] = + 2 * fy * fz * fz + + minus_2afy * original_values[12]; // d/dy(yyzz) - 2*exponent*yyzz*fy + gradient_values[28] = + fz * fz * fz + + minus_2afy * original_values[13]; // d/dy(yzzz) - 2*exponent*yzzz*fy + gradient_values[29] = + minus_2afy * original_values[14]; // d/dy(zzzz) - 2*exponent*zzzz*fy + // Third row: z gradient + gradient_values[30] = + minus_2afz * original_values[0]; // d/dz(xxxx) - 2*exponent*xxxx*fz + gradient_values[31] = + minus_2afz * original_values[1]; // d/dz(xxxy) - 2*exponent*xxxy*fz + gradient_values[32] = + fx * fx * fx + + minus_2afz * original_values[2]; // d/dz(xxxz) - 2*exponent*xxxz*fz + gradient_values[33] = + minus_2afz * original_values[3]; // d/dz(xxyy) - 2*exponent*xxyy*fz + gradient_values[34] = + fx * fx * fy + + minus_2afz * original_values[4]; // d/dz(xxyz) - 2*exponent*xxyz*fz + gradient_values[35] = + 2 * fx * fx * fz + + minus_2afz * original_values[5]; // d/dz(xxzz) - 2*exponent*xxzz*fz + gradient_values[36] = + minus_2afz * original_values[6]; // d/dz(xyyy) - 2*exponent*xyyy*fz + gradient_values[37] = + fx * fy * fy + + minus_2afz * original_values[7]; // d/dz(xyyz) - 2*exponent*xyyz*fz + gradient_values[38] = + 2 * fx * fy * fz + + minus_2afz * original_values[8]; // d/dz(xyzz) - 2*exponent*xyzz*fz + gradient_values[39] = + 3 * fx * fz * fz + + minus_2afz * original_values[9]; // d/dz(xzzz) - 2*exponent*xzzz*fz + gradient_values[40] = + minus_2afz * original_values[10]; // d/dz(yyyy) - 2*exponent*yyyy*fz + gradient_values[41] = + fy * fy * fy + + minus_2afz * original_values[11]; // d/dz(yyyz) - 2*exponent*yyyz*fz + gradient_values[42] = + 2 * fy * fy * fz + + minus_2afz * original_values[12]; // d/dz(yyzz) - 2*exponent*yyzz*fz + gradient_values[43] = + 3 * fy * fz * fz + + minus_2afz * original_values[13]; // d/dz(yzzz) - 2*exponent*yzzz*fz + gradient_values[44] = + 4 * fz * fz * fz + + minus_2afz * original_values[14]; // d/dz(zzzz) - 2*exponent*zzzz*fz + } +} + +template __forceinline__ __device__ +void gto_deriv2(double output[], double x, double y, double z, double exponent) +{ + // Output in (derivative, orbital) + // Where derivative in the order of xx, xy, xz, yy, yz, zz + // And orbital in the order of x,y,z; xx,xy,xz,yy,yz,zz; xxx,xxy,xxz,xyy,xyz,xzz,yyy,yyz,yzz,zzz; etc. + + double aa = exponent * 2; + double aa2 = aa * aa; + double x2 = x * x; + double y2 = y * y; + double z2 = z * z; + + if constexpr (ANG == 0) { + output[0] = aa*(aa*x2 - 1); + output[1] = aa2*x*y; + output[2] = aa2*x*z; + output[3] = aa*(aa*y2 - 1); + output[4] = aa2*y*z; + output[5] = aa*(aa*z2 - 1); + } else if constexpr (ANG == 1) { + output[ 0] = aa*x*(aa*x2 - 3); + output[ 1] = aa*y*(aa*x2 - 1); + output[ 2] = aa*z*(aa*x2 - 1); + output[ 3] = aa*y*(aa*x2 - 1); + output[ 4] = aa*x*(aa*y2 - 1); + output[ 5] = aa2*x*y*z; + output[ 6] = aa*z*(aa*x2 - 1); + output[ 7] = aa2*x*y*z; + output[ 8] = aa*x*(aa*z2 - 1); + output[ 9] = aa*x*(aa*y2 - 1); + output[10] = aa*y*(aa*y2 - 3); + output[11] = aa*z*(aa*y2 - 1); + output[12] = aa2*x*y*z; + output[13] = aa*z*(aa*y2 - 1); + output[14] = aa*y*(aa*z2 - 1); + output[15] = aa*x*(aa*z2 - 1); + output[16] = aa*y*(aa*z2 - 1); + output[17] = aa*z*(aa*z2 - 3); + } else if constexpr (ANG == 2) { + double x4 = x2 * x2; + double y4 = y2 * y2; + double z4 = z2 * z2; + output[ 0] = aa2*x4 - 5*aa*x2 + 2; + output[ 1] = aa*x*y*(aa*x2 - 3); + output[ 2] = aa*x*z*(aa*x2 - 3); + output[ 3] = aa*y2*(aa*x2 - 1); + output[ 4] = aa*y*z*(aa*x2 - 1); + output[ 5] = aa*z2*(aa*x2 - 1); + output[ 6] = aa*x*y*(aa*x2 - 2); + output[ 7] = (aa*x2 - 1)*(aa*y2 - 1); + output[ 8] = aa*y*z*(aa*x2 - 1); + output[ 9] = aa*x*y*(aa*y2 - 2); + output[10] = aa*x*z*(aa*y2 - 1); + output[11] = aa2*x*y*z2; + output[12] = aa*x*z*(aa*x2 - 2); + output[13] = aa*y*z*(aa*x2 - 1); + output[14] = (aa*x2 - 1)*(aa*z2 - 1); + output[15] = aa2*x*y2*z; + output[16] = aa*x*y*(aa*z2 - 1); + output[17] = aa*x*z*(aa*z2 - 2); + output[18] = aa*x2*(aa*y2 - 1); + output[19] = aa*x*y*(aa*y2 - 3); + output[20] = aa*x*z*(aa*y2 - 1); + output[21] = aa2*y4 - 5*aa*y2 + 2; + output[22] = aa*y*z*(aa*y2 - 3); + output[23] = aa*z2*(aa*y2 - 1); + output[24] = aa2*x2*y*z; + output[25] = aa*x*z*(aa*y2 - 1); + output[26] = aa*x*y*(aa*z2 - 1); + output[27] = aa*y*z*(aa*y2 - 2); + output[28] = (aa*y2 - 1)*(aa*z2 - 1); + output[29] = aa*y*z*(aa*z2 - 2); + output[30] = aa*x2*(aa*z2 - 1); + output[31] = aa*x*y*(aa*z2 - 1); + output[32] = aa*x*z*(aa*z2 - 3); + output[33] = aa*y2*(aa*z2 - 1); + output[34] = aa*y*z*(aa*z2 - 3); + output[35] = aa2*z4 - 5*aa*z2 + 2; + } else if constexpr (ANG == 3) { + double x3 = x * x2; + double y3 = y * y2; + double z3 = z * z2; + double x4 = x2 * x2; + double y4 = y2 * y2; + double z4 = z2 * z2; + output[ 0] = x*(aa2*x4 - 7*aa*x2 + 6); + output[ 1] = y*(aa2*x4 - 5*aa*x2 + 2); + output[ 2] = z*(aa2*x4 - 5*aa*x2 + 2); + output[ 3] = aa*x*y2*(aa*x2 - 3); + output[ 4] = aa*x*y*z*(aa*x2 - 3); + output[ 5] = aa*x*z2*(aa*x2 - 3); + output[ 6] = aa*y3*(aa*x2 - 1); + output[ 7] = aa*y2*z*(aa*x2 - 1); + output[ 8] = aa*y*z2*(aa*x2 - 1); + output[ 9] = aa*z3*(aa*x2 - 1); + output[10] = aa*x2*y*(aa*x2 - 3); + output[11] = x*(aa*x2 - 2)*(aa*y2 - 1); + output[12] = aa*x*y*z*(aa*x2 - 2); + output[13] = y*(aa*x2 - 1)*(aa*y2 - 2); + output[14] = z*(aa*x2 - 1)*(aa*y2 - 1); + output[15] = aa*y*z2*(aa*x2 - 1); + output[16] = aa*x*y2*(aa*y2 - 3); + output[17] = aa*x*y*z*(aa*y2 - 2); + output[18] = aa*x*z2*(aa*y2 - 1); + output[19] = aa2*x*y*z3; + output[20] = aa*x2*z*(aa*x2 - 3); + output[21] = aa*x*y*z*(aa*x2 - 2); + output[22] = x*(aa*x2 - 2)*(aa*z2 - 1); + output[23] = aa*y2*z*(aa*x2 - 1); + output[24] = y*(aa*x2 - 1)*(aa*z2 - 1); + output[25] = z*(aa*x2 - 1)*(aa*z2 - 2); + output[26] = aa2*x*y3*z; + output[27] = aa*x*y2*(aa*z2 - 1); + output[28] = aa*x*y*z*(aa*z2 - 2); + output[29] = aa*x*z2*(aa*z2 - 3); + output[30] = aa*x3*(aa*y2 - 1); + output[31] = aa*x2*y*(aa*y2 - 3); + output[32] = aa*x2*z*(aa*y2 - 1); + output[33] = x*(aa2*y4 - 5*aa*y2 + 2); + output[34] = aa*x*y*z*(aa*y2 - 3); + output[35] = aa*x*z2*(aa*y2 - 1); + output[36] = y*(aa2*y4 - 7*aa*y2 + 6); + output[37] = z*(aa2*y4 - 5*aa*y2 + 2); + output[38] = aa*y*z2*(aa*y2 - 3); + output[39] = aa*z3*(aa*y2 - 1); + output[40] = aa2*x3*y*z; + output[41] = aa*x2*z*(aa*y2 - 1); + output[42] = aa*x2*y*(aa*z2 - 1); + output[43] = aa*x*y*z*(aa*y2 - 2); + output[44] = x*(aa*y2 - 1)*(aa*z2 - 1); + output[45] = aa*x*y*z*(aa*z2 - 2); + output[46] = aa*y2*z*(aa*y2 - 3); + output[47] = y*(aa*y2 - 2)*(aa*z2 - 1); + output[48] = z*(aa*y2 - 1)*(aa*z2 - 2); + output[49] = aa*y*z2*(aa*z2 - 3); + output[50] = aa*x3*(aa*z2 - 1); + output[51] = aa*x2*y*(aa*z2 - 1); + output[52] = aa*x2*z*(aa*z2 - 3); + output[53] = aa*x*y2*(aa*z2 - 1); + output[54] = aa*x*y*z*(aa*z2 - 3); + output[55] = x*(aa2*z4 - 5*aa*z2 + 2); + output[56] = aa*y3*(aa*z2 - 1); + output[57] = aa*y2*z*(aa*z2 - 3); + output[58] = y*(aa2*z4 - 5*aa*z2 + 2); + output[59] = z*(aa2*z4 - 7*aa*z2 + 6); + } else if constexpr (ANG == 4) { + double x3 = x * x2; + double y3 = y * y2; + double z3 = z * z2; + double x4 = x2 * x2; + double y4 = y2 * y2; + double z4 = z2 * z2; + output[ 0] = x2*(aa2*x4 - 9*aa*x2 + 12); + output[ 1] = x*y*(aa2*x4 - 7*aa*x2 + 6); + output[ 2] = x*z*(aa2*x4 - 7*aa*x2 + 6); + output[ 3] = y2*(aa2*x4 - 5*aa*x2 + 2); + output[ 4] = y*z*(aa2*x4 - 5*aa*x2 + 2); + output[ 5] = z2*(aa2*x4 - 5*aa*x2 + 2); + output[ 6] = aa*x*y3*(aa*x2 - 3); + output[ 7] = aa*x*y2*z*(aa*x2 - 3); + output[ 8] = aa*x*y*z2*(aa*x2 - 3); + output[ 9] = aa*x*z3*(aa*x2 - 3); + output[10] = aa*y4*(aa*x2 - 1); + output[11] = aa*y3*z*(aa*x2 - 1); + output[12] = aa*y2*z2*(aa*x2 - 1); + output[13] = aa*y*z3*(aa*x2 - 1); + output[14] = aa*z4*(aa*x2 - 1); + output[15] = aa*x3*y*(aa*x2 - 4); + output[16] = x2*(aa*x2 - 3)*(aa*y2 - 1); + output[17] = aa*x2*y*z*(aa*x2 - 3); + output[18] = x*y*(aa*x2 - 2)*(aa*y2 - 2); + output[19] = x*z*(aa*x2 - 2)*(aa*y2 - 1); + output[20] = aa*x*y*z2*(aa*x2 - 2); + output[21] = y2*(aa*x2 - 1)*(aa*y2 - 3); + output[22] = y*z*(aa*x2 - 1)*(aa*y2 - 2); + output[23] = z2*(aa*x2 - 1)*(aa*y2 - 1); + output[24] = aa*y*z3*(aa*x2 - 1); + output[25] = aa*x*y3*(aa*y2 - 4); + output[26] = aa*x*y2*z*(aa*y2 - 3); + output[27] = aa*x*y*z2*(aa*y2 - 2); + output[28] = aa*x*z3*(aa*y2 - 1); + output[29] = aa2*x*y*z4; + output[30] = aa*x3*z*(aa*x2 - 4); + output[31] = aa*x2*y*z*(aa*x2 - 3); + output[32] = x2*(aa*x2 - 3)*(aa*z2 - 1); + output[33] = aa*x*y2*z*(aa*x2 - 2); + output[34] = x*y*(aa*x2 - 2)*(aa*z2 - 1); + output[35] = x*z*(aa*x2 - 2)*(aa*z2 - 2); + output[36] = aa*y3*z*(aa*x2 - 1); + output[37] = y2*(aa*x2 - 1)*(aa*z2 - 1); + output[38] = y*z*(aa*x2 - 1)*(aa*z2 - 2); + output[39] = z2*(aa*x2 - 1)*(aa*z2 - 3); + output[40] = aa2*x*y4*z; + output[41] = aa*x*y3*(aa*z2 - 1); + output[42] = aa*x*y2*z*(aa*z2 - 2); + output[43] = aa*x*y*z2*(aa*z2 - 3); + output[44] = aa*x*z3*(aa*z2 - 4); + output[45] = aa*x4*(aa*y2 - 1); + output[46] = aa*x3*y*(aa*y2 - 3); + output[47] = aa*x3*z*(aa*y2 - 1); + output[48] = x2*(aa2*y4 - 5*aa*y2 + 2); + output[49] = aa*x2*y*z*(aa*y2 - 3); + output[50] = aa*x2*z2*(aa*y2 - 1); + output[51] = x*y*(aa2*y4 - 7*aa*y2 + 6); + output[52] = x*z*(aa2*y4 - 5*aa*y2 + 2); + output[53] = aa*x*y*z2*(aa*y2 - 3); + output[54] = aa*x*z3*(aa*y2 - 1); + output[55] = y2*(aa2*y4 - 9*aa*y2 + 12); + output[56] = y*z*(aa2*y4 - 7*aa*y2 + 6); + output[57] = z2*(aa2*y4 - 5*aa*y2 + 2); + output[58] = aa*y*z3*(aa*y2 - 3); + output[59] = aa*z4*(aa*y2 - 1); + output[60] = aa2*x4*y*z; + output[61] = aa*x3*z*(aa*y2 - 1); + output[62] = aa*x3*y*(aa*z2 - 1); + output[63] = aa*x2*y*z*(aa*y2 - 2); + output[64] = x2*(aa*y2 - 1)*(aa*z2 - 1); + output[65] = aa*x2*y*z*(aa*z2 - 2); + output[66] = aa*x*y2*z*(aa*y2 - 3); + output[67] = x*y*(aa*y2 - 2)*(aa*z2 - 1); + output[68] = x*z*(aa*y2 - 1)*(aa*z2 - 2); + output[69] = aa*x*y*z2*(aa*z2 - 3); + output[70] = aa*y3*z*(aa*y2 - 4); + output[71] = y2*(aa*y2 - 3)*(aa*z2 - 1); + output[72] = y*z*(aa*y2 - 2)*(aa*z2 - 2); + output[73] = z2*(aa*y2 - 1)*(aa*z2 - 3); + output[74] = aa*y*z3*(aa*z2 - 4); + output[75] = aa*x4*(aa*z2 - 1); + output[76] = aa*x3*y*(aa*z2 - 1); + output[77] = aa*x3*z*(aa*z2 - 3); + output[78] = aa*x2*y2*(aa*z2 - 1); + output[79] = aa*x2*y*z*(aa*z2 - 3); + output[80] = x2*(aa2*z4 - 5*aa*z2 + 2); + output[81] = aa*x*y3*(aa*z2 - 1); + output[82] = aa*x*y2*z*(aa*z2 - 3); + output[83] = x*y*(aa2*z4 - 5*aa*z2 + 2); + output[84] = x*z*(aa2*z4 - 7*aa*z2 + 6); + output[85] = aa*y4*(aa*z2 - 1); + output[86] = aa*y3*z*(aa*z2 - 3); + output[87] = y2*(aa2*z4 - 5*aa*z2 + 2); + output[88] = y*z*(aa2*z4 - 7*aa*z2 + 6); + output[89] = z2*(aa2*z4 - 9*aa*z2 + 12); + } else { + output[0] = NAN; + } +} diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/constant_objects.cuh b/gpu4pyscf/lib/multigrid/multigrid_v3/constant_objects.cuh new file mode 100644 index 000000000..b36d71e00 --- /dev/null +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/constant_objects.cuh @@ -0,0 +1,25 @@ +/* + * Copyright 2025 The PySCF Developers. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +extern __constant__ double c_lattice_vectors[9]; +extern __constant__ double c_reciprocal_lattice_vectors[9]; // norm to 1 +extern __constant__ double c_dxyz_dabc[9]; +extern __constant__ int c_nf[]; +extern __constant__ float c_div_nf[]; + +#define NBAS_MAX 16777216 diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/drivers.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/drivers.cu new file mode 100644 index 000000000..ab8847dbb --- /dev/null +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/drivers.cu @@ -0,0 +1,153 @@ +/* + * Copyright 2025 The PySCF Developers. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include + +__constant__ double c_lattice_vectors[9]; +__constant__ double c_reciprocal_lattice_vectors[9]; +__constant__ double c_dxyz_dabc[9]; + +__constant__ int c_nf[] = { + 1, + 3, + 6, + 10, + 15, + 21, + 28, + 36, + 45, +}; + +__constant__ float c_div_nf[] = { + 1.f, + 0.333334f, + 0.166667f, + 0.100001f, + 0.066667f, + 0.047620f, + 0.035715f, + 0.027778f, + 0.022223f, +}; + +// input[nc,nx,ny,nz], output[nc,mx,my,mz] +__global__ static +void fft_take_kernel(double2* __restrict__ out, double2* __restrict__ in, + int mx, int my, int mz, int nx, int ny, int nz, int nc) +{ + int x = blockIdx.x; + int y = blockIdx.y * blockDim.y + threadIdx.y; + int tx = threadIdx.x; + int threadsx = blockDim.x; + if (x >= mx || y >= my) return; + + int sx = x; + int sy = y; + // fftfreq indexing + if (x > mx/2) sx = nx + x - mx; + if (y > my/2) sy = ny + y - my; + for (int z = tx; z < mz; z += threadsx) { + int sz = z; + if (z > mz/2) sz = nz + z - mz; + + for (int c = 0; c < nc; ++c) { + size_t src = (((size_t)c*nx + sx)*ny + sy)*nz + sz; + size_t dst = (((size_t)c*mx + x )*my + y )*mz + z; + out[dst] = in[src]; + } + } +} + +// output[nc,nx,ny,nz], input[nc,mx,my,mz] +__global__ static +void fft_takebak_kernel(double2* __restrict__ out, double2* __restrict__ in, + int mx, int my, int mz, int nx, int ny, int nz, int nc) +{ + int x = blockIdx.x; + int y = blockIdx.y * blockDim.y + threadIdx.y; + int tx = threadIdx.x; + int threadsx = blockDim.x; + if (x >= mx || y >= my) return; + + int sx = x; + int sy = y; + // fftfreq indexing + if (x > mx/2) sx = nx + x - mx; + if (y > my/2) sy = ny + y - my; + for (int z = tx; z < mz; z += threadsx) { + int sz = z; + if (z > mz/2) sz = nz + z - mz; + + for (int c = 0; c < nc; ++c) { + size_t dst = (((size_t)c*nx + sx)*ny + sy)*nz + sz; + size_t src = (((size_t)c*mx + x )*my + y )*mz + z; + out[dst].x += in[src].x; + out[dst].y += in[src].y; + } + } +} + +extern "C" { +void update_lattice_vectors(double *lattice_vectors, + double *reciprocal_lattice_vectors) +{ + cudaMemcpyToSymbol(c_lattice_vectors, lattice_vectors, 9 * sizeof(double)); + cudaMemcpyToSymbol(c_reciprocal_lattice_vectors, reciprocal_lattice_vectors, 9 * sizeof(double)); +} + +void update_dxyz_dabc(double *dxyz_dabc) { + cudaMemcpyToSymbol(c_dxyz_dabc, dxyz_dabc, 9 * sizeof(double)); +} + +int fft_take(double2 *out, double2 *in, int *out_shape, int *in_shape, int counts) +{ + int mx = out_shape[0]; + int my = out_shape[1]; + int mz = out_shape[2]; + dim3 threads(32, 16); + dim3 grids(mx, (my+15)/16); + fft_take_kernel<<>>( + out, in, mx, my, mz, in_shape[0], in_shape[1], in_shape[2], counts); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in fft_take kernel: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} + +int fft_takebak(double2 *out, double2 *in, int *out_shape, int *in_shape, int counts) +{ + int mx = in_shape[0]; + int my = in_shape[1]; + int mz = in_shape[2]; + dim3 threads(32, 16); + dim3 grids(mx, (my+15)/16); + fft_takebak_kernel<<>>( + out, in, mx, my, mz, out_shape[0], out_shape[1], out_shape[2], counts); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in fft_takebak kernel: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} +} diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/eval_density.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_density.cu new file mode 100644 index 000000000..2f893d7ac --- /dev/null +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_density.cu @@ -0,0 +1,385 @@ +/* + * Copyright 2025-2026 The PySCF Developers. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include "gvhf-rys/vhf.cuh" +#include "constant_objects.cuh" +#include "cartesian.cuh" +#include "utils.cuh" + +#define TILE 4 +#define WARP_SIZE 32 +#define THREADS 64 + +template +__global__ static +void eval_density_kernel(double *density, double *dm, PBCIntEnvVars envs, + double *supmol_img_coords, double factor, + int *shl_pair_offsets, int64_t *dressed_bas_ij_idx, + int *grid_tile_index, int ntiles, int tiles_per_block, + double a_dot_b, double a_dot_c, double b_dot_c, + double da_squared, double db_squared, double dc_squared, + int mesh_a, int mesh_b, int mesh_c, double negligible) +{ + constexpr int threads = THREADS; + constexpr int WARPS = THREADS / WARP_SIZE; + int thread_id = threadIdx.x; + int tile_id0 = blockIdx.x * tiles_per_block; + __shared__ int a_upper, b_upper, c_upper; + __shared__ double start_position_x, start_position_y, start_position_z; + __shared__ double density_value[TILE*TILE*TILE*WARPS]; + + constexpr int nfi = (LI + 1) * (LI + 2) / 2; + constexpr int nfj = (LJ + 1) * (LJ + 2) / 2; + +for (int tile_id = tile_id0; tile_id < min(tile_id0+tiles_per_block, ntiles); tile_id++) { + int tile_index = grid_tile_index[tile_id]; + int shl_pair0 = shl_pair_offsets[tile_id]; + int shl_pair1 = shl_pair_offsets[tile_id+1]; + int n_tiles_b = (mesh_b + TILE - 1) / TILE; + int n_tiles_c = (mesh_c + TILE - 1) / TILE; + int tile_ab_index = tile_index / n_tiles_c; + int tile_c_index = tile_index % n_tiles_c; + int tile_a_index = tile_ab_index / n_tiles_b; + int tile_b_index = tile_ab_index % n_tiles_b; + int a_start = tile_a_index * TILE; + int b_start = tile_b_index * TILE; + int c_start = tile_c_index * TILE; + + int *bas = envs.bas; + double *env = envs.env; + int nbas = envs.nbas; + int bvk_nbas = envs.bvk_ncells * envs.nbas; + + if (thread_id == 0) { + start_position_x = c_dxyz_dabc[0] * a_start + c_dxyz_dabc[3] * b_start + c_dxyz_dabc[6] * c_start; + start_position_y = c_dxyz_dabc[1] * a_start + c_dxyz_dabc[4] * b_start + c_dxyz_dabc[7] * c_start; + start_position_z = c_dxyz_dabc[2] * a_start + c_dxyz_dabc[5] * b_start + c_dxyz_dabc[8] * c_start; + a_upper = min(a_start + TILE, mesh_a) - a_start; + b_upper = min(b_start + TILE, mesh_b) - b_start; + c_upper = min(c_start + TILE, mesh_c) - c_start; + } + + int lane = thread_id % WARP_SIZE; + int warp = thread_id / WARP_SIZE; + for (int n = thread_id; n < TILE*TILE*TILE*WARPS; n += threads) { + density_value[n] = 0; + } + __syncthreads(); + + for (int pair_id = shl_pair0+thread_id; pair_id < shl_pair1+thread_id; pair_id += threads) { + int ish = 0; + int jsh = 0; + if (pair_id < shl_pair1) { + int64_t bas_ij = dressed_bas_ij_idx[pair_id]; + ish = bas_ij / NBAS_MAX; + jsh = bas_ij % NBAS_MAX; + } + int latsum_idx = ish / nbas; + ish = ish - nbas * latsum_idx; + int jL = jsh / bvk_nbas; + jsh = jsh - bvk_nbas * jL; + int ish_cell0 = ish; + int bvk_cell_id = jsh / nbas; + int jsh_cell0 = jsh - nbas * bvk_cell_id; + double Lx = supmol_img_coords[latsum_idx*3+0]; + double Ly = supmol_img_coords[latsum_idx*3+1]; + double Lz = supmol_img_coords[latsum_idx*3+2]; + int expi = bas[ish_cell0*BAS_SLOTS+PTR_EXP]; + int expj = bas[jsh_cell0*BAS_SLOTS+PTR_EXP]; + int ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; + int rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; + double xi = env[ri+0] - Lx; + double yi = env[ri+1] - Ly; + double zi = env[ri+2] - Lz; + double xj = env[rj+0] - Lx + envs.img_coords[jL*3+0]; + double yj = env[rj+1] - Ly + envs.img_coords[jL*3+1]; + double zj = env[rj+2] - Lz + envs.img_coords[jL*3+2]; + double xjxi = xj - xi; + double yjyi = yj - yi; + double zjzi = zj - zi; + double rr_ij = distance_squared(xjxi, yjyi, zjzi); + double ai = env[expi]; + double aj = env[expj]; + double aij = ai + aj; + double aj_aij = aj / aij; + double theta_ij = ai * aj_aij; + double xij = xjxi * aj_aij + xi; + double yij = yjyi * aj_aij + yi; + double zij = zjzi * aj_aij + zi; + double x0 = start_position_x - xij; + double y0 = start_position_y - yij; + double z0 = start_position_z - zij; + double gaussian_exponent_at_reference = aij * distance_squared(x0, y0, z0); + double gaussian_starting_exponent = theta_ij * rr_ij + gaussian_exponent_at_reference; + double gaussian_starting_point = 0.; + // BUG: when aij gets too large and x0 is negative and large, the + // exponential can overflow and return inf. + // ideally recursion should start from the nearest grid point to the pair + // center, instead of the fixed recursion path + // (min a, min b, min c) -> (max a, max b, max c) + // The inf ususally occurs when pseudo-potential is not used, + // and core electrons appear with large exponents. + // Potentially another fix is to have a better designed multi-grid + // structure, where the gaussians with large exponents are evaluated + // on a more dense grid. Around the boundary the numbers should be + // within the range of double precision. + double recursion_factor_a_start = 0.; + double recursion_factor_b_start = 0.; + double recursion_factor_c_start = 0.; + double exp_da_squared = 0.; + double exp_db_squared = 0.; + double exp_dc_squared = 0.; + double exp_dadb = 0.; + double exp_dadc = 0.; + double exp_dbdc = 0.; + if (gaussian_starting_exponent < 680.) { + double ci = env[bas[ish*BAS_SLOTS+PTR_COEFF]]; + double cj = env[bas[jsh*BAS_SLOTS+PTR_COEFF]]; + double cc = ci * cj; + gaussian_starting_point = exp(-gaussian_starting_exponent); + gaussian_starting_point *= factor * cc; + if (ish == jsh_cell0) { + gaussian_starting_point *= 0.5; + } + double cross_term_a = c_dxyz_dabc[0] * x0 + c_dxyz_dabc[1] * y0 + c_dxyz_dabc[2] * z0; + double cross_term_b = c_dxyz_dabc[3] * x0 + c_dxyz_dabc[4] * y0 + c_dxyz_dabc[5] * z0; + double cross_term_c = c_dxyz_dabc[6] * x0 + c_dxyz_dabc[7] * y0 + c_dxyz_dabc[8] * z0; + recursion_factor_a_start = exp(-aij * (2 * cross_term_a + da_squared)); + recursion_factor_b_start = exp(-aij * (2 * cross_term_b + db_squared)); + recursion_factor_c_start = exp(-aij * (2 * cross_term_c + dc_squared)); + exp_da_squared = exp(-2 * aij * da_squared); + exp_db_squared = exp(-2 * aij * db_squared); + exp_dc_squared = exp(-2 * aij * dc_squared); + exp_dadb = exp(-2 * aij * a_dot_b); + exp_dadc = exp(-2 * aij * a_dot_c); + exp_dbdc = exp(-2 * aij * b_dot_c); + } +#pragma unroll + for (int dm_i0 = 0; dm_i0 < nfi; dm_i0 += SLICE_SIZE_I) { +#pragma unroll + for (int dm_j0 = 0; dm_j0 < nfj; dm_j0 += SLICE_SIZE_J) { + size_t nao = envs.ao_loc[nbas]; + int i0 = envs.ao_loc[ish_cell0]; + int j0 = envs.ao_loc[jsh_cell0]; + size_t ij_offset = bvk_cell_id * nao * nao + (dm_i0+i0) * nao + dm_j0+j0; + double dm_cache[SLICE_SIZE_I * SLICE_SIZE_J]; + if (pair_id < shl_pair1) { +#pragma unroll + for (int i = 0; i < SLICE_SIZE_I; ++i) { + if (SLICE_SIZE_I < nfi && dm_i0 + i >= nfi) break; +#pragma unroll + for (int j = 0; j < SLICE_SIZE_J; ++j) { + if (SLICE_SIZE_J < nfj && dm_j0 + j >= nfj) break; + dm_cache[i*SLICE_SIZE_J+j] = dm[ij_offset + i*nao+j]; + } } + } + + double x, y, z; + double recursion_factor_ab_pow_a = 1; + double recursion_factor_ac_pow_a = 1; + double recursion_factor_bc_pow_b = 1; + + if constexpr (is_non_orthogonal) { + // recursion_factor_ab_pow_a = 1; + // recursion_factor_ac_pow_a = 1; + } else { + x = start_position_x; + } + double gaussian_x = gaussian_starting_point; + double recursion_factor_a = recursion_factor_a_start; + for (int a_index = 0; a_index < a_upper; a_index++, + gaussian_x *= recursion_factor_a, + recursion_factor_a *= exp_da_squared) { + + if constexpr (is_non_orthogonal) { + recursion_factor_bc_pow_b = 1; + } else { + y = start_position_y; + } + double gaussian_xy = gaussian_x; + double recursion_factor_b = recursion_factor_b_start * recursion_factor_ab_pow_a; + for (int b_index = 0; b_index < b_upper; b_index++, + gaussian_xy *= recursion_factor_b, + recursion_factor_b *= exp_db_squared) { + + if constexpr (is_non_orthogonal) { + x = start_position_x + a_index * c_dxyz_dabc[0] + b_index * c_dxyz_dabc[3]; + y = start_position_y + a_index * c_dxyz_dabc[1] + b_index * c_dxyz_dabc[4]; + z = start_position_z + a_index * c_dxyz_dabc[2] + b_index * c_dxyz_dabc[5]; + } else { + z = start_position_z; + } + double gaussian_xyz = gaussian_xy; + double recursion_factor_c = recursion_factor_c_start * + recursion_factor_ac_pow_a * recursion_factor_bc_pow_b; + for (int c_index = 0; c_index < c_upper; c_index++, + gaussian_xyz *= recursion_factor_c, + recursion_factor_c *= exp_dc_squared) { + + double val = 0; + if (pair_id < shl_pair1 && fabs(gaussian_xyz) > negligible) { + double i_cartesian[nfi]; + gto_cartesian
  • (i_cartesian, x - xi, y - yi, z - zi); + rename_registers(i_cartesian, dm_i0, nfi, SLICE_SIZE_I); + + double j_cartesian[nfj]; + gto_cartesian(j_cartesian, x - xj, y - yj, z - zj); + rename_registers(j_cartesian, dm_j0, nfj, SLICE_SIZE_J); +#pragma unroll + for (int i = 0; i < SLICE_SIZE_I; ++i) { + if (SLICE_SIZE_I < nfi && dm_i0 + i >= nfi) break; + double s = 0; +#pragma unroll + for (int j = 0; j < SLICE_SIZE_J; ++j) { + if (SLICE_SIZE_J < nfj && dm_j0 + j >= nfj) break; + s += dm_cache[i * SLICE_SIZE_J + j] * j_cartesian[j]; + } + val += s * i_cartesian[i]; + } + val *= gaussian_xyz; + } + for (int offset = WARP_SIZE/2; offset > 0; offset >>= 1) { + val += __shfl_down_sync(0xffffffff, val, offset); + } + if (lane == 0) { + int abc_index = a_index * TILE*TILE + b_index*TILE + c_index; + density_value[abc_index+TILE*TILE*TILE*warp] += val; + } + if constexpr (is_non_orthogonal) { + x += c_dxyz_dabc[6]; + y += c_dxyz_dabc[7]; + z += c_dxyz_dabc[8]; + } else { + z += c_dxyz_dabc[8]; + } + } + if constexpr (is_non_orthogonal) { + recursion_factor_bc_pow_b *= exp_dbdc; + } else { + y += c_dxyz_dabc[4]; + } + } + if constexpr (is_non_orthogonal) { + recursion_factor_ab_pow_a *= exp_dadb; + recursion_factor_ac_pow_a *= exp_dadc; + } else { + x += c_dxyz_dabc[0]; + } + } + } } + } + __syncthreads(); + + int a_idx = a_start + thread_id / (TILE*TILE); + int b_idx = b_start + thread_id / TILE % TILE; + int c_idx = c_start + thread_id % TILE; + if (a_idx < mesh_a && b_idx < mesh_b && c_idx < mesh_c) { + double val = density_value[thread_id]; + for (int i = 1; i < WARPS; i++) { + val += density_value[thread_id+i*TILE*TILE*TILE]; + } + size_t abc_idx = (a_idx * mesh_b + b_idx) * (size_t)mesh_c + c_idx; + // update the density.real, skip the imaginary part + atomicAdd(density + abc_idx*2, val); + } + __syncthreads(); +} +} + +extern "C" { +#define eval_density_kernel_case(li, lj, slice_i, slice_j, non_orth) \ + case (li * LMAX1 + lj): \ + eval_density_kernel<<>>( \ + density, dm, *envs, supmol_img_coords, factor, \ + shl_pair_offsets, dressed_bas_ij_idx, \ + grid_tile_index, n_contributing_tiles, tiles_per_block, \ + a_dot_b, a_dot_c, b_dot_c, da_squared, db_squared, dc_squared, \ + mesh_a, mesh_b, mesh_c, negligible); \ + break + +int evaluate_density(double *density, double *placeholder, + double *dm, PBCIntEnvVars *envs, + double *dxyz_dabc, double *supmol_img_coords, + int i_angular, int j_angular, int tiles_per_block, + int *shl_pair_offsets, int64_t *dressed_bas_ij_idx, + int *grid_tile_index, int n_contributing_tiles, int *mesh, + double factor, double negligible) +{ + int mesh_a = mesh[0]; + int mesh_b = mesh[1]; + int mesh_c = mesh[2]; + int block_grid = (n_contributing_tiles + tiles_per_block-1) / tiles_per_block; + double a_dot_b = dxyz_dabc[0] * dxyz_dabc[3] + dxyz_dabc[1] * dxyz_dabc[4] + dxyz_dabc[2] * dxyz_dabc[5]; + double a_dot_c = dxyz_dabc[0] * dxyz_dabc[6] + dxyz_dabc[1] * dxyz_dabc[7] + dxyz_dabc[2] * dxyz_dabc[8]; + double b_dot_c = dxyz_dabc[3] * dxyz_dabc[6] + dxyz_dabc[4] * dxyz_dabc[7] + dxyz_dabc[5] * dxyz_dabc[8]; + double da_squared = distance_squared(dxyz_dabc[0], dxyz_dabc[1], dxyz_dabc[2]); + double db_squared = distance_squared(dxyz_dabc[3], dxyz_dabc[4], dxyz_dabc[5]); + double dc_squared = distance_squared(dxyz_dabc[6], dxyz_dabc[7], dxyz_dabc[8]); + if (fabs(dxyz_dabc[1]) < 1e-14 && fabs(dxyz_dabc[2]) < 1e-14 && + fabs(dxyz_dabc[3]) < 1e-14 && fabs(dxyz_dabc[5]) < 1e-14 && + fabs(dxyz_dabc[6]) < 1e-14 && fabs(dxyz_dabc[8]) < 1e-14) { + // For most high-angular-momemtum basis, the performance of orthogonal + // and general cases are nearly identical. It's sufficient to specialize + // the orthogonal lattice for a few low-angular cases. + switch (i_angular * LMAX1 + j_angular) { + eval_density_kernel_case(0,0, 1, 1, 0); + eval_density_kernel_case(1,0, 3, 1, 0); + eval_density_kernel_case(1,1, 3, 3, 0); + eval_density_kernel_case(2,0, 6, 1, 0); + eval_density_kernel_case(2,1, 6, 3, 0); + eval_density_kernel_case(2,2, 6, 6, 1); + eval_density_kernel_case(3,0,10, 1, 0); + eval_density_kernel_case(3,1,10, 3, 1); + eval_density_kernel_case(3,2,10, 6, 1); + eval_density_kernel_case(3,3,10, 5, 1); + eval_density_kernel_case(4,0,15, 1, 0); + eval_density_kernel_case(4,1,15, 3, 1); + eval_density_kernel_case(4,2, 8, 6, 1); + eval_density_kernel_case(4,3,15, 5, 1); + eval_density_kernel_case(4,4,15, 5, 1); + } + } else { + switch (i_angular * LMAX1 + j_angular) { + eval_density_kernel_case(0,0, 1, 1, 1); + eval_density_kernel_case(1,0, 3, 1, 1); + eval_density_kernel_case(1,1, 3, 3, 1); + eval_density_kernel_case(2,0, 6, 1, 1); + eval_density_kernel_case(2,1, 6, 3, 1); + eval_density_kernel_case(2,2, 6, 6, 1); + eval_density_kernel_case(3,0,10, 1, 1); + eval_density_kernel_case(3,1,10, 3, 1); + eval_density_kernel_case(3,2,10, 6, 1); + eval_density_kernel_case(3,3,10, 5, 1); + eval_density_kernel_case(4,0,15, 1, 1); + eval_density_kernel_case(4,1,15, 3, 1); + eval_density_kernel_case(4,2, 8, 6, 1); + eval_density_kernel_case(4,3,15, 5, 1); + eval_density_kernel_case(4,4,15, 5, 1); + } + } + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in eval_density_kernel: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} +} diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/eval_lda_mat.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_lda_mat.cu new file mode 100644 index 000000000..3701e2803 --- /dev/null +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_lda_mat.cu @@ -0,0 +1,313 @@ +/* + * Copyright 2025-2026 The PySCF Developers. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include "gvhf-rys/vhf.cuh" +#include "constant_objects.cuh" +#include "cartesian.cuh" +#include "utils.cuh" + +#define TILE 4 +#define WARP_SIZE 32 +#define THREADS 64 + +template +__global__ static +void eval_lda_mat_kernel(double *out, double *vxc_weights, PBCIntEnvVars envs, + double *supmol_img_coords, + int *shl_pair_offsets, int64_t *dressed_bas_ij_idx, + int *grid_tile_index, int ntiles, int tiles_per_block, + double a_dot_b, double a_dot_c, double b_dot_c, + double da_squared, double db_squared, double dc_squared, + int mesh_a, int mesh_b, int mesh_c, double negligible) +{ + constexpr int threads = THREADS; + int thread_id = threadIdx.x; + int tile_id0 = blockIdx.x * tiles_per_block; + __shared__ int a_upper, b_upper, c_upper; + __shared__ double start_position_x, start_position_y, start_position_z; + __shared__ double vxc_cache[TILE*TILE*TILE]; + +for (int tile_id = tile_id0; tile_id < min(tile_id0+tiles_per_block, ntiles); tile_id++) { + int tile_index = grid_tile_index[tile_id]; + int shl_pair0 = shl_pair_offsets[tile_id]; + int shl_pair1 = shl_pair_offsets[tile_id+1]; + int n_tiles_b = (mesh_b + TILE - 1) / TILE; + int n_tiles_c = (mesh_c + TILE - 1) / TILE; + int tile_ab_index = tile_index / n_tiles_c; + int tile_c_index = tile_index % n_tiles_c; + int tile_a_index = tile_ab_index / n_tiles_b; + int tile_b_index = tile_ab_index % n_tiles_b; + int a_start = tile_a_index * TILE; + int b_start = tile_b_index * TILE; + int c_start = tile_c_index * TILE; + + constexpr int nfi = (LI + 1) * (LI + 2) / 2; + constexpr int nfj = (LJ + 1) * (LJ + 2) / 2; + + int *bas = envs.bas; + double *env = envs.env; + int nbas = envs.nbas; + int bvk_nbas = envs.bvk_ncells * envs.nbas; + + if (thread_id == 0) { + start_position_x = c_dxyz_dabc[0] * a_start + c_dxyz_dabc[3] * b_start + c_dxyz_dabc[6] * c_start; + start_position_y = c_dxyz_dabc[1] * a_start + c_dxyz_dabc[4] * b_start + c_dxyz_dabc[7] * c_start; + start_position_z = c_dxyz_dabc[2] * a_start + c_dxyz_dabc[5] * b_start + c_dxyz_dabc[8] * c_start; + a_upper = min(a_start + TILE, mesh_a) - a_start; + b_upper = min(b_start + TILE, mesh_b) - b_start; + c_upper = min(c_start + TILE, mesh_c) - c_start; + } + + int a_idx = a_start + thread_id / (TILE*TILE); + int b_idx = b_start + thread_id / TILE % TILE; + int c_idx = c_start + thread_id % TILE; + if (a_idx < mesh_a && b_idx < mesh_b && c_idx < mesh_c) { + size_t abc_idx = (a_idx * mesh_b + b_idx) * (size_t)mesh_c + c_idx; + vxc_cache[thread_id] = vxc_weights[abc_idx]; + } + __syncthreads(); + + for (int pair_id = shl_pair0+thread_id; pair_id < shl_pair1; pair_id += threads) { + int64_t bas_ij = dressed_bas_ij_idx[pair_id]; + int ish = bas_ij / NBAS_MAX; + int jsh = bas_ij % NBAS_MAX; + int latsum_idx = ish / nbas; + ish = ish - nbas * latsum_idx; + int jL = jsh / bvk_nbas; + jsh = jsh - bvk_nbas * jL; + double Lx = supmol_img_coords[latsum_idx*3+0]; + double Ly = supmol_img_coords[latsum_idx*3+1]; + double Lz = supmol_img_coords[latsum_idx*3+2]; + int expi = bas[ish*BAS_SLOTS+PTR_EXP]; + int expj = bas[jsh*BAS_SLOTS+PTR_EXP]; + int ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; + int rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; + double xi = env[ri+0] - Lx; + double yi = env[ri+1] - Ly; + double zi = env[ri+2] - Lz; + double xj = env[rj+0] - Lx + envs.img_coords[jL*3+0]; + double yj = env[rj+1] - Ly + envs.img_coords[jL*3+1]; + double zj = env[rj+2] - Lz + envs.img_coords[jL*3+2]; + double xjxi = xj - xi; + double yjyi = yj - yi; + double zjzi = zj - zi; + double rr_ij = distance_squared(xjxi, yjyi, zjzi); + double ai = env[expi]; + double aj = env[expj]; + double aij = ai + aj; + double aj_aij = aj / aij; + double theta_ij = ai * aj_aij; + double xij = xjxi * aj_aij + xi; + double yij = yjyi * aj_aij + yi; + double zij = zjzi * aj_aij + zi; + double x0 = start_position_x - xij; + double y0 = start_position_y - yij; + double z0 = start_position_z - zij; + double gaussian_exponent_at_reference = aij * distance_squared(x0, y0, z0); + double gaussian_starting_exponent = theta_ij * rr_ij + gaussian_exponent_at_reference; + if (gaussian_starting_exponent > 680.) { + continue; + } + int ish_cell0 = ish; + int bvk_cell_id = jsh / nbas; + int jsh_cell0 = jsh - nbas * bvk_cell_id; + double ci = env[bas[ish_cell0*BAS_SLOTS+PTR_COEFF]]; + double cj = env[bas[jsh_cell0*BAS_SLOTS+PTR_COEFF]]; + double cc = ci * cj; + if (ish_cell0 == jsh_cell0) { + cc *= .5; + } + double gaussian_starting_point = exp(-gaussian_starting_exponent) * cc; + double cross_term_a = c_dxyz_dabc[0] * x0 + c_dxyz_dabc[1] * y0 + c_dxyz_dabc[2] * z0; + double cross_term_b = c_dxyz_dabc[3] * x0 + c_dxyz_dabc[4] * y0 + c_dxyz_dabc[5] * z0; + double cross_term_c = c_dxyz_dabc[6] * x0 + c_dxyz_dabc[7] * y0 + c_dxyz_dabc[8] * z0; + double recursion_factor_a_start = exp(-aij * (2 * cross_term_a + da_squared)); + double recursion_factor_b_start = exp(-aij * (2 * cross_term_b + db_squared)); + double recursion_factor_c_start = exp(-aij * (2 * cross_term_c + dc_squared)); + double exp_da_squared = exp(-2 * aij * da_squared); + double exp_db_squared = exp(-2 * aij * db_squared); + double exp_dc_squared = exp(-2 * aij * dc_squared); + double exp_dadb = exp(-2 * aij * a_dot_b); + double exp_dadc = exp(-2 * aij * a_dot_c); + double exp_dbdc = exp(-2 * aij * b_dot_c); +#pragma unroll + for (int dm_i0 = 0; dm_i0 < nfi; dm_i0 += SLICE_SIZE_I) { +#pragma unroll + for (int dm_j0 = 0; dm_j0 < nfj; dm_j0 += SLICE_SIZE_J) { + double vj_cache[SLICE_SIZE_I * SLICE_SIZE_J]; +#pragma unroll + for (int n = 0; n < SLICE_SIZE_I*SLICE_SIZE_J; ++n) { + vj_cache[n] = 0; + } + + double x, y, z; + double recursion_factor_ab_pow_a = 1; + double recursion_factor_ac_pow_a = 1; + double recursion_factor_bc_pow_b = 1; + + if constexpr (is_non_orthogonal) { + } else { + x = start_position_x; + } + double gaussian_x = gaussian_starting_point; + double recursion_factor_a = recursion_factor_a_start; + for (int a_index = 0; a_index < a_upper; a_index++, + gaussian_x *= recursion_factor_a, + recursion_factor_a *= exp_da_squared) { + + if constexpr (is_non_orthogonal) { + recursion_factor_bc_pow_b = 1; + } else { + y = start_position_y; + } + double gaussian_xy = gaussian_x; + double recursion_factor_b = recursion_factor_b_start * recursion_factor_ab_pow_a; + for (int b_index = 0; b_index < b_upper; b_index++, + gaussian_xy *= recursion_factor_b, + recursion_factor_b *= exp_db_squared) { + + if constexpr (is_non_orthogonal) { + x = start_position_x + a_index * c_dxyz_dabc[0] + b_index * c_dxyz_dabc[3]; + y = start_position_y + a_index * c_dxyz_dabc[1] + b_index * c_dxyz_dabc[4]; + z = start_position_z + a_index * c_dxyz_dabc[2] + b_index * c_dxyz_dabc[5]; + } else { + z = start_position_z; + } + double gaussian_xyz = gaussian_xy; + double recursion_factor_c = recursion_factor_c_start * + recursion_factor_ac_pow_a * recursion_factor_bc_pow_b; + for (int c_index = 0; c_index < c_upper; c_index++, + gaussian_xyz *= recursion_factor_c, + recursion_factor_c *= exp_dc_squared) { + + if (fabs(gaussian_xyz) > negligible) { + double i_cartesian[nfi]; + gto_cartesian
  • (i_cartesian, x - xi, y - yi, z - zi); + rename_registers(i_cartesian, dm_i0, nfi, SLICE_SIZE_I); + + double j_cartesian[nfj]; + gto_cartesian(j_cartesian, x - xj, y - yj, z - zj); + rename_registers(j_cartesian, dm_j0, nfj, SLICE_SIZE_J); + + int abc_index = a_index * TILE*TILE + b_index*TILE + c_index; + double fac = gaussian_xyz * vxc_cache[abc_index]; +#pragma unroll + for (int i = 0; i < SLICE_SIZE_I; ++i) { + if (SLICE_SIZE_I < nfi && dm_i0 + i >= nfi) break; + double s = fac * i_cartesian[i]; +#pragma unroll + for (int j = 0; j < SLICE_SIZE_J; ++j) { + if (SLICE_SIZE_J < nfj && dm_j0 + j >= nfj) break; + vj_cache[i*SLICE_SIZE_J+j] += s * j_cartesian[j]; + } + } + } + if constexpr (is_non_orthogonal) { + x += c_dxyz_dabc[6]; + y += c_dxyz_dabc[7]; + z += c_dxyz_dabc[8]; + } else { + z += c_dxyz_dabc[8]; + } + } + if constexpr (is_non_orthogonal) { + recursion_factor_bc_pow_b *= exp_dbdc; + } else { + y += c_dxyz_dabc[4]; + } + } + if constexpr (is_non_orthogonal) { + recursion_factor_ab_pow_a *= exp_dadb; + recursion_factor_ac_pow_a *= exp_dadc; + } else { + x += c_dxyz_dabc[0]; + } + } + + size_t nao = envs.ao_loc[nbas]; + int i0 = envs.ao_loc[ish_cell0]; + int j0 = envs.ao_loc[jsh_cell0]; + double *pout = out + bvk_cell_id * nao * nao + (dm_i0+i0) * nao + dm_j0+j0; +#pragma unroll + for (int i = 0; i < SLICE_SIZE_I; ++i) { + if (SLICE_SIZE_I < nfi && dm_i0 + i >= nfi) break; +#pragma unroll + for (int j = 0; j < SLICE_SIZE_J; ++j) { + if (SLICE_SIZE_J < nfj && dm_j0 + j >= nfj) break; + atomicAdd(pout + i*nao+j, vj_cache[i*SLICE_SIZE_J+j]); + } } + } } + } +} +} + +extern "C" { +#define eval_lda_mat_kernel_case(li, lj, slice_i, slice_j, non_orth) \ + case (li * LMAX1 + lj): \ + eval_lda_mat_kernel<<>>( \ + out, vxc, *envs, supmol_img_coords, shl_pair_offsets, dressed_bas_ij_idx, \ + grid_tile_index, n_contributing_tiles, tiles_per_block, \ + a_dot_b, a_dot_c, b_dot_c, da_squared, db_squared, dc_squared, \ + mesh_a, mesh_b, mesh_c, negligible); \ + break + +int evaluate_lda_mat(double *out, double *vxc, double *placeholder, PBCIntEnvVars *envs, + double *dxyz_dabc, double *supmol_img_coords, + int i_angular, int j_angular, int tiles_per_block, + int *shl_pair_offsets, int64_t *dressed_bas_ij_idx, + int *grid_tile_index, int n_contributing_tiles, int *mesh, + double negligible) +{ + int mesh_a = mesh[0]; + int mesh_b = mesh[1]; + int mesh_c = mesh[2]; + int block_grid = (n_contributing_tiles + tiles_per_block-1) / tiles_per_block; + double a_dot_b = dxyz_dabc[0] * dxyz_dabc[3] + dxyz_dabc[1] * dxyz_dabc[4] + dxyz_dabc[2] * dxyz_dabc[5]; + double a_dot_c = dxyz_dabc[0] * dxyz_dabc[6] + dxyz_dabc[1] * dxyz_dabc[7] + dxyz_dabc[2] * dxyz_dabc[8]; + double b_dot_c = dxyz_dabc[3] * dxyz_dabc[6] + dxyz_dabc[4] * dxyz_dabc[7] + dxyz_dabc[5] * dxyz_dabc[8]; + double da_squared = distance_squared(dxyz_dabc[0], dxyz_dabc[1], dxyz_dabc[2]); + double db_squared = distance_squared(dxyz_dabc[3], dxyz_dabc[4], dxyz_dabc[5]); + double dc_squared = distance_squared(dxyz_dabc[6], dxyz_dabc[7], dxyz_dabc[8]); + switch (i_angular * LMAX1 + j_angular) { + eval_lda_mat_kernel_case(0,0, 1, 1, 1); + eval_lda_mat_kernel_case(1,0, 3, 1, 1); + eval_lda_mat_kernel_case(1,1, 3, 3, 1); + eval_lda_mat_kernel_case(2,0, 6, 1, 1); + eval_lda_mat_kernel_case(2,1, 6, 3, 1); + eval_lda_mat_kernel_case(2,2, 6, 6, 1); + eval_lda_mat_kernel_case(3,0,10, 1, 1); + eval_lda_mat_kernel_case(3,1,10, 3, 1); + eval_lda_mat_kernel_case(3,2,10, 6, 1); + eval_lda_mat_kernel_case(3,3, 5,10, 1); + eval_lda_mat_kernel_case(4,0,15, 1, 1); + eval_lda_mat_kernel_case(4,1,15, 3, 1); + eval_lda_mat_kernel_case(4,2, 8, 6, 1); + eval_lda_mat_kernel_case(4,3,15, 5, 1); + eval_lda_mat_kernel_case(4,4,15, 5, 1); + } + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in eval_lda_mat_kernel: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} +} diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/eval_lda_mat_v2.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_lda_mat_v2.cu new file mode 100644 index 000000000..bd4f917bf --- /dev/null +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_lda_mat_v2.cu @@ -0,0 +1,291 @@ +/* + * Copyright 2025-2026 The PySCF Developers. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include "gvhf-rys/vhf.cuh" +#include "constant_objects.cuh" +#include "cartesian.cuh" +#include "utils.cuh" + +template +__global__ static +void eval_lda_mat_kernel_v2(double *out, double *vxc_weights, PBCIntEnvVars envs, + int64_t *bas_ij_idx, float2 *grid_frac_ranges, + double da_squared, double db_squared, double dc_squared, + int mesh_a, int mesh_b, int mesh_c, int npairs, + double negligible) +{ + constexpr int tile = 16; + int tx = threadIdx.x; + int ty = threadIdx.y; + int thread_id = ty * tile + tx; + int pair_id = blockIdx.x; + + __shared__ int a_start, a_stop, a_center; + __shared__ int b_start, b_stop; + __shared__ int c_start, c_stop; + __shared__ uint32_t ij_offset; + __shared__ double cc, exp_da_squared; + __shared__ double xi, yi, zi; + __shared__ double xj, yj, zj; + __shared__ double xij, yij, zij, aij, theta_rr; + __shared__ double swap[8]; + + int *bas = envs.bas; + double *env = envs.env; + int nbas = envs.nbas; + int bvk_nbas = envs.bvk_ncells * envs.nbas; + + int64_t bas_ij = bas_ij_idx[pair_id]; + int ish = bas_ij / NBAS_MAX; + int jsh = bas_ij % NBAS_MAX; + int jL = jsh / bvk_nbas; + jsh = jsh - bvk_nbas * jL; + if (thread_id == 0) { + int ish_cell0 = ish; + int bvk_cell_id = jsh / nbas; + int jsh_cell0 = jsh - nbas * bvk_cell_id; + uint32_t nao = envs.ao_loc[nbas]; + uint32_t i0 = envs.ao_loc[ish_cell0]; + uint32_t j0 = envs.ao_loc[jsh_cell0]; + ij_offset = bvk_cell_id * nao * nao + i0 * nao + j0; + + int ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; + int rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; + double ai = env[bas[ish_cell0*BAS_SLOTS+PTR_EXP]]; + double aj = env[bas[jsh_cell0*BAS_SLOTS+PTR_EXP]]; + double ci = env[bas[ish_cell0*BAS_SLOTS+PTR_COEFF]]; + double cj = env[bas[jsh_cell0*BAS_SLOTS+PTR_COEFF]]; + cc = ci * cj; + if (ish_cell0 == jsh_cell0) { + cc *= .5; + } + aij = ai + aj; + double aj_aij = aj / aij; + xi = env[ri+0]; + yi = env[ri+1]; + zi = env[ri+2]; + xj = env[rj+0] + envs.img_coords[jL*3+0]; + yj = env[rj+1] + envs.img_coords[jL*3+1]; + zj = env[rj+2] + envs.img_coords[jL*3+2]; + double xjxi = xj - xi; + double yjyi = yj - yi; + double zjzi = zj - zi; + double rr = distance_squared(xjxi, yjyi, zjzi); + xij = xjxi * aj_aij + xi; + yij = yjyi * aj_aij + yi; + zij = zjzi * aj_aij + zi; + theta_rr = ai * aj_aij * rr; + + float2 range = grid_frac_ranges[pair_id]; + float xfrac_lower = range.x; + float xfrac_upper = range.y; + range = grid_frac_ranges[npairs+pair_id]; + float yfrac_lower = range.x; + float yfrac_upper = range.y; + range = grid_frac_ranges[npairs*2+pair_id]; + float zfrac_lower = range.x; + float zfrac_upper = range.y; + a_start = ceil(xfrac_lower * mesh_a); + b_start = ceil(yfrac_lower * mesh_b); + c_start = ceil(zfrac_lower * mesh_c); + a_stop = floor(xfrac_upper * mesh_a); + b_stop = floor(yfrac_upper * mesh_b); + c_stop = floor(zfrac_upper * mesh_c); + a_center = (a_start + a_stop) / 2; + exp_da_squared = exp(-2 * aij * da_squared); + } + __syncthreads(); + + constexpr int nfi = (LI + 1) * (LI + 2) / 2; + constexpr int nfj = (LJ + 1) * (LJ + 2) / 2; + +#pragma unroll + for (int dm_i0 = 0; dm_i0 < nfi; dm_i0 += SLICE_SIZE_I) { +#pragma unroll + for (int dm_j0 = 0; dm_j0 < nfj; dm_j0 += SLICE_SIZE_J) { + double vj_cache[SLICE_SIZE_I * SLICE_SIZE_J]; +#pragma unroll + for (int n = 0; n < SLICE_SIZE_I*SLICE_SIZE_J; ++n) { + vj_cache[n] = 0; + } + + for (int b_index = b_start+ty; b_index <= b_stop; b_index += tile) { + for (int c_index = c_start+tx; c_index <= c_stop; c_index += tile) { + double x_start = a_center * c_dxyz_dabc[0] + b_index * c_dxyz_dabc[3] + c_index * c_dxyz_dabc[6]; + double y_start = a_center * c_dxyz_dabc[1] + b_index * c_dxyz_dabc[4] + c_index * c_dxyz_dabc[7]; + double z_start = a_center * c_dxyz_dabc[2] + b_index * c_dxyz_dabc[5] + c_index * c_dxyz_dabc[8]; + double x = x_start; + double y = y_start; + double z = z_start; + double x_xij = x - xij; + double y_yij = y - yij; + double z_zij = z - zij; + double e = theta_rr + aij * distance_squared(x_xij, y_yij, z_zij); + if (e > 50.) continue; // ~1e-22 + double gaussian_starting_point = exp(-e) * cc; + double cross_term_a = c_dxyz_dabc[0] * x_xij + c_dxyz_dabc[1] * y_yij + c_dxyz_dabc[2] * z_zij; + double recursion_factor_a_start = exp(-aij * (2 * cross_term_a + da_squared)); + + int mesh_bc = mesh_b * mesh_c; + int mesh_abc = mesh_a * mesh_bc; + // mod(negative_number, N) leads to a negative value. Adding a large + // multiplier to a_index, b_index, c_index to avoid the negative modulo. + // Images spreads in supmol are typically < 10. A shift of 100* images + // should be enough. + int64_t abc_idx_start = (a_center + 100 * mesh_a) % mesh_a * (int64_t)mesh_bc + + (b_index + 100 * mesh_b) % mesh_b * mesh_c + + (c_index + 100 * mesh_c) % mesh_c; + int64_t abc_idx = abc_idx_start; + double gaussian_xyz = gaussian_starting_point; + double recursion_factor_a = recursion_factor_a_start; + for (int a_index = a_center; a_index <= a_stop; a_index++, + gaussian_xyz *= recursion_factor_a, + recursion_factor_a *= exp_da_squared) { + if (fabs(gaussian_xyz) < negligible) break; + + double v = vxc_weights[abc_idx] * gaussian_xyz; + double i_cartesian[nfi]; + gto_cartesian
  • (i_cartesian, x - xi, y - yi, z - zi); + rename_registers(i_cartesian, dm_i0, nfi, SLICE_SIZE_I); + + double j_cartesian[nfj]; + gto_cartesian(j_cartesian, x - xj, y - yj, z - zj); + rename_registers(j_cartesian, dm_j0, nfj, SLICE_SIZE_J); +#pragma unroll + for (int i = 0; i < SLICE_SIZE_I; ++i) { + if (SLICE_SIZE_I < nfi && dm_i0 + i >= nfi) break; + double s = v * i_cartesian[i]; +#pragma unroll + for (int j = 0; j < SLICE_SIZE_J; ++j) { + if (SLICE_SIZE_J < nfj && dm_j0 + j >= nfj) break; + vj_cache[i*SLICE_SIZE_J+j] += s * j_cartesian[j]; + } + } + x += c_dxyz_dabc[0]; + y += c_dxyz_dabc[1]; + z += c_dxyz_dabc[2]; + abc_idx += mesh_bc; + if (abc_idx >= mesh_abc) { + abc_idx -= mesh_abc; + } + } + + x = x_start; + y = y_start; + z = z_start; + gaussian_xyz = gaussian_starting_point; + double inv_recursion_factor_a = exp_da_squared / recursion_factor_a_start; + abc_idx = abc_idx_start; + for (int a_index = a_center - 1; a_index >= a_start; a_index--, + inv_recursion_factor_a *= exp_da_squared) { + gaussian_xyz *= inv_recursion_factor_a; + if (fabs(gaussian_xyz) < negligible) break; + x -= c_dxyz_dabc[0]; + y -= c_dxyz_dabc[1]; + z -= c_dxyz_dabc[2]; + abc_idx -= mesh_bc; + if (abc_idx < 0) { + abc_idx += mesh_abc; + } + double v = vxc_weights[abc_idx] * gaussian_xyz; + double i_cartesian[nfi]; + gto_cartesian
  • (i_cartesian, x - xi, y - yi, z - zi); + rename_registers(i_cartesian, dm_i0, nfi, SLICE_SIZE_I); + + double j_cartesian[nfj]; + gto_cartesian(j_cartesian, x - xj, y - yj, z - zj); + rename_registers(j_cartesian, dm_j0, nfj, SLICE_SIZE_J); +#pragma unroll + for (int i = 0; i < SLICE_SIZE_I; ++i) { + if (SLICE_SIZE_I < nfi && dm_i0 + i >= nfi) break; + double s = v * i_cartesian[i]; +#pragma unroll + for (int j = 0; j < SLICE_SIZE_J; ++j) { + if (SLICE_SIZE_J < nfj && dm_j0 + j >= nfj) break; + vj_cache[i*SLICE_SIZE_J+j] += s * j_cartesian[j]; + } + } + } + } } + + uint32_t nao = envs.ao_loc[nbas]; + double *pout = out + ij_offset + dm_i0 * nao + dm_j0; +#pragma unroll + for (int i = 0; i < SLICE_SIZE_I; ++i) { + if (SLICE_SIZE_I < nfi && dm_i0 + i >= nfi) break; +#pragma unroll + for (int j = 0; j < SLICE_SIZE_J; ++j) { + if (SLICE_SIZE_J < nfj && dm_j0 + j >= nfj) break; + double val = reduce(vj_cache[i*SLICE_SIZE_J+j], swap, thread_id); + if (thread_id == 0) { + atomicAdd(pout + i*nao+j, val); + } + } + } + } } +} + +extern "C" { +#define eval_lda_mat_kernel_v2_case(li, lj, slice_i, slice_j) \ + case (li * LMAX1 + lj): \ + eval_lda_mat_kernel_v2<<>>( \ + out, vxc, *envs, bas_ij_idx, grid_frac_ranges, \ + da_squared, db_squared, dc_squared, mesh_a, mesh_b, mesh_c, npairs, \ + negligible); \ + break + +int evaluate_lda_mat_v2(double *out, double *vxc, double *placeholder, PBCIntEnvVars *envs, + double *dxyz_dabc, int li, int lj, int64_t *bas_ij_idx, + float2 *grid_frac_ranges, int *mesh, int npairs, + double negligible) +{ + int mesh_a = mesh[0]; + int mesh_b = mesh[1]; + int mesh_c = mesh[2]; + double da_squared = distance_squared(dxyz_dabc[0], dxyz_dabc[1], dxyz_dabc[2]); + double db_squared = distance_squared(dxyz_dabc[3], dxyz_dabc[4], dxyz_dabc[5]); + double dc_squared = distance_squared(dxyz_dabc[6], dxyz_dabc[7], dxyz_dabc[8]); + dim3 threads(16, 16); + switch (li * LMAX1 + lj) { + eval_lda_mat_kernel_v2_case(0,0, 1, 1); + eval_lda_mat_kernel_v2_case(1,0, 3, 1); + eval_lda_mat_kernel_v2_case(1,1, 3, 3); + eval_lda_mat_kernel_v2_case(2,0, 6, 1); + eval_lda_mat_kernel_v2_case(2,1, 6, 3); + eval_lda_mat_kernel_v2_case(2,2, 6, 6); + eval_lda_mat_kernel_v2_case(3,0,10, 1); + eval_lda_mat_kernel_v2_case(3,1,10, 3); + eval_lda_mat_kernel_v2_case(3,2,10, 6); + eval_lda_mat_kernel_v2_case(3,3,10, 5); + eval_lda_mat_kernel_v2_case(4,0,15, 1); + eval_lda_mat_kernel_v2_case(4,1,15, 3); + eval_lda_mat_kernel_v2_case(4,2, 8, 6); + eval_lda_mat_kernel_v2_case(4,3,15, 5); + eval_lda_mat_kernel_v2_case(4,4,15, 5); + } + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in eval_lda_mat_kernel: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} +} diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/eval_lda_strain_grad.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_lda_strain_grad.cu new file mode 100644 index 000000000..1a53bc3f2 --- /dev/null +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_lda_strain_grad.cu @@ -0,0 +1,399 @@ +/* + * Copyright 2025-2026 The PySCF Developers. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include "gvhf-rys/vhf.cuh" +#include "constant_objects.cuh" +#include "cartesian.cuh" +#include "utils.cuh" + +#define TILE 4 +#define WARP_SIZE 32 +#define THREADS 64 + +template +__global__ static +void eval_lda_grad_kernel(double *grad, double *strain, double *dm, + double *vxc_weights, PBCIntEnvVars envs, + int64_t *bas_ij_idx, float2 *grid_frac_ranges, + double da_squared, double db_squared, double dc_squared, + int mesh_a, int mesh_b, int mesh_c, int npairs, + double factor, double negligible) +{ + constexpr int tile = 16; + int tx = threadIdx.x; + int ty = threadIdx.y; + int thread_id = ty * tile + tx; + int pair_id = blockIdx.x; + + constexpr int nfi = (LI + 1) * (LI + 2) / 2; + constexpr int nfj = (LJ + 1) * (LJ + 2) / 2; + + __shared__ int a_start, a_stop, a_center; + __shared__ int b_start, b_stop; + __shared__ int c_start, c_stop; + __shared__ double cc, exp_da_squared; + __shared__ double xi, yi, zi; + __shared__ double xj, yj, zj; + __shared__ double xij, yij, zij, ai, aj, aij, theta_rr; + __shared__ double xjxi, yjyi, zjzi; + __shared__ double dm_cache[nfi*nfj]; + + int *bas = envs.bas; + double *env = envs.env; + int nbas = envs.nbas; + int bvk_nbas = envs.bvk_ncells * envs.nbas; + + int64_t bas_ij = bas_ij_idx[pair_id]; + int ish = bas_ij / NBAS_MAX; + int jsh = bas_ij % NBAS_MAX; + int jL = jsh / bvk_nbas; + jsh = jsh - bvk_nbas * jL; + if (thread_id == 0) { + int ish_cell0 = ish; + int bvk_cell_id = jsh / nbas; + int jsh_cell0 = jsh - nbas * bvk_cell_id; + int ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; + int rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; + ai = env[bas[ish_cell0*BAS_SLOTS+PTR_EXP]]; + aj = env[bas[jsh_cell0*BAS_SLOTS+PTR_EXP]]; + double ci = env[bas[ish_cell0*BAS_SLOTS+PTR_COEFF]]; + double cj = env[bas[jsh_cell0*BAS_SLOTS+PTR_COEFF]]; + cc = ci * cj; + if (ish_cell0 == jsh_cell0) { + cc *= .5; + } + aij = ai + aj; + double aj_aij = aj / aij; + xi = env[ri+0]; + yi = env[ri+1]; + zi = env[ri+2]; + xj = env[rj+0] + envs.img_coords[jL*3+0]; + yj = env[rj+1] + envs.img_coords[jL*3+1]; + zj = env[rj+2] + envs.img_coords[jL*3+2]; + xjxi = xj - xi; + yjyi = yj - yi; + zjzi = zj - zi; + double rr = distance_squared(xjxi, yjyi, zjzi); + xij = xjxi * aj_aij + xi; + yij = yjyi * aj_aij + yi; + zij = zjzi * aj_aij + zi; + theta_rr = ai * aj_aij * rr; + + float2 range = grid_frac_ranges[pair_id]; + float xfrac_lower = range.x; + float xfrac_upper = range.y; + range = grid_frac_ranges[npairs+pair_id]; + float yfrac_lower = range.x; + float yfrac_upper = range.y; + range = grid_frac_ranges[npairs*2+pair_id]; + float zfrac_lower = range.x; + float zfrac_upper = range.y; + a_start = ceil(xfrac_lower * mesh_a); + b_start = ceil(yfrac_lower * mesh_b); + c_start = ceil(zfrac_lower * mesh_c); + a_stop = floor(xfrac_upper * mesh_a); + b_stop = floor(yfrac_upper * mesh_b); + c_stop = floor(zfrac_upper * mesh_c); + a_center = (a_start + a_stop) / 2; + exp_da_squared = exp(-2 * aij * da_squared); + } + __syncthreads(); + + for (int n = thread_id; n < nfi * nfj; n += THREADS) { + int ish_cell0 = ish; + int bvk_cell_id = jsh / nbas; + int jsh_cell0 = jsh - nbas * bvk_cell_id; + size_t nao = envs.ao_loc[nbas]; + int i0 = envs.ao_loc[ish_cell0]; + int j0 = envs.ao_loc[jsh_cell0]; + int i = n * c_div_nf[LJ]; + int j = n - nfj * i; + dm_cache[n] = dm[bvk_cell_id*nao*nao + (i0+i)*nao + j0+j] * factor; + } + __syncthreads(); + + double grad_i[3] = {}; + double grad_j[3] = {}; + double sigma[9] = {}; + +#pragma unroll + for (int dm_i0 = 0; dm_i0 < nfi; dm_i0 += SLICE_SIZE_I) { +#pragma unroll + for (int dm_j0 = 0; dm_j0 < nfj; dm_j0 += SLICE_SIZE_J) { + + for (int b_index = b_start+ty; b_index <= b_stop; b_index += tile) { + for (int c_index = c_start+tx; c_index <= c_stop; c_index += tile) { + double x_start = a_center * c_dxyz_dabc[0] + b_index * c_dxyz_dabc[3] + c_index * c_dxyz_dabc[6]; + double y_start = a_center * c_dxyz_dabc[1] + b_index * c_dxyz_dabc[4] + c_index * c_dxyz_dabc[7]; + double z_start = a_center * c_dxyz_dabc[2] + b_index * c_dxyz_dabc[5] + c_index * c_dxyz_dabc[8]; + double x = x_start; + double y = y_start; + double z = z_start; + double x_xij = x - xij; + double y_yij = y - yij; + double z_zij = z - zij; + double e = theta_rr + aij * distance_squared(x_xij, y_yij, z_zij); + if (e > 50.) continue; // ~1e-22 + double gaussian_starting_point = exp(-e) * cc; + double cross_term_a = c_dxyz_dabc[0] * x_xij + c_dxyz_dabc[1] * y_yij + c_dxyz_dabc[2] * z_zij; + double recursion_factor_a_start = exp(-aij * (2 * cross_term_a + da_squared)); + + int mesh_bc = mesh_b * mesh_c; + int mesh_abc = mesh_a * mesh_bc; + // mod(negative_number, N) leads to a negative value. Adding a large + // multiplier to a_index, b_index, c_index to avoid the negative modulo. + // Images spreads in supmol are typically < 10. A shift of 100* images + // should be enough. + int64_t abc_idx_start = (a_center + 100 * mesh_a) % mesh_a * (int64_t)mesh_bc + + (b_index + 100 * mesh_b) % mesh_b * mesh_c + + (c_index + 100 * mesh_c) % mesh_c; + int64_t abc_idx = abc_idx_start; + double gaussian_xyz = gaussian_starting_point; + double recursion_factor_a = recursion_factor_a_start; + for (int a_index = a_center; a_index <= a_stop; a_index++, + gaussian_xyz *= recursion_factor_a, + recursion_factor_a *= exp_da_squared) { + if (fabs(gaussian_xyz) < negligible) break; + + double v = vxc_weights[abc_idx] * gaussian_xyz; + double i_deriv0[nfi]; + double i_deriv1[3*nfi]; + double x_xi = x - xi; + double y_yi = y - yi; + double z_zi = z - zi; + gto_cartesian
  • (i_deriv0, x_xi, y_yi, z_zi); + gto_deriv1
  • (i_deriv1, i_deriv0, x_xi, y_yi, z_zi, ai); + + double j_deriv0[nfj]; + double j_deriv1[3*nfj]; + double x_xj = x - xj; + double y_yj = y - yj; + double z_zj = z - zj; + gto_cartesian(j_deriv0, x_xj, y_yj, z_zj); + gto_deriv1(j_deriv1, j_deriv0, x_xj, y_yj, z_zj, aj); + double rho_i[3] = {}; +#pragma unroll + for (int i = dm_i0; i < min(dm_i0+SLICE_SIZE_I, nfi); ++i) { + double s = 0; +#pragma unroll + for (int j = dm_j0; j < min(dm_j0+SLICE_SIZE_J, nfj); ++j) { + s += dm_cache[i*nfj+j] * j_deriv0[j]; + } + rho_i[0] += s * i_deriv1[i ]; + rho_i[1] += s * i_deriv1[i+nfi ]; + rho_i[2] += s * i_deriv1[i+nfi*2]; + } + for (int n = 0; n < 3; ++n) { + rho_i[n] *= v; + grad_i[n] -= rho_i[n]; + } + + double rho_j[3] = {}; +#pragma unroll + for (int j = dm_j0; j < min(dm_j0+SLICE_SIZE_J, nfj); ++j) { + double s = 0; +#pragma unroll + for (int i = dm_i0; i < min(dm_i0+SLICE_SIZE_I, nfi); ++i) { + s += dm_cache[i*nfj+j] * i_deriv0[i]; + } + rho_j[0] += s * j_deriv1[j ]; + rho_j[1] += s * j_deriv1[j+nfj ]; + rho_j[2] += s * j_deriv1[j+nfj*2]; + } + for (int n = 0; n < 3; ++n) { + rho_j[n] *= v; + grad_j[n] -= rho_j[n]; + rho_i[n] += rho_j[n]; +// Grid-response contributions. +// Note that these grid coordinates are taken from the entire integration space, +// not restricted to the unit cell. In the NumInt implementation, numerical +// integration is performed only within the unit cell. In that case, only the +// grid response inside the unit cell contributes, yielding a simple term: +// sum_{r in unit cell} (\nabla rho) * Vxc * r. Although the expression below +// looks similar, it cannot be reused here because it is evaluated over a +// different integration domain. + sigma[n*3+0] += rho_i[n] * x; + sigma[n*3+1] += rho_i[n] * y; + sigma[n*3+2] += rho_i[n] * z; + } + + x += c_dxyz_dabc[0]; + y += c_dxyz_dabc[1]; + z += c_dxyz_dabc[2]; + abc_idx += mesh_bc; + if (abc_idx >= mesh_abc) { + abc_idx -= mesh_abc; + } + } + + x = x_start; + y = y_start; + z = z_start; + gaussian_xyz = gaussian_starting_point; + double inv_recursion_factor_a = exp_da_squared / recursion_factor_a_start; + abc_idx = abc_idx_start; + for (int a_index = a_center - 1; a_index >= a_start; a_index--, + inv_recursion_factor_a *= exp_da_squared) { + gaussian_xyz *= inv_recursion_factor_a; + if (fabs(gaussian_xyz) < negligible) break; + x -= c_dxyz_dabc[0]; + y -= c_dxyz_dabc[1]; + z -= c_dxyz_dabc[2]; + abc_idx -= mesh_bc; + if (abc_idx < 0) { + abc_idx += mesh_abc; + } + double v = vxc_weights[abc_idx] * gaussian_xyz; + double i_deriv0[nfi]; + double i_deriv1[3*nfi]; + double x_xi = x - xi; + double y_yi = y - yi; + double z_zi = z - zi; + gto_cartesian
  • (i_deriv0, x_xi, y_yi, z_zi); + gto_deriv1
  • (i_deriv1, i_deriv0, x_xi, y_yi, z_zi, ai); + + double j_deriv0[nfj]; + double j_deriv1[3*nfj]; + double x_xj = x - xj; + double y_yj = y - yj; + double z_zj = z - zj; + gto_cartesian(j_deriv0, x_xj, y_yj, z_zj); + gto_deriv1(j_deriv1, j_deriv0, x_xj, y_yj, z_zj, aj); + double rho_i[3] = {}; +#pragma unroll + for (int i = dm_i0; i < min(dm_i0+SLICE_SIZE_I, nfi); ++i) { + double s = 0; +#pragma unroll + for (int j = dm_j0; j < min(dm_j0+SLICE_SIZE_J, nfj); ++j) { + s += dm_cache[i*nfj+j] * j_deriv0[j]; + } + rho_i[0] += s * i_deriv1[i ]; + rho_i[1] += s * i_deriv1[i+nfi ]; + rho_i[2] += s * i_deriv1[i+nfi*2]; + } + for (int n = 0; n < 3; ++n) { + rho_i[n] *= v; + grad_i[n] -= rho_i[n]; + } + + double rho_j[3] = {}; +#pragma unroll + for (int j = dm_j0; j < min(dm_j0+SLICE_SIZE_J, nfj); ++j) { + double s = 0; +#pragma unroll + for (int i = dm_i0; i < min(dm_i0+SLICE_SIZE_I, nfi); ++i) { + s += dm_cache[i*nfj+j] * i_deriv0[i]; + } + rho_j[0] += s * j_deriv1[j ]; + rho_j[1] += s * j_deriv1[j+nfj ]; + rho_j[2] += s * j_deriv1[j+nfj*2]; + } + for (int n = 0; n < 3; ++n) { + rho_j[n] *= v; + grad_j[n] -= rho_j[n]; + rho_i[n] += rho_j[n]; + sigma[n*3+0] += rho_i[n] * x; + sigma[n*3+1] += rho_i[n] * y; + sigma[n*3+2] += rho_i[n] * z; + } + } + } } + } } + + for (int n = 0; n < 3; ++n) { + sigma[n*3+0] += grad_i[n] * xi + grad_j[n] * xj; + sigma[n*3+1] += grad_i[n] * yi + grad_j[n] * yj; + sigma[n*3+2] += grad_i[n] * zi + grad_j[n] * zj; + } + + for (int offset = 16; offset > 0; offset >>= 1) { + for (int n = 0; n < 3; ++n) { + grad_i[n] += __shfl_down_sync(0xffffffff, grad_i[n], offset); + grad_j[n] += __shfl_down_sync(0xffffffff, grad_j[n], offset); + } + for (int n = 0; n < 9; ++n) { + sigma[n] += __shfl_down_sync(0xffffffff, sigma[n], offset); + } + } + int lane = thread_id % WARP_SIZE; + int ish_cell0 = ish; + int bvk_cell_id = jsh / nbas; + int jsh_cell0 = jsh - nbas * bvk_cell_id; + int ia = bas[ish_cell0*BAS_SLOTS+ATOM_OF]; + int ja = bas[jsh_cell0*BAS_SLOTS+ATOM_OF]; + if (lane == 0) { + for (int n = 0; n < 3; ++n) { + atomicAdd(grad+ia*3+n, grad_i[n]); + atomicAdd(grad+ja*3+n, grad_j[n]); + } + for (int n = 0; n < 9; ++n) { + atomicAdd(strain+n, sigma[n]); + } + } +} + +extern "C" { +#define eval_lda_grad_kernel_case(li, lj, slice_i, slice_j) \ + case (li * LMAX1 + lj): \ + eval_lda_grad_kernel<<>>( \ + grad, strain, dm, vxc, *envs, bas_ij_idx, grid_frac_ranges, \ + da_squared, db_squared, dc_squared, mesh_a, mesh_b, mesh_c, npairs, \ + factor, negligible); \ + break + +int evaluate_lda_grad(double *grad, double *strain, double *dm, + double *vxc, double *placeholder, PBCIntEnvVars *envs, + double *dxyz_dabc, int li, int lj, int64_t *bas_ij_idx, + float2 *grid_frac_ranges, int *mesh, int npairs, + double factor, double negligible) +{ + int mesh_a = mesh[0]; + int mesh_b = mesh[1]; + int mesh_c = mesh[2]; + double da_squared = distance_squared(dxyz_dabc[0], dxyz_dabc[1], dxyz_dabc[2]); + double db_squared = distance_squared(dxyz_dabc[3], dxyz_dabc[4], dxyz_dabc[5]); + double dc_squared = distance_squared(dxyz_dabc[6], dxyz_dabc[7], dxyz_dabc[8]); + dim3 threads(16, 16); + switch (li * LMAX1 + lj) { + eval_lda_grad_kernel_case(0,0, 1, 1); + eval_lda_grad_kernel_case(1,0, 3, 1); + eval_lda_grad_kernel_case(1,1, 3, 3); + eval_lda_grad_kernel_case(2,0, 6, 1); + eval_lda_grad_kernel_case(2,1, 6, 3); + eval_lda_grad_kernel_case(2,2, 6, 6); + eval_lda_grad_kernel_case(3,0,10, 1); + eval_lda_grad_kernel_case(3,1,10, 3); + eval_lda_grad_kernel_case(3,2, 5, 6); + eval_lda_grad_kernel_case(3,3,10, 5); + eval_lda_grad_kernel_case(4,0,15, 1); + eval_lda_grad_kernel_case(4,1,15, 3); + eval_lda_grad_kernel_case(4,2, 8, 6); + eval_lda_grad_kernel_case(4,3, 8, 5); + eval_lda_grad_kernel_case(4,4, 8, 5); + } + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in eval_lda_grad_kernel: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} +} diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/eval_mgga_mat.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_mgga_mat.cu new file mode 100644 index 000000000..f64d3f1fa --- /dev/null +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_mgga_mat.cu @@ -0,0 +1,332 @@ +/* + * Copyright 2025-2026 The PySCF Developers. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include "gvhf-rys/vhf.cuh" +#include "constant_objects.cuh" +#include "cartesian.cuh" +#include "utils.cuh" + +#define TILE 4 +#define WARP_SIZE 32 +#define THREADS 64 + +template +__global__ static +void eval_mgga_mat_kernel(double *out, double *vrho_weights, double *vtau_weights, + PBCIntEnvVars envs, double *supmol_img_coords, + int *shl_pair_offsets, int64_t *dressed_bas_ij_idx, + int *grid_tile_index, int ntiles, int tiles_per_block, + double a_dot_b, double a_dot_c, double b_dot_c, + double da_squared, double db_squared, double dc_squared, + int mesh_a, int mesh_b, int mesh_c, double negligible) +{ + constexpr int threads = THREADS; + int thread_id = threadIdx.x; + int tile_id0 = blockIdx.x * tiles_per_block; + __shared__ int a_upper, b_upper, c_upper; + __shared__ double start_position_x, start_position_y, start_position_z; + __shared__ double vrho_cache[TILE*TILE*TILE]; + __shared__ double vtau_cache[TILE*TILE*TILE]; + +for (int tile_id = tile_id0; tile_id < min(tile_id0+tiles_per_block, ntiles); tile_id++) { + int tile_index = grid_tile_index[tile_id]; + int shl_pair0 = shl_pair_offsets[tile_id]; + int shl_pair1 = shl_pair_offsets[tile_id+1]; + int n_tiles_b = (mesh_b + TILE - 1) / TILE; + int n_tiles_c = (mesh_c + TILE - 1) / TILE; + int tile_ab_index = tile_index / n_tiles_c; + int tile_c_index = tile_index % n_tiles_c; + int tile_a_index = tile_ab_index / n_tiles_b; + int tile_b_index = tile_ab_index % n_tiles_b; + int a_start = tile_a_index * TILE; + int b_start = tile_b_index * TILE; + int c_start = tile_c_index * TILE; + + constexpr int nfi = (LI + 1) * (LI + 2) / 2; + constexpr int nfj = (LJ + 1) * (LJ + 2) / 2; + + int *bas = envs.bas; + double *env = envs.env; + int nbas = envs.nbas; + int bvk_nbas = envs.bvk_ncells * envs.nbas; + + if (thread_id == 0) { + start_position_x = c_dxyz_dabc[0] * a_start + c_dxyz_dabc[3] * b_start + c_dxyz_dabc[6] * c_start; + start_position_y = c_dxyz_dabc[1] * a_start + c_dxyz_dabc[4] * b_start + c_dxyz_dabc[7] * c_start; + start_position_z = c_dxyz_dabc[2] * a_start + c_dxyz_dabc[5] * b_start + c_dxyz_dabc[8] * c_start; + a_upper = min(a_start + TILE, mesh_a) - a_start; + b_upper = min(b_start + TILE, mesh_b) - b_start; + c_upper = min(c_start + TILE, mesh_c) - c_start; + } + + int a_idx = a_start + thread_id / (TILE*TILE); + int b_idx = b_start + thread_id / TILE % TILE; + int c_idx = c_start + thread_id % TILE; + if (a_idx < mesh_a && b_idx < mesh_b && c_idx < mesh_c) { + size_t abc_idx = (a_idx * mesh_b + b_idx) * (size_t)mesh_c + c_idx; + vrho_cache[thread_id] = vrho_weights[abc_idx]; + vtau_cache[thread_id] = vtau_weights[abc_idx] / 2; + } + __syncthreads(); + + for (int pair_id = shl_pair0+thread_id; pair_id < shl_pair1; pair_id += threads) { + int64_t bas_ij = dressed_bas_ij_idx[pair_id]; + int ish = bas_ij / NBAS_MAX; + int jsh = bas_ij % NBAS_MAX; + int latsum_idx = ish / nbas; + ish = ish - nbas * latsum_idx; + int jL = jsh / bvk_nbas; + jsh = jsh - bvk_nbas * jL; + double Lx = supmol_img_coords[latsum_idx*3+0]; + double Ly = supmol_img_coords[latsum_idx*3+1]; + double Lz = supmol_img_coords[latsum_idx*3+2]; + int expi = bas[ish*BAS_SLOTS+PTR_EXP]; + int expj = bas[jsh*BAS_SLOTS+PTR_EXP]; + int ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; + int rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; + double xi = env[ri+0] - Lx; + double yi = env[ri+1] - Ly; + double zi = env[ri+2] - Lz; + double xj = env[rj+0] - Lx + envs.img_coords[jL*3+0]; + double yj = env[rj+1] - Ly + envs.img_coords[jL*3+1]; + double zj = env[rj+2] - Lz + envs.img_coords[jL*3+2]; + double xjxi = xj - xi; + double yjyi = yj - yi; + double zjzi = zj - zi; + double rr_ij = distance_squared(xjxi, yjyi, zjzi); + double ai = env[expi]; + double aj = env[expj]; + double aij = ai + aj; + double aj_aij = aj / aij; + double theta_ij = ai * aj_aij; + double xij = xjxi * aj_aij + xi; + double yij = yjyi * aj_aij + yi; + double zij = zjzi * aj_aij + zi; + double x0 = start_position_x - xij; + double y0 = start_position_y - yij; + double z0 = start_position_z - zij; + double gaussian_exponent_at_reference = aij * distance_squared(x0, y0, z0); + double gaussian_starting_exponent = theta_ij * rr_ij + gaussian_exponent_at_reference; + if (gaussian_starting_exponent > 680.) { + continue; + } + int ish_cell0 = ish; + int bvk_cell_id = jsh / nbas; + int jsh_cell0 = jsh - nbas * bvk_cell_id; + double ci = env[bas[ish_cell0*BAS_SLOTS+PTR_COEFF]]; + double cj = env[bas[jsh_cell0*BAS_SLOTS+PTR_COEFF]]; + double cc = ci * cj; + if (ish_cell0 == jsh_cell0) { + cc *= .5; + } + double gaussian_starting_point = exp(-gaussian_starting_exponent) * cc; + double cross_term_a = c_dxyz_dabc[0] * x0 + c_dxyz_dabc[1] * y0 + c_dxyz_dabc[2] * z0; + double cross_term_b = c_dxyz_dabc[3] * x0 + c_dxyz_dabc[4] * y0 + c_dxyz_dabc[5] * z0; + double cross_term_c = c_dxyz_dabc[6] * x0 + c_dxyz_dabc[7] * y0 + c_dxyz_dabc[8] * z0; + double recursion_factor_a_start = exp(-aij * (2 * cross_term_a + da_squared)); + double recursion_factor_b_start = exp(-aij * (2 * cross_term_b + db_squared)); + double recursion_factor_c_start = exp(-aij * (2 * cross_term_c + dc_squared)); + double exp_da_squared = exp(-2 * aij * da_squared); + double exp_db_squared = exp(-2 * aij * db_squared); + double exp_dc_squared = exp(-2 * aij * dc_squared); + double exp_dadb = exp(-2 * aij * a_dot_b); + double exp_dadc = exp(-2 * aij * a_dot_c); + double exp_dbdc = exp(-2 * aij * b_dot_c); +#pragma unroll + for (int dm_i0 = 0; dm_i0 < nfi; dm_i0 += SLICE_SIZE_I) { +#pragma unroll + for (int dm_j0 = 0; dm_j0 < nfj; dm_j0 += SLICE_SIZE_J) { + double vj_cache[SLICE_SIZE_I * SLICE_SIZE_J]; +#pragma unroll + for (int n = 0; n < SLICE_SIZE_I*SLICE_SIZE_J; ++n) { + vj_cache[n] = 0; + } + + double x, y, z; + double recursion_factor_ab_pow_a = 1; + double recursion_factor_ac_pow_a = 1; + double recursion_factor_bc_pow_b = 1; + + if constexpr (is_non_orthogonal) { + } else { + x = start_position_x; + } + double gaussian_x = gaussian_starting_point; + double recursion_factor_a = recursion_factor_a_start; + for (int a_index = 0; a_index < a_upper; a_index++, + gaussian_x *= recursion_factor_a, + recursion_factor_a *= exp_da_squared) { + + if constexpr (is_non_orthogonal) { + recursion_factor_bc_pow_b = 1; + } else { + y = start_position_y; + } + double gaussian_xy = gaussian_x; + double recursion_factor_b = recursion_factor_b_start * recursion_factor_ab_pow_a; + for (int b_index = 0; b_index < b_upper; b_index++, + gaussian_xy *= recursion_factor_b, + recursion_factor_b *= exp_db_squared) { + + if constexpr (is_non_orthogonal) { + x = start_position_x + a_index * c_dxyz_dabc[0] + b_index * c_dxyz_dabc[3]; + y = start_position_y + a_index * c_dxyz_dabc[1] + b_index * c_dxyz_dabc[4]; + z = start_position_z + a_index * c_dxyz_dabc[2] + b_index * c_dxyz_dabc[5]; + } else { + z = start_position_z; + } + double gaussian_xyz = gaussian_xy; + double recursion_factor_c = recursion_factor_c_start * + recursion_factor_ac_pow_a * recursion_factor_bc_pow_b; + for (int c_index = 0; c_index < c_upper; c_index++, + gaussian_xyz *= recursion_factor_c, + recursion_factor_c *= exp_dc_squared) { + + if (fabs(gaussian_xyz) > negligible) { + double i_cartesian[nfi]; + double i_gradient[3*nfi]; + gto_cartesian
  • (i_cartesian, x - xi, y - yi, z - zi); + gto_deriv1
  • (i_gradient, i_cartesian, x - xi, y - yi, z - zi, ai); + rename_registers(i_cartesian, dm_i0, nfi, SLICE_SIZE_I); + rename_registers(i_gradient , dm_i0, nfi, SLICE_SIZE_I); + rename_registers(i_gradient+nfi , dm_i0, nfi, SLICE_SIZE_I); + rename_registers(i_gradient+nfi*2, dm_i0, nfi, SLICE_SIZE_I); + + double j_cartesian[nfj]; + double j_gradient[3*nfj]; + gto_cartesian(j_cartesian, x - xj, y - yj, z - zj); + gto_deriv1(j_gradient, j_cartesian, x - xj, y - yj, z - zj, aj); + rename_registers(j_cartesian, dm_j0, nfj, SLICE_SIZE_J); + rename_registers(j_gradient , dm_j0, nfj, SLICE_SIZE_J); + rename_registers(j_gradient+nfj , dm_j0, nfj, SLICE_SIZE_J); + rename_registers(j_gradient+nfj*2, dm_j0, nfj, SLICE_SIZE_J); + + int abc_index = a_index * TILE*TILE + b_index*TILE + c_index; + double rho_fac = gaussian_xyz * vrho_cache[abc_index]; + double tau_fac = gaussian_xyz * vtau_cache[abc_index] / 2; +#pragma unroll + for (int i = 0; i < SLICE_SIZE_I; ++i) { + if (SLICE_SIZE_I < nfi && dm_i0 + i >= nfi) break; + double s0 = rho_fac * i_cartesian[i]; + double s1 = tau_fac * i_gradient[i]; + double s2 = tau_fac * i_gradient[i+nfi]; + double s3 = tau_fac * i_gradient[i+nfi*2]; +#pragma unroll + for (int j = 0; j < SLICE_SIZE_J; ++j) { + if (SLICE_SIZE_J < nfj && dm_j0 + j >= nfj) break; + vj_cache[i*SLICE_SIZE_J+j] += s0 * j_cartesian[j]; + vj_cache[i*SLICE_SIZE_J+j] += s1 * j_gradient[j]; + vj_cache[i*SLICE_SIZE_J+j] += s2 * j_gradient[j+nfj]; + vj_cache[i*SLICE_SIZE_J+j] += s3 * j_gradient[j+nfj*2]; + } + } + } + if constexpr (is_non_orthogonal) { + x += c_dxyz_dabc[6]; + y += c_dxyz_dabc[7]; + z += c_dxyz_dabc[8]; + } else { + z += c_dxyz_dabc[8]; + } + } + if constexpr (is_non_orthogonal) { + recursion_factor_bc_pow_b *= exp_dbdc; + } else { + y += c_dxyz_dabc[4]; + } + } + if constexpr (is_non_orthogonal) { + recursion_factor_ab_pow_a *= exp_dadb; + recursion_factor_ac_pow_a *= exp_dadc; + } else { + x += c_dxyz_dabc[0]; + } + } + + size_t nao = envs.ao_loc[nbas]; + int i0 = envs.ao_loc[ish_cell0]; + int j0 = envs.ao_loc[jsh_cell0]; + double *pout = out + bvk_cell_id * nao * nao + (dm_i0+i0) * nao + dm_j0+j0; +#pragma unroll + for (int i = 0; i < SLICE_SIZE_I; ++i) { + if (SLICE_SIZE_I < nfi && dm_i0 + i >= nfi) break; +#pragma unroll + for (int j = 0; j < SLICE_SIZE_J; ++j) { + if (SLICE_SIZE_J < nfj && dm_j0 + j >= nfj) break; + atomicAdd(pout + i*nao+j, vj_cache[i*SLICE_SIZE_J+j]); + } } + } } + } +} +} + +extern "C" { +#define eval_mgga_mat_kernel_case(li, lj, slice_i, slice_j, non_orth) \ + case (li * LMAX1 + lj): \ + eval_mgga_mat_kernel<<>>( \ + out, vxc, vtau, *envs, supmol_img_coords, shl_pair_offsets, dressed_bas_ij_idx, \ + grid_tile_index, n_contributing_tiles, tiles_per_block, \ + a_dot_b, a_dot_c, b_dot_c, da_squared, db_squared, dc_squared, \ + mesh_a, mesh_b, mesh_c, negligible); \ + break + +int evaluate_mgga_mat(double *out, double *vxc, double *vtau, PBCIntEnvVars *envs, + double *dxyz_dabc, double *supmol_img_coords, + int i_angular, int j_angular, int tiles_per_block, + int *shl_pair_offsets, int64_t *dressed_bas_ij_idx, + int *grid_tile_index, int n_contributing_tiles, int *mesh, + double negligible) +{ + int mesh_a = mesh[0]; + int mesh_b = mesh[1]; + int mesh_c = mesh[2]; + int block_grid = (n_contributing_tiles + tiles_per_block-1) / tiles_per_block; + double a_dot_b = dxyz_dabc[0] * dxyz_dabc[3] + dxyz_dabc[1] * dxyz_dabc[4] + dxyz_dabc[2] * dxyz_dabc[5]; + double a_dot_c = dxyz_dabc[0] * dxyz_dabc[6] + dxyz_dabc[1] * dxyz_dabc[7] + dxyz_dabc[2] * dxyz_dabc[8]; + double b_dot_c = dxyz_dabc[3] * dxyz_dabc[6] + dxyz_dabc[4] * dxyz_dabc[7] + dxyz_dabc[5] * dxyz_dabc[8]; + double da_squared = distance_squared(dxyz_dabc[0], dxyz_dabc[1], dxyz_dabc[2]); + double db_squared = distance_squared(dxyz_dabc[3], dxyz_dabc[4], dxyz_dabc[5]); + double dc_squared = distance_squared(dxyz_dabc[6], dxyz_dabc[7], dxyz_dabc[8]); + switch (i_angular * LMAX1 + j_angular) { + eval_mgga_mat_kernel_case(0,0, 1, 1, 1); + eval_mgga_mat_kernel_case(1,0, 3, 1, 1); + eval_mgga_mat_kernel_case(1,1, 3, 3, 1); + eval_mgga_mat_kernel_case(2,0, 6, 1, 1); + eval_mgga_mat_kernel_case(2,1, 6, 3, 1); + eval_mgga_mat_kernel_case(2,2, 6, 6, 1); + eval_mgga_mat_kernel_case(3,0,10, 1, 1); + eval_mgga_mat_kernel_case(3,1,10, 3, 1); + eval_mgga_mat_kernel_case(3,2, 5, 6, 1); + eval_mgga_mat_kernel_case(3,3,10, 5, 1); + eval_mgga_mat_kernel_case(4,0,15, 1, 1); + eval_mgga_mat_kernel_case(4,1,15, 3, 1); + eval_mgga_mat_kernel_case(4,2,15, 3, 1); + eval_mgga_mat_kernel_case(4,3, 8, 5, 1); + eval_mgga_mat_kernel_case(4,4,15, 3, 1); + } + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in eval_mgga_mat_kernel: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} +} diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/eval_mgga_mat_v2.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_mgga_mat_v2.cu new file mode 100644 index 000000000..164f7f67b --- /dev/null +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_mgga_mat_v2.cu @@ -0,0 +1,338 @@ +/* + * Copyright 2025-2026 The PySCF Developers. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include "gvhf-rys/vhf.cuh" +#include "constant_objects.cuh" +#include "cartesian.cuh" +#include "utils.cuh" + +template +__global__ static +void eval_mgga_mat_kernel_v2(double *out, double *vrho_weights, double *vtau_weights, + PBCIntEnvVars envs, + int64_t *bas_ij_idx, float2 *grid_frac_ranges, + double da_squared, double db_squared, double dc_squared, + int mesh_a, int mesh_b, int mesh_c, int npairs, + double negligible) +{ + constexpr int tile = 16; + int tx = threadIdx.x; + int ty = threadIdx.y; + int thread_id = ty * tile + tx; + int pair_id = blockIdx.x; + + __shared__ int a_start, a_stop, a_center; + __shared__ int b_start, b_stop; + __shared__ int c_start, c_stop; + __shared__ uint32_t ij_offset; + __shared__ double cc, exp_da_squared; + __shared__ double xi, yi, zi; + __shared__ double xj, yj, zj; + __shared__ double xij, yij, zij, ai, aj, aij, theta_rr; + __shared__ double swap[8]; + + int *bas = envs.bas; + double *env = envs.env; + int nbas = envs.nbas; + int bvk_nbas = envs.bvk_ncells * envs.nbas; + + int64_t bas_ij = bas_ij_idx[pair_id]; + int ish = bas_ij / NBAS_MAX; + int jsh = bas_ij % NBAS_MAX; + int jL = jsh / bvk_nbas; + jsh = jsh - bvk_nbas * jL; + if (thread_id == 0) { + int ish_cell0 = ish; + int bvk_cell_id = jsh / nbas; + int jsh_cell0 = jsh - nbas * bvk_cell_id; + uint32_t nao = envs.ao_loc[nbas]; + uint32_t i0 = envs.ao_loc[ish_cell0]; + uint32_t j0 = envs.ao_loc[jsh_cell0]; + ij_offset = bvk_cell_id * nao * nao + i0 * nao + j0; + + int ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; + int rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; + ai = env[bas[ish_cell0*BAS_SLOTS+PTR_EXP]]; + aj = env[bas[jsh_cell0*BAS_SLOTS+PTR_EXP]]; + double ci = env[bas[ish_cell0*BAS_SLOTS+PTR_COEFF]]; + double cj = env[bas[jsh_cell0*BAS_SLOTS+PTR_COEFF]]; + cc = ci * cj; + if (ish_cell0 == jsh_cell0) { + cc *= .5; + } + aij = ai + aj; + double aj_aij = aj / aij; + xi = env[ri+0]; + yi = env[ri+1]; + zi = env[ri+2]; + xj = env[rj+0] + envs.img_coords[jL*3+0]; + yj = env[rj+1] + envs.img_coords[jL*3+1]; + zj = env[rj+2] + envs.img_coords[jL*3+2]; + double xjxi = xj - xi; + double yjyi = yj - yi; + double zjzi = zj - zi; + double rr = distance_squared(xjxi, yjyi, zjzi); + xij = xjxi * aj_aij + xi; + yij = yjyi * aj_aij + yi; + zij = zjzi * aj_aij + zi; + theta_rr = ai * aj_aij * rr; + + float2 range = grid_frac_ranges[pair_id]; + float xfrac_lower = range.x; + float xfrac_upper = range.y; + range = grid_frac_ranges[npairs+pair_id]; + float yfrac_lower = range.x; + float yfrac_upper = range.y; + range = grid_frac_ranges[npairs*2+pair_id]; + float zfrac_lower = range.x; + float zfrac_upper = range.y; + a_start = ceil(xfrac_lower * mesh_a); + b_start = ceil(yfrac_lower * mesh_b); + c_start = ceil(zfrac_lower * mesh_c); + a_stop = floor(xfrac_upper * mesh_a); + b_stop = floor(yfrac_upper * mesh_b); + c_stop = floor(zfrac_upper * mesh_c); + a_center = (a_start + a_stop) / 2; + exp_da_squared = exp(-2 * aij * da_squared); + } + __syncthreads(); + + constexpr int nfi = (LI + 1) * (LI + 2) / 2; + constexpr int nfj = (LJ + 1) * (LJ + 2) / 2; + +#pragma unroll + for (int dm_i0 = 0; dm_i0 < nfi; dm_i0 += SLICE_SIZE_I) { +#pragma unroll + for (int dm_j0 = 0; dm_j0 < nfj; dm_j0 += SLICE_SIZE_J) { + double vj_cache[SLICE_SIZE_I * SLICE_SIZE_J]; +#pragma unroll + for (int n = 0; n < SLICE_SIZE_I*SLICE_SIZE_J; ++n) { + vj_cache[n] = 0; + } + + for (int b_index = b_start+ty; b_index <= b_stop; b_index += tile) { + for (int c_index = c_start+tx; c_index <= c_stop; c_index += tile) { + double x_start = a_center * c_dxyz_dabc[0] + b_index * c_dxyz_dabc[3] + c_index * c_dxyz_dabc[6]; + double y_start = a_center * c_dxyz_dabc[1] + b_index * c_dxyz_dabc[4] + c_index * c_dxyz_dabc[7]; + double z_start = a_center * c_dxyz_dabc[2] + b_index * c_dxyz_dabc[5] + c_index * c_dxyz_dabc[8]; + double x = x_start; + double y = y_start; + double z = z_start; + double x_xij = x - xij; + double y_yij = y - yij; + double z_zij = z - zij; + double e = theta_rr + aij * distance_squared(x_xij, y_yij, z_zij); + if (e > 50.) continue; // ~1e-22 + double gaussian_starting_point = exp(-e) * cc; + double cross_term_a = c_dxyz_dabc[0] * x_xij + c_dxyz_dabc[1] * y_yij + c_dxyz_dabc[2] * z_zij; + double recursion_factor_a_start = exp(-aij * (2 * cross_term_a + da_squared)); + + int mesh_bc = mesh_b * mesh_c; + int mesh_abc = mesh_a * mesh_bc; + // mod(negative_number, N) leads to a negative value. Adding a large + // multiplier to a_index, b_index, c_index to avoid the negative modulo. + // Images spreads in supmol are typically < 10. A shift of 100* images + // should be enough. + int64_t abc_idx_start = (a_center + 100 * mesh_a) % mesh_a * (int64_t)mesh_bc + + (b_index + 100 * mesh_b) % mesh_b * mesh_c + + (c_index + 100 * mesh_c) % mesh_c; + int64_t abc_idx = abc_idx_start; + double gaussian_xyz = gaussian_starting_point; + double recursion_factor_a = recursion_factor_a_start; + for (int a_index = a_center; a_index <= a_stop; a_index++, + gaussian_xyz *= recursion_factor_a, + recursion_factor_a *= exp_da_squared) { + if (fabs(gaussian_xyz) < negligible) break; + + double rho_fac = vrho_weights[abc_idx] * gaussian_xyz; + double tau_fac = vtau_weights[abc_idx] * gaussian_xyz / 2; + double i_cartesian[nfi]; + double i_gradient[3*nfi]; + double x_xi = x - xi; + double y_yi = y - yi; + double z_zi = z - zi; + gto_cartesian
  • (i_cartesian, x_xi, y_yi, z_zi); + gto_deriv1
  • (i_gradient, i_cartesian, x_xi, y_yi, z_zi, ai); + rename_registers(i_cartesian, dm_i0, nfi, SLICE_SIZE_I); + rename_registers(i_gradient , dm_i0, nfi, SLICE_SIZE_I); + rename_registers(i_gradient+nfi , dm_i0, nfi, SLICE_SIZE_I); + rename_registers(i_gradient+nfi*2, dm_i0, nfi, SLICE_SIZE_I); + + double x_xj = x - xj; + double y_yj = y - yj; + double z_zj = z - zj; + double j_cartesian[nfj]; + double j_gradient[3*nfj]; + gto_cartesian(j_cartesian, x_xj, y_yj, z_zj); + gto_deriv1(j_gradient, j_cartesian, x_xj, y_yj, z_zj, aj); + rename_registers(j_cartesian, dm_j0, nfj, SLICE_SIZE_J); + rename_registers(j_gradient , dm_j0, nfj, SLICE_SIZE_J); + rename_registers(j_gradient+nfj , dm_j0, nfj, SLICE_SIZE_J); + rename_registers(j_gradient+nfj*2, dm_j0, nfj, SLICE_SIZE_J); +#pragma unroll + for (int i = 0; i < SLICE_SIZE_I; ++i) { + if (SLICE_SIZE_I < nfi && dm_i0 + i >= nfi) break; + double s0 = rho_fac * i_cartesian[i]; + double s1 = tau_fac * i_gradient[i]; + double s2 = tau_fac * i_gradient[i+nfi]; + double s3 = tau_fac * i_gradient[i+nfi*2]; +#pragma unroll + for (int j = 0; j < SLICE_SIZE_J; ++j) { + if (SLICE_SIZE_J < nfj && dm_j0 + j >= nfj) break; + vj_cache[i*SLICE_SIZE_J+j] += s0 * j_cartesian[j]; + vj_cache[i*SLICE_SIZE_J+j] += s1 * j_gradient[j]; + vj_cache[i*SLICE_SIZE_J+j] += s2 * j_gradient[j+nfj]; + vj_cache[i*SLICE_SIZE_J+j] += s3 * j_gradient[j+nfj*2]; + } + } + x += c_dxyz_dabc[0]; + y += c_dxyz_dabc[1]; + z += c_dxyz_dabc[2]; + abc_idx += mesh_bc; + if (abc_idx >= mesh_abc) { + abc_idx -= mesh_abc; + } + } + + x = x_start; + y = y_start; + z = z_start; + gaussian_xyz = gaussian_starting_point; + double inv_recursion_factor_a = exp_da_squared / recursion_factor_a_start; + abc_idx = abc_idx_start; + for (int a_index = a_center - 1; a_index >= a_start; a_index--, + inv_recursion_factor_a *= exp_da_squared) { + gaussian_xyz *= inv_recursion_factor_a; + if (fabs(gaussian_xyz) < negligible) break; + x -= c_dxyz_dabc[0]; + y -= c_dxyz_dabc[1]; + z -= c_dxyz_dabc[2]; + abc_idx -= mesh_bc; + if (abc_idx < 0) { + abc_idx += mesh_abc; + } + double rho_fac = vrho_weights[abc_idx] * gaussian_xyz; + double tau_fac = vtau_weights[abc_idx] * gaussian_xyz / 2; + double i_cartesian[nfi]; + double i_gradient[3*nfi]; + double x_xi = x - xi; + double y_yi = y - yi; + double z_zi = z - zi; + gto_cartesian
  • (i_cartesian, x_xi, y_yi, z_zi); + gto_deriv1
  • (i_gradient, i_cartesian, x_xi, y_yi, z_zi, ai); + rename_registers(i_cartesian, dm_i0, nfi, SLICE_SIZE_I); + rename_registers(i_gradient , dm_i0, nfi, SLICE_SIZE_I); + rename_registers(i_gradient+nfi , dm_i0, nfi, SLICE_SIZE_I); + rename_registers(i_gradient+nfi*2, dm_i0, nfi, SLICE_SIZE_I); + + double x_xj = x - xj; + double y_yj = y - yj; + double z_zj = z - zj; + double j_cartesian[nfj]; + double j_gradient[3*nfj]; + gto_cartesian(j_cartesian, x_xj, y_yj, z_zj); + gto_deriv1(j_gradient, j_cartesian, x_xj, y_yj, z_zj, aj); + rename_registers(j_cartesian, dm_j0, nfj, SLICE_SIZE_J); + rename_registers(j_gradient , dm_j0, nfj, SLICE_SIZE_J); + rename_registers(j_gradient+nfj , dm_j0, nfj, SLICE_SIZE_J); + rename_registers(j_gradient+nfj*2, dm_j0, nfj, SLICE_SIZE_J); +#pragma unroll + for (int i = 0; i < SLICE_SIZE_I; ++i) { + if (SLICE_SIZE_I < nfi && dm_i0 + i >= nfi) break; + double s0 = rho_fac * i_cartesian[i]; + double s1 = tau_fac * i_gradient[i]; + double s2 = tau_fac * i_gradient[i+nfi]; + double s3 = tau_fac * i_gradient[i+nfi*2]; +#pragma unroll + for (int j = 0; j < SLICE_SIZE_J; ++j) { + if (SLICE_SIZE_J < nfj && dm_j0 + j >= nfj) break; + vj_cache[i*SLICE_SIZE_J+j] += s0 * j_cartesian[j]; + vj_cache[i*SLICE_SIZE_J+j] += s1 * j_gradient[j]; + vj_cache[i*SLICE_SIZE_J+j] += s2 * j_gradient[j+nfj]; + vj_cache[i*SLICE_SIZE_J+j] += s3 * j_gradient[j+nfj*2]; + } + } + } + } } + + uint32_t nao = envs.ao_loc[nbas]; + double *pout = out + ij_offset + dm_i0 * nao + dm_j0; +#pragma unroll + for (int i = 0; i < SLICE_SIZE_I; ++i) { + if (SLICE_SIZE_I < nfi && dm_i0 + i >= nfi) break; +#pragma unroll + for (int j = 0; j < SLICE_SIZE_J; ++j) { + if (SLICE_SIZE_J < nfj && dm_j0 + j >= nfj) break; + double val = reduce(vj_cache[i*SLICE_SIZE_J+j], swap, thread_id); + if (thread_id == 0) { + atomicAdd(pout + i*nao+j, val); + } + } + } + } } +} + +extern "C" { +#define eval_mgga_mat_kernel_v2_case(li, lj, slice_i, slice_j) \ + case (li * LMAX1 + lj): \ + eval_mgga_mat_kernel_v2<<>>( \ + out, vxc, tau, *envs, bas_ij_idx, grid_frac_ranges, \ + da_squared, db_squared, dc_squared, mesh_a, mesh_b, mesh_c, npairs, \ + negligible); \ + break + +int evaluate_mgga_mat_v2(double *out, double *vxc, double *tau, PBCIntEnvVars *envs, + double *dxyz_dabc, int li, int lj, int64_t *bas_ij_idx, + float2 *grid_frac_ranges, int *mesh, int npairs, + double negligible) +{ + int mesh_a = mesh[0]; + int mesh_b = mesh[1]; + int mesh_c = mesh[2]; + double da_squared = distance_squared(dxyz_dabc[0], dxyz_dabc[1], dxyz_dabc[2]); + double db_squared = distance_squared(dxyz_dabc[3], dxyz_dabc[4], dxyz_dabc[5]); + double dc_squared = distance_squared(dxyz_dabc[6], dxyz_dabc[7], dxyz_dabc[8]); + dim3 threads(16, 16); + switch (li * LMAX1 + lj) { + eval_mgga_mat_kernel_v2_case(0,0, 1, 1); + eval_mgga_mat_kernel_v2_case(1,0, 3, 1); + eval_mgga_mat_kernel_v2_case(1,1, 3, 3); + eval_mgga_mat_kernel_v2_case(2,0, 6, 1); + eval_mgga_mat_kernel_v2_case(2,1, 6, 3); + eval_mgga_mat_kernel_v2_case(2,2, 6, 6); + eval_mgga_mat_kernel_v2_case(3,0,10, 1); + eval_mgga_mat_kernel_v2_case(3,1,10, 3); + eval_mgga_mat_kernel_v2_case(3,2,10, 6); + eval_mgga_mat_kernel_v2_case(3,3,10, 5); + eval_mgga_mat_kernel_v2_case(4,0,15, 1); + eval_mgga_mat_kernel_v2_case(4,1,15, 3); + eval_mgga_mat_kernel_v2_case(4,2,15, 3); + eval_mgga_mat_kernel_v2_case(4,3, 8, 5); + eval_mgga_mat_kernel_v2_case(4,4,15, 3); + } + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in eval_mgga_mat_kernel: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} +} diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/eval_mgga_strain_grad.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_mgga_strain_grad.cu new file mode 100644 index 000000000..2db72f30e --- /dev/null +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_mgga_strain_grad.cu @@ -0,0 +1,531 @@ +/* + * Copyright 2025-2026 The PySCF Developers. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include "gvhf-rys/vhf.cuh" +#include "constant_objects.cuh" +#include "cartesian.cuh" +#include "utils.cuh" + +#define TILE 4 +#define WARP_SIZE 32 +#define THREADS 64 + +template +__global__ static +void eval_mgga_grad_kernel(double *grad, double *strain, double *dm, + double *vrho_weights, double *vtau_weights, PBCIntEnvVars envs, + int64_t *bas_ij_idx, float2 *grid_frac_ranges, + double da_squared, double db_squared, double dc_squared, + int mesh_a, int mesh_b, int mesh_c, int npairs, + double factor, double negligible) +{ + constexpr int tile = 16; + int tx = threadIdx.x; + int ty = threadIdx.y; + int thread_id = ty * tile + tx; + int pair_id = blockIdx.x; + + constexpr int nfi = (LI + 1) * (LI + 2) / 2; + constexpr int nfj = (LJ + 1) * (LJ + 2) / 2; + + __shared__ int a_start, a_stop, a_center; + __shared__ int b_start, b_stop; + __shared__ int c_start, c_stop; + __shared__ double cc, exp_da_squared; + __shared__ double xi, yi, zi; + __shared__ double xj, yj, zj; + __shared__ double xij, yij, zij, ai, aj, aij, theta_rr; + __shared__ double xjxi, yjyi, zjzi; + __shared__ double dm_cache[nfi*nfj]; + + int *bas = envs.bas; + double *env = envs.env; + int nbas = envs.nbas; + int bvk_nbas = envs.bvk_ncells * envs.nbas; + + int64_t bas_ij = bas_ij_idx[pair_id]; + int ish = bas_ij / NBAS_MAX; + int jsh = bas_ij % NBAS_MAX; + int jL = jsh / bvk_nbas; + jsh = jsh - bvk_nbas * jL; + if (thread_id == 0) { + int ish_cell0 = ish; + int bvk_cell_id = jsh / nbas; + int jsh_cell0 = jsh - nbas * bvk_cell_id; + int ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; + int rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; + ai = env[bas[ish_cell0*BAS_SLOTS+PTR_EXP]]; + aj = env[bas[jsh_cell0*BAS_SLOTS+PTR_EXP]]; + double ci = env[bas[ish_cell0*BAS_SLOTS+PTR_COEFF]]; + double cj = env[bas[jsh_cell0*BAS_SLOTS+PTR_COEFF]]; + cc = ci * cj; + if (ish_cell0 == jsh_cell0) { + cc *= .5; + } + aij = ai + aj; + double aj_aij = aj / aij; + xi = env[ri+0]; + yi = env[ri+1]; + zi = env[ri+2]; + xj = env[rj+0] + envs.img_coords[jL*3+0]; + yj = env[rj+1] + envs.img_coords[jL*3+1]; + zj = env[rj+2] + envs.img_coords[jL*3+2]; + xjxi = xj - xi; + yjyi = yj - yi; + zjzi = zj - zi; + double rr = distance_squared(xjxi, yjyi, zjzi); + xij = xjxi * aj_aij + xi; + yij = yjyi * aj_aij + yi; + zij = zjzi * aj_aij + zi; + theta_rr = ai * aj_aij * rr; + + float2 range = grid_frac_ranges[pair_id]; + float xfrac_lower = range.x; + float xfrac_upper = range.y; + range = grid_frac_ranges[npairs+pair_id]; + float yfrac_lower = range.x; + float yfrac_upper = range.y; + range = grid_frac_ranges[npairs*2+pair_id]; + float zfrac_lower = range.x; + float zfrac_upper = range.y; + a_start = ceil(xfrac_lower * mesh_a); + b_start = ceil(yfrac_lower * mesh_b); + c_start = ceil(zfrac_lower * mesh_c); + a_stop = floor(xfrac_upper * mesh_a); + b_stop = floor(yfrac_upper * mesh_b); + c_stop = floor(zfrac_upper * mesh_c); + a_center = (a_start + a_stop) / 2; + exp_da_squared = exp(-2 * aij * da_squared); + } + __syncthreads(); + + for (int n = thread_id; n < nfi * nfj; n += THREADS) { + int ish_cell0 = ish; + int bvk_cell_id = jsh / nbas; + int jsh_cell0 = jsh - nbas * bvk_cell_id; + size_t nao = envs.ao_loc[nbas]; + int i0 = envs.ao_loc[ish_cell0]; + int j0 = envs.ao_loc[jsh_cell0]; + int i = n * c_div_nf[LJ]; + int j = n - nfj * i; + dm_cache[n] = dm[bvk_cell_id*nao*nao + (i0+i)*nao + j0+j] * factor; + } + __syncthreads(); + + constexpr int XX = 0; + constexpr int XY = 1; + constexpr int XZ = 2; + constexpr int YX = 1; + constexpr int YY = 3; + constexpr int YZ = 4; + constexpr int ZX = 2; + constexpr int ZY = 4; + constexpr int ZZ = 5; + + double grad_i[3] = {}; + double grad_j[3] = {}; + double sigma[9] = {}; + +#pragma unroll + for (int dm_i0 = 0; dm_i0 < nfi; dm_i0 += SLICE_SIZE_I) { +#pragma unroll + for (int dm_j0 = 0; dm_j0 < nfj; dm_j0 += SLICE_SIZE_J) { + + for (int b_index = b_start+ty; b_index <= b_stop; b_index += tile) { + for (int c_index = c_start+tx; c_index <= c_stop; c_index += tile) { + double x_start = a_center * c_dxyz_dabc[0] + b_index * c_dxyz_dabc[3] + c_index * c_dxyz_dabc[6]; + double y_start = a_center * c_dxyz_dabc[1] + b_index * c_dxyz_dabc[4] + c_index * c_dxyz_dabc[7]; + double z_start = a_center * c_dxyz_dabc[2] + b_index * c_dxyz_dabc[5] + c_index * c_dxyz_dabc[8]; + double x = x_start; + double y = y_start; + double z = z_start; + double x_xij = x - xij; + double y_yij = y - yij; + double z_zij = z - zij; + double e = theta_rr + aij * distance_squared(x_xij, y_yij, z_zij); + if (e > 50.) continue; // ~1e-22 + double gaussian_starting_point = exp(-e) * cc; + double cross_term_a = c_dxyz_dabc[0] * x_xij + c_dxyz_dabc[1] * y_yij + c_dxyz_dabc[2] * z_zij; + double recursion_factor_a_start = exp(-aij * (2 * cross_term_a + da_squared)); + + int mesh_bc = mesh_b * mesh_c; + int mesh_abc = mesh_a * mesh_bc; + // mod(negative_number, N) leads to a negative value. Adding a large + // multiplier to a_index, b_index, c_index to avoid the negative modulo. + // Images spreads in supmol are typically < 10. A shift of 100* images + // should be enough. + int64_t abc_idx_start = (a_center + 100 * mesh_a) % mesh_a * (int64_t)mesh_bc + + (b_index + 100 * mesh_b) % mesh_b * mesh_c + + (c_index + 100 * mesh_c) % mesh_c; + int64_t abc_idx = abc_idx_start; + double gaussian_xyz = gaussian_starting_point; + double recursion_factor_a = recursion_factor_a_start; + for (int a_index = a_center; a_index <= a_stop; a_index++, + gaussian_xyz *= recursion_factor_a, + recursion_factor_a *= exp_da_squared) { + if (fabs(gaussian_xyz) < negligible) break; + + double rho_fac = vrho_weights[abc_idx] * gaussian_xyz; + double x_xi = x - xi; + double y_yi = y - yi; + double z_zi = z - zi; + double i_deriv0[nfi]; + double i_deriv1[3*nfi]; + gto_cartesian
  • (i_deriv0, x_xi, y_yi, z_zi); + gto_deriv1
  • (i_deriv1, i_deriv0, x_xi, y_yi, z_zi, ai); + + double x_xj = x - xj; + double y_yj = y - yj; + double z_zj = z - zj; + double j_deriv0[nfj]; + double j_deriv1[3*nfj]; + gto_cartesian(j_deriv0, x_xj, y_yj, z_zj); + gto_deriv1(j_deriv1, j_deriv0, x_xj, y_yj, z_zj, aj); + double rho_i[3] = {}; +#pragma unroll + for (int i = dm_i0; i < min(dm_i0+SLICE_SIZE_I, nfi); ++i) { + double s0 = 0; +#pragma unroll + for (int j = dm_j0; j < min(dm_j0+SLICE_SIZE_J, nfj); ++j) { + s0 += dm_cache[i*nfj+j] * j_deriv0[j]; + } + rho_i[0] += s0 * i_deriv1[i ]; + rho_i[1] += s0 * i_deriv1[i+nfi ]; + rho_i[2] += s0 * i_deriv1[i+nfi*2]; + } + for (int n = 0; n < 3; ++n) { + rho_i[n] *= rho_fac; + grad_i[n] -= rho_i[n]; + } + + double rho_j[3] = {}; +#pragma unroll + for (int j = dm_j0; j < min(dm_j0+SLICE_SIZE_J, nfj); ++j) { + double s0 = 0; +#pragma unroll + for (int i = dm_i0; i < min(dm_i0+SLICE_SIZE_I, nfi); ++i) { + s0 += dm_cache[i*nfj+j] * i_deriv0[i]; + } + rho_j[0] += s0 * j_deriv1[j ]; + rho_j[1] += s0 * j_deriv1[j+nfj ]; + rho_j[2] += s0 * j_deriv1[j+nfj*2]; + } + for (int n = 0; n < 3; ++n) { + rho_j[n] *= rho_fac; + grad_j[n] -= rho_j[n]; + rho_i[n] += rho_j[n]; + } + + double tau_fac = vtau_weights[abc_idx] * gaussian_xyz / 2; + double i_deriv2[6*nfi]; + gto_deriv2
  • (i_deriv2, x_xi, y_yi, z_zi, ai); + double tau[3] = {}; +#pragma unroll + for (int i = dm_i0; i < min(dm_i0+SLICE_SIZE_I, nfi); ++i) { + double sx = 0; + double sy = 0; + double sz = 0; +#pragma unroll + for (int j = dm_j0; j < min(dm_j0+SLICE_SIZE_J, nfj); ++j) { + double dm_fac = dm_cache[i*nfj+j]; + sx += dm_fac * j_deriv1[j ]; + sy += dm_fac * j_deriv1[j+nfj ]; + sz += dm_fac * j_deriv1[j+nfj*2]; + } + tau[0] += sx * i_deriv2[i+nfi*XX]; + tau[1] += sx * i_deriv2[i+nfi*XY]; + tau[2] += sx * i_deriv2[i+nfi*XZ]; + tau[0] += sy * i_deriv2[i+nfi*YX]; + tau[1] += sy * i_deriv2[i+nfi*YY]; + tau[2] += sy * i_deriv2[i+nfi*YZ]; + tau[0] += sz * i_deriv2[i+nfi*ZX]; + tau[1] += sz * i_deriv2[i+nfi*ZY]; + tau[2] += sz * i_deriv2[i+nfi*ZZ]; + } + for (int n = 0; n < 3; ++n) { + tau[n] *= tau_fac; + grad_i[n] -= tau[n]; + rho_i[n] += tau[n]; + } + + double j_deriv2[6*nfj]; + gto_deriv2(j_deriv2, x_xj, y_yj, z_zj, aj); + tau[0] = 0; + tau[1] = 0; + tau[2] = 0; +#pragma unroll + for (int j = dm_j0; j < min(dm_j0+SLICE_SIZE_J, nfj); ++j) { + double sx = 0; + double sy = 0; + double sz = 0; +#pragma unroll + for (int i = dm_i0; i < min(dm_i0+SLICE_SIZE_I, nfi); ++i) { + double dm_fac = dm_cache[i*nfj+j]; + sx += dm_fac * i_deriv1[i ]; + sy += dm_fac * i_deriv1[i+nfi ]; + sz += dm_fac * i_deriv1[i+nfi*2]; + } + tau[0] += sx * j_deriv2[j+nfj*XX]; + tau[1] += sx * j_deriv2[j+nfj*XY]; + tau[2] += sx * j_deriv2[j+nfj*XZ]; + tau[0] += sy * j_deriv2[j+nfj*YX]; + tau[1] += sy * j_deriv2[j+nfj*YY]; + tau[2] += sy * j_deriv2[j+nfj*YZ]; + tau[0] += sz * j_deriv2[j+nfj*ZX]; + tau[1] += sz * j_deriv2[j+nfj*ZY]; + tau[2] += sz * j_deriv2[j+nfj*ZZ]; + } + for (int n = 0; n < 3; ++n) { + tau[n] *= tau_fac; + grad_j[n] -= tau[n]; + rho_i[n] += tau[n]; + sigma[n*3+0] += rho_i[n] * x; + sigma[n*3+1] += rho_i[n] * y; + sigma[n*3+2] += rho_i[n] * z; + } + + x += c_dxyz_dabc[0]; + y += c_dxyz_dabc[1]; + z += c_dxyz_dabc[2]; + abc_idx += mesh_bc; + if (abc_idx >= mesh_abc) { + abc_idx -= mesh_abc; + } + } + + x = x_start; + y = y_start; + z = z_start; + gaussian_xyz = gaussian_starting_point; + double inv_recursion_factor_a = exp_da_squared / recursion_factor_a_start; + abc_idx = abc_idx_start; + for (int a_index = a_center - 1; a_index >= a_start; a_index--, + inv_recursion_factor_a *= exp_da_squared) { + gaussian_xyz *= inv_recursion_factor_a; + if (fabs(gaussian_xyz) < negligible) break; + x -= c_dxyz_dabc[0]; + y -= c_dxyz_dabc[1]; + z -= c_dxyz_dabc[2]; + abc_idx -= mesh_bc; + if (abc_idx < 0) { + abc_idx += mesh_abc; + } + double rho_fac = vrho_weights[abc_idx] * gaussian_xyz; + double x_xi = x - xi; + double y_yi = y - yi; + double z_zi = z - zi; + double i_deriv0[nfi]; + double i_deriv1[3*nfi]; + gto_cartesian
  • (i_deriv0, x_xi, y_yi, z_zi); + gto_deriv1
  • (i_deriv1, i_deriv0, x_xi, y_yi, z_zi, ai); + + double x_xj = x - xj; + double y_yj = y - yj; + double z_zj = z - zj; + double j_deriv0[nfj]; + double j_deriv1[3*nfj]; + gto_cartesian(j_deriv0, x_xj, y_yj, z_zj); + gto_deriv1(j_deriv1, j_deriv0, x_xj, y_yj, z_zj, aj); + double rho_i[3] = {}; +#pragma unroll + for (int i = dm_i0; i < min(dm_i0+SLICE_SIZE_I, nfi); ++i) { + double s0 = 0; +#pragma unroll + for (int j = dm_j0; j < min(dm_j0+SLICE_SIZE_J, nfj); ++j) { + s0 += dm_cache[i*nfj+j] * j_deriv0[j]; + } + rho_i[0] += s0 * i_deriv1[i ]; + rho_i[1] += s0 * i_deriv1[i+nfi ]; + rho_i[2] += s0 * i_deriv1[i+nfi*2]; + } + for (int n = 0; n < 3; ++n) { + rho_i[n] *= rho_fac; + grad_i[n] -= rho_i[n]; + } + + double rho_j[3] = {}; +#pragma unroll + for (int j = dm_j0; j < min(dm_j0+SLICE_SIZE_J, nfj); ++j) { + double s0 = 0; +#pragma unroll + for (int i = dm_i0; i < min(dm_i0+SLICE_SIZE_I, nfi); ++i) { + s0 += dm_cache[i*nfj+j] * i_deriv0[i]; + } + rho_j[0] += s0 * j_deriv1[j ]; + rho_j[1] += s0 * j_deriv1[j+nfj ]; + rho_j[2] += s0 * j_deriv1[j+nfj*2]; + } + for (int n = 0; n < 3; ++n) { + rho_j[n] *= rho_fac; + grad_j[n] -= rho_j[n]; + rho_i[n] += rho_j[n]; + } + + double tau_fac = vtau_weights[abc_idx] * gaussian_xyz / 2; + double i_deriv2[6*nfi]; + gto_deriv2
  • (i_deriv2, x_xi, y_yi, z_zi, ai); + double tau[3] = {}; +#pragma unroll + for (int i = dm_i0; i < min(dm_i0+SLICE_SIZE_I, nfi); ++i) { + double sx = 0; + double sy = 0; + double sz = 0; +#pragma unroll + for (int j = dm_j0; j < min(dm_j0+SLICE_SIZE_J, nfj); ++j) { + double dm_fac = dm_cache[i*nfj+j]; + sx += dm_fac * j_deriv1[j ]; + sy += dm_fac * j_deriv1[j+nfj ]; + sz += dm_fac * j_deriv1[j+nfj*2]; + } + tau[0] += sx * i_deriv2[i+nfi*XX]; + tau[1] += sx * i_deriv2[i+nfi*XY]; + tau[2] += sx * i_deriv2[i+nfi*XZ]; + tau[0] += sy * i_deriv2[i+nfi*YX]; + tau[1] += sy * i_deriv2[i+nfi*YY]; + tau[2] += sy * i_deriv2[i+nfi*YZ]; + tau[0] += sz * i_deriv2[i+nfi*ZX]; + tau[1] += sz * i_deriv2[i+nfi*ZY]; + tau[2] += sz * i_deriv2[i+nfi*ZZ]; + } + for (int n = 0; n < 3; ++n) { + tau[n] *= tau_fac; + grad_i[n] -= tau[n]; + rho_i[n] += tau[n]; + } + + double j_deriv2[6*nfj]; + gto_deriv2(j_deriv2, x_xj, y_yj, z_zj, aj); + tau[0] = 0; + tau[1] = 0; + tau[2] = 0; +#pragma unroll + for (int j = dm_j0; j < min(dm_j0+SLICE_SIZE_J, nfj); ++j) { + double sx = 0; + double sy = 0; + double sz = 0; +#pragma unroll + for (int i = dm_i0; i < min(dm_i0+SLICE_SIZE_I, nfi); ++i) { + double dm_fac = dm_cache[i*nfj+j]; + sx += dm_fac * i_deriv1[i ]; + sy += dm_fac * i_deriv1[i+nfi ]; + sz += dm_fac * i_deriv1[i+nfi*2]; + } + tau[0] += sx * j_deriv2[j+nfj*XX]; + tau[1] += sx * j_deriv2[j+nfj*XY]; + tau[2] += sx * j_deriv2[j+nfj*XZ]; + tau[0] += sy * j_deriv2[j+nfj*YX]; + tau[1] += sy * j_deriv2[j+nfj*YY]; + tau[2] += sy * j_deriv2[j+nfj*YZ]; + tau[0] += sz * j_deriv2[j+nfj*ZX]; + tau[1] += sz * j_deriv2[j+nfj*ZY]; + tau[2] += sz * j_deriv2[j+nfj*ZZ]; + } + for (int n = 0; n < 3; ++n) { + tau[n] *= tau_fac; + grad_j[n] -= tau[n]; + rho_i[n] += tau[n]; + sigma[n*3+0] += rho_i[n] * x; + sigma[n*3+1] += rho_i[n] * y; + sigma[n*3+2] += rho_i[n] * z; + } + } + } } + } } + + for (int n = 0; n < 3; ++n) { + sigma[n*3+0] += grad_i[n] * xi + grad_j[n] * xj; + sigma[n*3+1] += grad_i[n] * yi + grad_j[n] * yj; + sigma[n*3+2] += grad_i[n] * zi + grad_j[n] * zj; + } + + for (int offset = 16; offset > 0; offset >>= 1) { + for (int n = 0; n < 3; ++n) { + grad_i[n] += __shfl_down_sync(0xffffffff, grad_i[n], offset); + grad_j[n] += __shfl_down_sync(0xffffffff, grad_j[n], offset); + } + for (int n = 0; n < 9; ++n) { + sigma[n] += __shfl_down_sync(0xffffffff, sigma[n], offset); + } + } + int lane = thread_id % WARP_SIZE; + int ish_cell0 = ish; + int bvk_cell_id = jsh / nbas; + int jsh_cell0 = jsh - nbas * bvk_cell_id; + int ia = bas[ish_cell0*BAS_SLOTS+ATOM_OF]; + int ja = bas[jsh_cell0*BAS_SLOTS+ATOM_OF]; + if (lane == 0) { + for (int n = 0; n < 3; ++n) { + atomicAdd(grad+ia*3+n, grad_i[n]); + atomicAdd(grad+ja*3+n, grad_j[n]); + } + for (int n = 0; n < 9; ++n) { + atomicAdd(strain+n, sigma[n]); + } + } +} + +extern "C" { +#define eval_mgga_grad_kernel_case(li, lj, slice_i, slice_j) \ + case (li * LMAX1 + lj): \ + eval_mgga_grad_kernel<<>>( \ + grad, strain, dm, vxc, tau, *envs, bas_ij_idx, grid_frac_ranges, \ + da_squared, db_squared, dc_squared, mesh_a, mesh_b, mesh_c, npairs, \ + factor, negligible); \ + break + +int evaluate_mgga_grad(double *grad, double *strain, double *dm, + double *vxc, double *tau, PBCIntEnvVars *envs, + double *dxyz_dabc, int li, int lj, int64_t *bas_ij_idx, + float2 *grid_frac_ranges, int *mesh, int npairs, + double factor, double negligible) +{ + int mesh_a = mesh[0]; + int mesh_b = mesh[1]; + int mesh_c = mesh[2]; + double da_squared = distance_squared(dxyz_dabc[0], dxyz_dabc[1], dxyz_dabc[2]); + double db_squared = distance_squared(dxyz_dabc[3], dxyz_dabc[4], dxyz_dabc[5]); + double dc_squared = distance_squared(dxyz_dabc[6], dxyz_dabc[7], dxyz_dabc[8]); + dim3 threads(16, 16); + switch (li * LMAX1 + lj) { + eval_mgga_grad_kernel_case(0,0, 1, 1); + eval_mgga_grad_kernel_case(1,0, 3, 1); + eval_mgga_grad_kernel_case(1,1, 3, 3); + eval_mgga_grad_kernel_case(2,0, 6, 1); + eval_mgga_grad_kernel_case(2,1, 6, 3); + eval_mgga_grad_kernel_case(2,2, 6, 3); + eval_mgga_grad_kernel_case(3,0,10, 1); + eval_mgga_grad_kernel_case(3,1, 5, 3); + eval_mgga_grad_kernel_case(3,2, 5, 3); + eval_mgga_grad_kernel_case(3,3, 5, 3); + eval_mgga_grad_kernel_case(4,0, 8, 1); + eval_mgga_grad_kernel_case(4,1, 5, 3); + eval_mgga_grad_kernel_case(4,2, 5, 3); + eval_mgga_grad_kernel_case(4,3, 3, 5); + eval_mgga_grad_kernel_case(4,4, 3, 5); + } + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in eval_mgga_grad_kernel: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} +} diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/eval_tau.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_tau.cu new file mode 100644 index 000000000..06d9a6f23 --- /dev/null +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/eval_tau.cu @@ -0,0 +1,399 @@ +/* + * Copyright 2025-2026 The PySCF Developers. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include "gvhf-rys/vhf.cuh" +#include "constant_objects.cuh" +#include "cartesian.cuh" +#include "utils.cuh" + +#define TILE 4 +#define WARP_SIZE 32 +#define THREADS 64 + +template +__global__ static +void eval_tau_kernel(double *density, double *tau, double *dm, PBCIntEnvVars envs, + double *supmol_img_coords, double factor, + int *shl_pair_offsets, int64_t *dressed_bas_ij_idx, + int *grid_tile_index, int ntiles, int tiles_per_block, + double a_dot_b, double a_dot_c, double b_dot_c, + double da_squared, double db_squared, double dc_squared, + int mesh_a, int mesh_b, int mesh_c, double negligible) +{ + constexpr int threads = THREADS; + constexpr int WARPS = THREADS / WARP_SIZE; + int thread_id = threadIdx.x; + int tile_id0 = blockIdx.x * tiles_per_block; + __shared__ int a_upper, b_upper, c_upper; + __shared__ double start_position_x, start_position_y, start_position_z; + __shared__ double rho_value[TILE*TILE*TILE*WARPS]; + __shared__ double tau_value[TILE*TILE*TILE*WARPS]; + + constexpr int nfi = (LI + 1) * (LI + 2) / 2; + constexpr int nfj = (LJ + 1) * (LJ + 2) / 2; + +for (int tile_id = tile_id0; tile_id < min(tile_id0+tiles_per_block, ntiles); tile_id++) { + int tile_index = grid_tile_index[tile_id]; + int shl_pair0 = shl_pair_offsets[tile_id]; + int shl_pair1 = shl_pair_offsets[tile_id+1]; + int n_tiles_b = (mesh_b + TILE - 1) / TILE; + int n_tiles_c = (mesh_c + TILE - 1) / TILE; + int tile_ab_index = tile_index / n_tiles_c; + int tile_c_index = tile_index % n_tiles_c; + int tile_a_index = tile_ab_index / n_tiles_b; + int tile_b_index = tile_ab_index % n_tiles_b; + int a_start = tile_a_index * TILE; + int b_start = tile_b_index * TILE; + int c_start = tile_c_index * TILE; + + int *bas = envs.bas; + double *env = envs.env; + int nbas = envs.nbas; + int bvk_nbas = envs.bvk_ncells * envs.nbas; + + if (thread_id == 0) { + start_position_x = c_dxyz_dabc[0] * a_start + c_dxyz_dabc[3] * b_start + c_dxyz_dabc[6] * c_start; + start_position_y = c_dxyz_dabc[1] * a_start + c_dxyz_dabc[4] * b_start + c_dxyz_dabc[7] * c_start; + start_position_z = c_dxyz_dabc[2] * a_start + c_dxyz_dabc[5] * b_start + c_dxyz_dabc[8] * c_start; + a_upper = min(a_start + TILE, mesh_a) - a_start; + b_upper = min(b_start + TILE, mesh_b) - b_start; + c_upper = min(c_start + TILE, mesh_c) - c_start; + } + + int lane = thread_id % WARP_SIZE; + int warp = thread_id / WARP_SIZE; + for (int n = thread_id; n < TILE*TILE*TILE*WARPS; n += threads) { + rho_value[n] = 0; + tau_value[n] = 0; + } + __syncthreads(); + + for (int pair_id = shl_pair0+thread_id; pair_id < shl_pair1+thread_id; pair_id += threads) { + int ish = 0; + int jsh = 0; + if (pair_id < shl_pair1) { + int64_t bas_ij = dressed_bas_ij_idx[pair_id]; + ish = bas_ij / NBAS_MAX; + jsh = bas_ij % NBAS_MAX; + } + int latsum_idx = ish / nbas; + ish = ish - nbas * latsum_idx; + int jL = jsh / bvk_nbas; + jsh = jsh - bvk_nbas * jL; + int ish_cell0 = ish; + int bvk_cell_id = jsh / nbas; + int jsh_cell0 = jsh - nbas * bvk_cell_id; + double Lx = supmol_img_coords[latsum_idx*3+0]; + double Ly = supmol_img_coords[latsum_idx*3+1]; + double Lz = supmol_img_coords[latsum_idx*3+2]; + int expi = bas[ish_cell0*BAS_SLOTS+PTR_EXP]; + int expj = bas[jsh_cell0*BAS_SLOTS+PTR_EXP]; + int ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; + int rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; + double xi = env[ri+0] - Lx; + double yi = env[ri+1] - Ly; + double zi = env[ri+2] - Lz; + double xj = env[rj+0] - Lx + envs.img_coords[jL*3+0]; + double yj = env[rj+1] - Ly + envs.img_coords[jL*3+1]; + double zj = env[rj+2] - Lz + envs.img_coords[jL*3+2]; + double xjxi = xj - xi; + double yjyi = yj - yi; + double zjzi = zj - zi; + double rr_ij = distance_squared(xjxi, yjyi, zjzi); + double ai = env[expi]; + double aj = env[expj]; + double aij = ai + aj; + double aj_aij = aj / aij; + double theta_ij = ai * aj_aij; + double xij = xjxi * aj_aij + xi; + double yij = yjyi * aj_aij + yi; + double zij = zjzi * aj_aij + zi; + double x0 = start_position_x - xij; + double y0 = start_position_y - yij; + double z0 = start_position_z - zij; + double gaussian_exponent_at_reference = aij * distance_squared(x0, y0, z0); + double gaussian_starting_exponent = theta_ij * rr_ij + gaussian_exponent_at_reference; + double gaussian_starting_point = 0.; + double recursion_factor_a_start = 0.; + double recursion_factor_b_start = 0.; + double recursion_factor_c_start = 0.; + double exp_da_squared = 0.; + double exp_db_squared = 0.; + double exp_dc_squared = 0.; + double exp_dadb = 0.; + double exp_dadc = 0.; + double exp_dbdc = 0.; + if (gaussian_starting_exponent < 680.) { + double ci = env[bas[ish*BAS_SLOTS+PTR_COEFF]]; + double cj = env[bas[jsh*BAS_SLOTS+PTR_COEFF]]; + double cc = ci * cj; + gaussian_starting_point = exp(-gaussian_starting_exponent); + gaussian_starting_point *= factor * cc; + if (ish == jsh_cell0) { + gaussian_starting_point *= 0.5; + } + double cross_term_a = c_dxyz_dabc[0] * x0 + c_dxyz_dabc[1] * y0 + c_dxyz_dabc[2] * z0; + double cross_term_b = c_dxyz_dabc[3] * x0 + c_dxyz_dabc[4] * y0 + c_dxyz_dabc[5] * z0; + double cross_term_c = c_dxyz_dabc[6] * x0 + c_dxyz_dabc[7] * y0 + c_dxyz_dabc[8] * z0; + recursion_factor_a_start = exp(-aij * (2 * cross_term_a + da_squared)); + recursion_factor_b_start = exp(-aij * (2 * cross_term_b + db_squared)); + recursion_factor_c_start = exp(-aij * (2 * cross_term_c + dc_squared)); + exp_da_squared = exp(-2 * aij * da_squared); + exp_db_squared = exp(-2 * aij * db_squared); + exp_dc_squared = exp(-2 * aij * dc_squared); + exp_dadb = exp(-2 * aij * a_dot_b); + exp_dadc = exp(-2 * aij * a_dot_c); + exp_dbdc = exp(-2 * aij * b_dot_c); + } +#pragma unroll + for (int dm_i0 = 0; dm_i0 < nfi; dm_i0 += SLICE_SIZE_I) { +#pragma unroll + for (int dm_j0 = 0; dm_j0 < nfj; dm_j0 += SLICE_SIZE_J) { + size_t nao = envs.ao_loc[nbas]; + int i0 = envs.ao_loc[ish_cell0]; + int j0 = envs.ao_loc[jsh_cell0]; + size_t ij_offset = bvk_cell_id * nao * nao + (dm_i0+i0) * nao + dm_j0+j0; + double dm_cache[SLICE_SIZE_I * SLICE_SIZE_J]; + if (pair_id < shl_pair1) { +#pragma unroll + for (int i = 0; i < SLICE_SIZE_I; ++i) { + if (SLICE_SIZE_I < nfi && dm_i0 + i >= nfi) break; +#pragma unroll + for (int j = 0; j < SLICE_SIZE_J; ++j) { + if (SLICE_SIZE_J < nfj && dm_j0 + j >= nfj) break; + dm_cache[i*SLICE_SIZE_J+j] = dm[ij_offset + i*nao+j]; + } } + } + + double x, y, z; + double recursion_factor_ab_pow_a = 1; + double recursion_factor_ac_pow_a = 1; + double recursion_factor_bc_pow_b = 1; + + if constexpr (is_non_orthogonal) { + // recursion_factor_ab_pow_a = 1; + // recursion_factor_ac_pow_a = 1; + } else { + x = start_position_x; + } + double gaussian_x = gaussian_starting_point; + double recursion_factor_a = recursion_factor_a_start; + for (int a_index = 0; a_index < a_upper; a_index++, + gaussian_x *= recursion_factor_a, + recursion_factor_a *= exp_da_squared) { + + if constexpr (is_non_orthogonal) { + recursion_factor_bc_pow_b = 1; + } else { + y = start_position_y; + } + double gaussian_xy = gaussian_x; + double recursion_factor_b = recursion_factor_b_start * recursion_factor_ab_pow_a; + for (int b_index = 0; b_index < b_upper; b_index++, + gaussian_xy *= recursion_factor_b, + recursion_factor_b *= exp_db_squared) { + + if constexpr (is_non_orthogonal) { + x = start_position_x + a_index * c_dxyz_dabc[0] + b_index * c_dxyz_dabc[3]; + y = start_position_y + a_index * c_dxyz_dabc[1] + b_index * c_dxyz_dabc[4]; + z = start_position_z + a_index * c_dxyz_dabc[2] + b_index * c_dxyz_dabc[5]; + } else { + z = start_position_z; + } + double gaussian_xyz = gaussian_xy; + double recursion_factor_c = recursion_factor_c_start * + recursion_factor_ac_pow_a * recursion_factor_bc_pow_b; + for (int c_index = 0; c_index < c_upper; c_index++, + gaussian_xyz *= recursion_factor_c, + recursion_factor_c *= exp_dc_squared) { + + double rho = 0; + double val = 0; + if (pair_id < shl_pair1 && fabs(gaussian_xyz) > negligible) { + double i_cartesian[nfi]; + double i_gradient[3*nfi]; + gto_cartesian
  • (i_cartesian, x - xi, y - yi, z - zi); + gto_deriv1
  • (i_gradient, i_cartesian, x - xi, y - yi, z - zi, ai); + rename_registers(i_cartesian, dm_i0, nfi, SLICE_SIZE_I); + rename_registers(i_gradient , dm_i0, nfi, SLICE_SIZE_I); + rename_registers(i_gradient+nfi , dm_i0, nfi, SLICE_SIZE_I); + rename_registers(i_gradient+nfi*2, dm_i0, nfi, SLICE_SIZE_I); + + double j_cartesian[nfj]; + double j_gradient[3*nfj]; + gto_cartesian(j_cartesian, x - xj, y - yj, z - zj); + gto_deriv1(j_gradient, j_cartesian, x - xj, y - yj, z - zj, aj); + rename_registers(j_cartesian, dm_j0, nfj, SLICE_SIZE_J); + rename_registers(j_gradient , dm_j0, nfj, SLICE_SIZE_J); + rename_registers(j_gradient+nfj , dm_j0, nfj, SLICE_SIZE_J); + rename_registers(j_gradient+nfj*2, dm_j0, nfj, SLICE_SIZE_J); +#pragma unroll + for (int i = 0; i < SLICE_SIZE_I; ++i) { + if (SLICE_SIZE_I < nfi && dm_i0 + i >= nfi) break; + double s0 = 0; + double s1 = 0; + double s2 = 0; + double s3 = 0; +#pragma unroll + for (int j = 0; j < SLICE_SIZE_J; ++j) { + if (SLICE_SIZE_J < nfj && dm_j0 + j >= nfj) break; + double dm_val = dm_cache[i * SLICE_SIZE_J + j]; + s0 += dm_val * j_cartesian[j]; + s1 += dm_val * j_gradient[j ]; + s2 += dm_val * j_gradient[j+nfj ]; + s3 += dm_val * j_gradient[j+nfj*2]; + } + rho += s0 * i_cartesian[i]; + val += s1 * i_gradient[i ]; + val += s2 * i_gradient[i+nfi ]; + val += s3 * i_gradient[i+nfi*2]; + } + rho *= gaussian_xyz; + val *= gaussian_xyz; + } + for (int offset = WARP_SIZE/2; offset > 0; offset >>= 1) { + rho += __shfl_down_sync(0xffffffff, rho, offset); + val += __shfl_down_sync(0xffffffff, val, offset); + } + if (lane == 0) { + int abc_index = a_index * TILE*TILE + b_index*TILE + c_index; + rho_value[abc_index+TILE*TILE*TILE*warp] += rho; + tau_value[abc_index+TILE*TILE*TILE*warp] += val; + } + if constexpr (is_non_orthogonal) { + x += c_dxyz_dabc[6]; + y += c_dxyz_dabc[7]; + z += c_dxyz_dabc[8]; + } else { + z += c_dxyz_dabc[8]; + } + } + if constexpr (is_non_orthogonal) { + recursion_factor_bc_pow_b *= exp_dbdc; + } else { + y += c_dxyz_dabc[4]; + } + } + if constexpr (is_non_orthogonal) { + recursion_factor_ab_pow_a *= exp_dadb; + recursion_factor_ac_pow_a *= exp_dadc; + } else { + x += c_dxyz_dabc[0]; + } + } + } } + } + __syncthreads(); + + int a_idx = a_start + thread_id / (TILE*TILE); + int b_idx = b_start + thread_id / TILE % TILE; + int c_idx = c_start + thread_id % TILE; + if (a_idx < mesh_a && b_idx < mesh_b && c_idx < mesh_c) { + double rho = rho_value[thread_id]; + double val = tau_value[thread_id]; + for (int i = 1; i < WARPS; i++) { + rho += rho_value[thread_id+i*TILE*TILE*TILE]; + val += tau_value[thread_id+i*TILE*TILE*TILE]; + } + size_t abc_idx = (a_idx * mesh_b + b_idx) * (size_t)mesh_c + c_idx; + // update the density.real, skip the imaginary part + atomicAdd(density + abc_idx*2, rho); + atomicAdd(tau + abc_idx*2, val/2); + } + __syncthreads(); +} +} + +extern "C" { +#define eval_tau_kernel_case(li, lj, slice_i, slice_j, non_orth) \ + case (li * LMAX1 + lj): \ + eval_tau_kernel<<>>( \ + density, tau, dm, *envs, supmol_img_coords, factor, \ + shl_pair_offsets, dressed_bas_ij_idx, \ + grid_tile_index, n_contributing_tiles, tiles_per_block, \ + a_dot_b, a_dot_c, b_dot_c, da_squared, db_squared, dc_squared, \ + mesh_a, mesh_b, mesh_c, negligible); \ + break + +int evaluate_tau(double *density, double *tau, double *dm, PBCIntEnvVars *envs, + double *dxyz_dabc, double *supmol_img_coords, + int i_angular, int j_angular, int tiles_per_block, + int *shl_pair_offsets, int64_t *dressed_bas_ij_idx, + int *grid_tile_index, int n_contributing_tiles, int *mesh, + double factor, double negligible) +{ + int mesh_a = mesh[0]; + int mesh_b = mesh[1]; + int mesh_c = mesh[2]; + int block_grid = (n_contributing_tiles + tiles_per_block-1) / tiles_per_block; + double a_dot_b = dxyz_dabc[0] * dxyz_dabc[3] + dxyz_dabc[1] * dxyz_dabc[4] + dxyz_dabc[2] * dxyz_dabc[5]; + double a_dot_c = dxyz_dabc[0] * dxyz_dabc[6] + dxyz_dabc[1] * dxyz_dabc[7] + dxyz_dabc[2] * dxyz_dabc[8]; + double b_dot_c = dxyz_dabc[3] * dxyz_dabc[6] + dxyz_dabc[4] * dxyz_dabc[7] + dxyz_dabc[5] * dxyz_dabc[8]; + double da_squared = distance_squared(dxyz_dabc[0], dxyz_dabc[1], dxyz_dabc[2]); + double db_squared = distance_squared(dxyz_dabc[3], dxyz_dabc[4], dxyz_dabc[5]); + double dc_squared = distance_squared(dxyz_dabc[6], dxyz_dabc[7], dxyz_dabc[8]); + if (fabs(dxyz_dabc[1]) < 1e-14 && fabs(dxyz_dabc[2]) < 1e-14 && + fabs(dxyz_dabc[3]) < 1e-14 && fabs(dxyz_dabc[5]) < 1e-14 && + fabs(dxyz_dabc[6]) < 1e-14 && fabs(dxyz_dabc[8]) < 1e-14) { + switch (i_angular * LMAX1 + j_angular) { + eval_tau_kernel_case(0,0, 1, 1, 0); + eval_tau_kernel_case(1,0, 3, 1, 0); + eval_tau_kernel_case(1,1, 3, 3, 0); + eval_tau_kernel_case(2,0, 6, 1, 0); + eval_tau_kernel_case(2,1, 6, 3, 0); + eval_tau_kernel_case(2,2, 6, 6, 1); + eval_tau_kernel_case(3,0,10, 1, 0); + eval_tau_kernel_case(3,1,10, 3, 1); + eval_tau_kernel_case(3,2, 5, 6, 1); + eval_tau_kernel_case(3,3,10, 5, 1); + eval_tau_kernel_case(4,0,15, 1, 0); + eval_tau_kernel_case(4,1,15, 3, 1); + eval_tau_kernel_case(4,2, 8, 6, 1); + eval_tau_kernel_case(4,3, 8, 5, 1); + eval_tau_kernel_case(4,4, 5, 5, 1); + } + } else { + switch (i_angular * LMAX1 + j_angular) { + eval_tau_kernel_case(0,0, 1, 1, 1); + eval_tau_kernel_case(1,0, 3, 1, 1); + eval_tau_kernel_case(1,1, 3, 3, 1); + eval_tau_kernel_case(2,0, 6, 1, 1); + eval_tau_kernel_case(2,1, 6, 3, 1); + eval_tau_kernel_case(2,2, 6, 6, 1); + eval_tau_kernel_case(3,0,10, 1, 1); + eval_tau_kernel_case(3,1,10, 3, 1); + eval_tau_kernel_case(3,2, 5, 6, 1); + eval_tau_kernel_case(3,3,10, 5, 1); + eval_tau_kernel_case(4,0,15, 1, 1); + eval_tau_kernel_case(4,1,15, 3, 1); + eval_tau_kernel_case(4,2, 8, 6, 1); + eval_tau_kernel_case(4,3, 8, 5, 1); + eval_tau_kernel_case(4,4, 5, 5, 1); + } + } + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in eval_tau_kernel: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} +} diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/screen.cu b/gpu4pyscf/lib/multigrid/multigrid_v3/screen.cu new file mode 100644 index 000000000..9fb1662b9 --- /dev/null +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/screen.cu @@ -0,0 +1,582 @@ +/* + * Copyright 2025-2026 The PySCF Developers. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include +#include +#include "gvhf-rys/vhf.cuh" +#include "constant_objects.cuh" +#include "cartesian.cuh" +#include "utils.cuh" + +#define TILE 4 +#define THREADS 256 + +template +__device__ static +T estimate_rcut(int li, int lj, T x, T aij, T xpi, T xpj, T log_factor) +{ + // let s = r - Rp + // rho[r-Rp] ~ ci*cj * exp(-theta*(Ri-Rj)**2) * (s+Rp-Ri)**li * (s+Rp-Rj)**lj * exp(-aij*s**2) + // ~= ovlp * (s+Rp-Ri)**li * (s+Rp-Rj)**lj * exp(-aij*s**2) + // radius can be solved using fixed iteration + // radius = (log(ovlp/precision * (s+Rpi)**li * (s+Rpj)**lj) / aij)**.5 + T aij_ss = log_factor + li * std::log(x + std::abs(xpi)) + lj * std::log(x + std::abs(xpj)); + return std::sqrt(max(aij_ss, static_cast(0)) / aij); +} + +template +__device__ inline +void accumulate(T lower, T upper, T c, T& min_val, T& max_val) +{ + T a = c * lower; + T b = c * upper; + min_val += min(a, b); + max_val += max(a, b); +} + +__global__ static +void grid_ranges_kernel(float2 *grid_frac_ranges, float *pair_ke, + float *Ecut_by_shell, PBCIntEnvVars envs, + int64_t *bas_ij_idx, int li_inc, int lj_inc, + int npairs, float log_threshold, + float undressed_threshold, float ke_max) +{ + int pair_id = blockIdx.x * blockDim.x + threadIdx.x; + if (pair_id >= npairs) return; + + int *bas = envs.bas; + double *env = envs.env; + int nbas = envs.nbas; + int bvk_nbas = envs.nbas * envs.bvk_ncells; + int64_t bas_ij = bas_ij_idx[pair_id]; + int ish = bas_ij / NBAS_MAX; + int jsh = bas_ij % NBAS_MAX; + int jL = jsh / bvk_nbas; + jsh = jsh % bvk_nbas; + // li_inc and lj_inc to account for derivatives + int li = bas[ish*BAS_SLOTS+ANG_OF] + li_inc; + int lj = bas[jsh*BAS_SLOTS+ANG_OF] + lj_inc; + int expi = bas[ish*BAS_SLOTS+PTR_EXP]; + int expj = bas[jsh*BAS_SLOTS+PTR_EXP]; + int ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; + int rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; + float ai = env[expi]; + float aj = env[expj]; + float aij = ai + aj; + float aj_aij = aj / aij; + float xi = env[ri+0]; + float yi = env[ri+1]; + float zi = env[ri+2]; + float xj = env[rj+0] + envs.img_coords[jL*3+0]; + float yj = env[rj+1] + envs.img_coords[jL*3+1]; + float zj = env[rj+2] + envs.img_coords[jL*3+2]; + float xjxi = xj - xi; + float yjyi = yj - yi; + float zjzi = zj - zi; + float rr = distance_squared(xjxi, yjyi, zjzi); + float xp = xjxi * aj_aij + xi; + float yp = yjyi * aj_aij + yi; + float zp = zjzi * aj_aij + zi; + float theta_rr = ai * aj_aij * rr; + float ci = env[bas[ish*BAS_SLOTS+PTR_COEFF]]; + float cj = env[bas[jsh*BAS_SLOTS+PTR_COEFF]]; + float log_cicj = logf(fabsf(ci * cj)); + float derivative_penalty = li_inc * logf(2*ai) + lj_inc * logf(2*aj); + float xpi = xp - xi; + float xpj = xp - xj; + float ypi = yp - yi; + float ypj = yp - yj; + float zpi = zp - zi; + float zpj = zp - zj; + + // let s = r - Rp + // rho[r-Rp] ~ ci*cj * exp(-theta*(Ri-Rj)**2) * (s+Rp-Ri)**li * (s+Rp-Rj)**lj * exp(-aij*s**2) + float log_factor = log_cicj + 1.717f - 1.5f * logf(aij) - log_threshold - theta_rr; + log_factor += derivative_penalty; + // initial guess + float log_r = 2.303f; // log(10) + float radius = sqrtf(max(log_factor + (li+lj)*log_r, 1e-20f) / aij); + // to encounter the integral over remaining space ~ int_radius^inf 4*pi*r^2 exp(-aij*r^2); + float penalty = 6.283f * radius / (2*aij); + // To accurately integrate a gaussian, the required resolution (Ngrid/Bohr) ~ 2*a**.5 + float resolution = 2*sqrtf(aij); + log_factor += max(logf(max(penalty, 12.56f*radius*radius/resolution)), 0.f); + float x_cut = estimate_rcut(li, lj, radius, aij, xpi, xpj, log_factor); + float y_cut = estimate_rcut(li, lj, radius, aij, ypi, ypj, log_factor); + float z_cut = estimate_rcut(li, lj, radius, aij, zpi, zpj, log_factor); + + float b00 = c_reciprocal_lattice_vectors[0]; + float b01 = c_reciprocal_lattice_vectors[1]; + float b02 = c_reciprocal_lattice_vectors[2]; + float b10 = c_reciprocal_lattice_vectors[3]; + float b11 = c_reciprocal_lattice_vectors[4]; + float b12 = c_reciprocal_lattice_vectors[5]; + float b20 = c_reciprocal_lattice_vectors[6]; + float b21 = c_reciprocal_lattice_vectors[7]; + float b22 = c_reciprocal_lattice_vectors[8]; + + float xp_frac = xp * b00 + yp * b01 + zp * b02; + float yp_frac = xp * b10 + yp * b11 + zp * b12; + float zp_frac = xp * b20 + yp * b21 + zp * b22; + + float bnorm_0 = sqrtf(distance_squared(b00, b01, b02)); + float bnorm_1 = sqrtf(distance_squared(b10, b11, b12)); + float bnorm_2 = sqrtf(distance_squared(b20, b21, b22)); + + float xcut_frac = x_cut * bnorm_0; + float ycut_frac = y_cut * bnorm_1; + float zcut_frac = z_cut * bnorm_2; + + float2 *xfrac_range = grid_frac_ranges; + float2 *yfrac_range = grid_frac_ranges + npairs; + float2 *zfrac_range = grid_frac_ranges + npairs * 2; + xfrac_range[pair_id] = {xp_frac - xcut_frac, xp_frac + xcut_frac}; + yfrac_range[pair_id] = {yp_frac - ycut_frac, yp_frac + ycut_frac}; + zfrac_range[pair_id] = {zp_frac - zcut_frac, zp_frac + zcut_frac}; + + // When cutoff radius is 0, the contribution of this orbital pair is small. + // By setting its pair_ke to 0, this orbital pair will be discarded when + // filtering orbitals in _partition_ke_for_fft function. + if (x_cut < 1e-3 || y_cut < 1e-3 || z_cut < 1e-3) { + pair_ke[pair_id] = -1.f; + } else { + float ish_ke = Ecut_by_shell[ish]; + float jsh_ke = Ecut_by_shell[jsh % nbas]; + float ke_two_centers = max(ish_ke, jsh_ke); + if (ri == rj && jL == 0) { + // Higher resolution is required for orbitals located on the same center. + // Ecut ~= 2 * ke_two_centers + // (Ecut/2/aij)**((li+lj)/2) * exp(-Ecut/(2*aij)) ~ cell.threshold + //float log_factor = (li+lj)*.5f * logf(ke_two_centers/aij) - + // logf(undressed_threshold); + float log_factor = -logf(undressed_threshold); + float ke_raw = log_factor * aij * 2; + ke_raw = min(ke_raw, ke_max); + pair_ke[pair_id] = max(ke_raw, ke_two_centers); + } else { + pair_ke[pair_id] = ke_two_centers; + } + } +} + +__global__ static +void grid_range_to_tiles_kernel(int *grid_tile_idx, int64_t *dressed_bas_ij, + int64_t *bas_ij_idx, float2 *grid_frac_ranges, + int nimgs_x, int nimgs_y, int nimgs_z, + int mesh_x, int mesh_y, int mesh_z, int npairs, + int nbas, int *head) +{ + int pair_id = blockIdx.x * blockDim.x + threadIdx.x; + if (pair_id >= npairs) return; + + int64_t bas_ij = bas_ij_idx[pair_id]; + float2 range = grid_frac_ranges[pair_id]; + float xfrac_lower = range.x; + float xfrac_upper = range.y; + range = grid_frac_ranges[npairs+pair_id]; + float yfrac_lower = range.x; + float yfrac_upper = range.y; + range = grid_frac_ranges[npairs*2+pair_id]; + float zfrac_lower = range.x; + float zfrac_upper = range.y; + + int tiles_x = (mesh_x + TILE - 1) / TILE; + int tiles_y = (mesh_y + TILE - 1) / TILE; + int tiles_z = (mesh_z + TILE - 1) / TILE; + int tile_size_x = min(TILE, mesh_x); + int tile_size_y = min(TILE, mesh_y); + int tile_size_z = min(TILE, mesh_z); + + int img_x_lower = floor(xfrac_lower); + int img_y_lower = floor(yfrac_lower); + int img_z_lower = floor(zfrac_lower); + int img_x_upper = floor(xfrac_upper); + int img_y_upper = floor(yfrac_upper); + int img_z_upper = floor(zfrac_upper); + img_x_lower = max(img_x_lower, -nimgs_x); + img_y_lower = max(img_y_lower, -nimgs_y); + img_z_lower = max(img_z_lower, -nimgs_z); + img_x_upper = min(img_x_upper, nimgs_x); + img_y_upper = min(img_y_upper, nimgs_y); + img_z_upper = min(img_z_upper, nimgs_z); + int rem_x_lower = floor(max(xfrac_lower - img_x_lower, 0.f) * mesh_x / tile_size_x); + int rem_y_lower = floor(max(yfrac_lower - img_y_lower, 0.f) * mesh_y / tile_size_y); + int rem_z_lower = floor(max(zfrac_lower - img_z_lower, 0.f) * mesh_z / tile_size_z); + int rem_x_upper = ceil (min(xfrac_upper - img_x_upper, 1.f) * mesh_x / tile_size_x); + int rem_y_upper = ceil (min(yfrac_upper - img_y_upper, 1.f) * mesh_y / tile_size_y); + int rem_z_upper = ceil (min(zfrac_upper - img_z_upper, 1.f) * mesh_z / tile_size_z); + int count_x = rem_x_upper - rem_x_lower + (img_x_upper - img_x_lower) * tiles_x; + int count_y = rem_y_upper - rem_y_lower + (img_y_upper - img_y_lower) * tiles_y; + int count_z = rem_z_upper - rem_z_lower + (img_z_upper - img_z_lower) * tiles_z; + // TODO: tiles in the corners sometimes are out of the cutoff radius. + // They can be discarded and counts can be reduced + int counts = count_x * count_y * count_z; + int n = atomicAdd(head, counts); + int Ny = nimgs_y * 2 + 1; + int Nz = nimgs_z * 2 + 1; + // lattice sum spans over [-nimgs_x, nimgs_x], [-nimgs_y, nimgs_y], [-nimgs_z, nimgs_z], + // Add img_offset to avoid negative indexing + int img_offset = nimgs_x * Ny * Nz + nimgs_y * Nz + nimgs_z; + int64_t Nbas = nbas; + for (int x = rem_x_lower, img_x = img_x_lower; x < rem_x_upper || img_x < img_x_upper;) { + for (int y = rem_y_lower, img_y = img_y_lower; y < rem_y_upper || img_y < img_y_upper;) { + for (int z = rem_z_lower, img_z = img_z_lower; z < rem_z_upper || img_z < img_z_upper;) { + // when (x, y, z) lies out of the unit cell, they can be repositioned + // by shifting the lattice sum index on bra + int64_t latsum_idx = img_offset + (img_x * Ny + img_y) * Nz + img_z; + // dressed_bas_ij stores (latsum_idx*nbas+ish, jL*bvk_nbas+jsh). + // latsum_idx is the image index to reposition bra whereas jL is + // the image index relative to bra. + dressed_bas_ij[n] = latsum_idx * Nbas * NBAS_MAX + bas_ij; + grid_tile_idx[n] = (x * tiles_y + y) * tiles_z + z; + + n++; + z++; + if (z >= tiles_z && img_z < img_z_upper) { + z = 0; + img_z++; + } + } + y++; + if (y >= tiles_y && img_y < img_y_upper) { + y = 0; + img_y++; + } + } + x++; + if (x >= tiles_x && img_x < img_x_upper) { + x = 0; + img_x++; + } + } +} + +// An estimation of the upper bound of the overlap || for +// shell pairs between the primitve cell and the super-mol +__global__ static +void ovlp_mask_estimation_kernel(int8_t *ovlp_mask, PBCIntEnvVars envs, + double *img_coords, int nimgs, float log_cutoff) +{ + int jsh = blockIdx.x * blockDim.x + threadIdx.x; + int ish = blockIdx.y * blockDim.y + threadIdx.y; + int nbas = envs.nbas; + int bvk_nbas = envs.nbas * envs.bvk_ncells; + if (ish >= nbas || jsh >= bvk_nbas) { + return; + } + int jsh_cell0 = jsh % nbas; + if (ish < jsh_cell0) { + return; + } + int *bas = envs.bas; + double *env = envs.env; + int li = bas[ish*BAS_SLOTS+ANG_OF]; + int lj = bas[jsh*BAS_SLOTS+ANG_OF]; + float ai = env[bas[ish*BAS_SLOTS+PTR_EXP]]; + float aj = env[bas[jsh*BAS_SLOTS+PTR_EXP]]; + float ci = env[bas[ish*BAS_SLOTS+PTR_COEFF]]; + float cj = env[bas[jsh*BAS_SLOTS+PTR_COEFF]]; + double *ri = env + bas[ish*BAS_SLOTS+PTR_BAS_COORD]; + double *rj = env + bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; + float xi = ri[0]; + float yi = ri[1]; + float zi = ri[2]; + float xj = rj[0]; + float yj = rj[1]; + float zj = rj[2]; + float aij = ai + aj; + float fi = ai / aij; + float fj = aj / aij; + float theta = ai * fj; + float xjxi = xj - xi; + float yjyi = yj - yi; + float zjzi = zj - zi; + + float log_cicj = logf(fabsf(ci * cj)); + float log_fac = log_cicj + 1.717f - 1.5f * logf(aij) - log_cutoff; + log_fac = max(log_fac, 1e-9f); + float rr_raw = log_fac / theta; + float Ecut_raw = log_fac * (2*aij); + float Ecut_2a = Ecut_raw / (2*aij*aij); + float dri_fac = .5f * logf(.5f*li/aij + fi*fi*rr_raw + Ecut_2a); + float drj_fac = .5f * logf(.5f*lj/aij + fj*fj*rr_raw + Ecut_2a); + // An approximate penalty for the polynomial part of the gaussian product + log_fac += li * dri_fac + lj * drj_fac; + float rr_cutoff = log_fac / theta; + + for (int img = 0; img < nimgs; ++img) { + float xjLxi = xjxi + img_coords[img*3+0]; + float yjLyi = yjyi + img_coords[img*3+1]; + float zjLzi = zjzi + img_coords[img*3+2]; + float rr_ij = xjLxi * xjLxi + yjLyi * yjLyi + zjLzi * zjLzi; + if (rr_ij < rr_cutoff) { + ovlp_mask[ish * bvk_nbas + jsh] = 1; + break; + } + } +} + +__global__ static +void estimate_aft_Ecut_kernel(float *Ecut, int64_t *bas_ij_idx, PBCIntEnvVars envs, + double *img_coords, int nimgs, int npairs, + float log_cutoff, float Ecut_max, int is_mgga) +{ + int pair_id = blockIdx.x * blockDim.x + threadIdx.x; + if (pair_id >= npairs) { + return; + } + int *bas = envs.bas; + double *env = envs.env; + int64_t bas_ij = bas_ij_idx[pair_id]; + int ish = bas_ij / NBAS_MAX; + int jsh = bas_ij % NBAS_MAX; + int li = bas[ish*BAS_SLOTS+ANG_OF]; + int lj = bas[jsh*BAS_SLOTS+ANG_OF]; + float ai = env[bas[ish*BAS_SLOTS+PTR_EXP]]; + float aj = env[bas[jsh*BAS_SLOTS+PTR_EXP]]; + float ci = env[bas[ish*BAS_SLOTS+PTR_COEFF]]; + float cj = env[bas[jsh*BAS_SLOTS+PTR_COEFF]]; + double *ri = env + bas[ish*BAS_SLOTS+PTR_BAS_COORD]; + double *rj = env + bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; + float aij = ai + aj; + float fi = ai / aij; + float fj = aj / aij; + float theta = ai * fj; + if (is_mgga) { + li += 1; + lj += 1; + ci *= ai * 2; + cj *= aj * 2; + } + float log_cicj = logf(fabsf(ci * cj)); + float log_fac = log_cicj + 1.717f - 1.5f * logf(aij) - log_cutoff; + log_fac = max(log_fac, 1e-9f); + float rr_raw = log_fac / theta; + float Ecut_raw = log_fac * (2*aij); + float Ecut_2a = Ecut_raw / (2*aij*aij); + float dri_fac = .5f * logf(.5f*li/aij + fi*fi*rr_raw + Ecut_2a); + float drj_fac = .5f * logf(.5f*lj/aij + fj*fj*rr_raw + Ecut_2a); + // An approximate penalty for the polynomial part of the gaussian product + log_fac += li * dri_fac + lj * drj_fac; + + float xi = ri[0]; + float yi = ri[1]; + float zi = ri[2]; + float xj = rj[0]; + float yj = rj[1]; + float zj = rj[2]; + float xjxi = xj - xi; + float yjyi = yj - yi; + float zjzi = zj - zi; + float Ecut_required = 0.f; + for (int img = 0; img < nimgs; ++img) { + float xjLxi = xjxi + img_coords[img*3+0]; + float yjLyi = yjyi + img_coords[img*3+1]; + float zjLzi = zjzi + img_coords[img*3+2]; + float rr = xjLxi * xjLxi + yjLyi * yjLyi + zjLzi * zjLzi; +// Ecut estimation based on pyscf.pbc.gto.cell.estimate_ke_cutoff +// Factors for Ecut estimation should be +// fac = cs[:,None]*cs * cp.exp(-theta*dr**2) * fac_dri * fac_drj * fl +// where +// fac_dri = (li * .5/aij + dri**2 + Ecut/2/aij**2)**(li*.5) +// ~= (li * .5/aij + dri**2 + log(1./precision)/aij)**(li*.5) +// fac_drj = (lj * .5/aij + drj**2 + Ecut/2/aij**2)**(lj*.5) +// ~= (lj * .5/aij + drj**2 + log(1./precision)/aij)**(lj*.5) +// Here, this fac is approximately derived from the overlap integral +// Ecut ~= log(fac / precision) * 2*aij + float Ecut_estimate = (log_fac - theta*rr) * (2*aij); + Ecut_required = max(Ecut_estimate, Ecut_required); + } + Ecut[pair_id] = min(Ecut_max, Ecut_required); +} + +__global__ static +void supmol_non_trivial_pairs_kernel(int64_t *supmol_bas_ij, int64_t *bas_ij_idx, + PBCIntEnvVars envs, int npairs, float log_cutoff, + int is_mgga, int *head) +{ + int thread_id = threadIdx.x; + int pair_id = blockIdx.x * blockDim.x + thread_id; + if (pair_id >= npairs) { + return; + } + constexpr int batch_size = 64; + __shared__ int8_t img_cache[THREADS*batch_size]; + int bvk_nbas = envs.nbas * envs.bvk_ncells; + int nimgs = envs.nimgs; + int *bas = envs.bas; + double *env = envs.env; + double *img_coords = envs.img_coords; + + int64_t bas_ij = bas_ij_idx[pair_id]; + int ish = bas_ij / NBAS_MAX; + int jsh = bas_ij % NBAS_MAX; + int li = bas[ish*BAS_SLOTS+ANG_OF]; + int lj = bas[jsh*BAS_SLOTS+ANG_OF]; + float ai = env[bas[ish*BAS_SLOTS+PTR_EXP]]; + float aj = env[bas[jsh*BAS_SLOTS+PTR_EXP]]; + float ci = env[bas[ish*BAS_SLOTS+PTR_COEFF]]; + float cj = env[bas[jsh*BAS_SLOTS+PTR_COEFF]]; + double *ri = env + bas[ish*BAS_SLOTS+PTR_BAS_COORD]; + double *rj = env + bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; + float aij = ai + aj; + float fi = ai / aij; + float fj = aj / aij; + float theta = ai * fj; + if (is_mgga) { + li += 1; + lj += 1; + ci *= ai * 2; + cj *= aj * 2; + } + float log_cicj = logf(fabsf(ci * cj)); + float log_fac = log_cicj + 1.717f - 1.5f * logf(aij) - log_cutoff; + log_fac = max(log_fac, 1e-9f); + float rr_raw = log_fac / theta; + float dri_fac = .5f * logf(.5f*li/aij + fi*fi*rr_raw); + float drj_fac = .5f * logf(.5f*lj/aij + fj*fj*rr_raw); + // An approximate penalty for the polynomial part of the gaussian product + log_fac += li * dri_fac + lj * drj_fac; + log_fac = max(log_fac, 0.f); + float rr_cutoff = log_fac / theta; + + float xi = ri[0]; + float yi = ri[1]; + float zi = ri[2]; + float xj = rj[0]; + float yj = rj[1]; + float zj = rj[2]; + float xjxi = xj - xi; + float yjyi = yj - yi; + float zjzi = zj - zi; + for (int img0 = 0; img0 < nimgs; img0 += batch_size) { + int count = 0; + for (int i = 0; i < min(batch_size, nimgs-img0); ++i) { + int img = img0 + i; + float xjLxi = xjxi + img_coords[img*3+0]; + float yjLyi = yjyi + img_coords[img*3+1]; + float zjLzi = zjzi + img_coords[img*3+2]; + float rr = xjLxi * xjLxi + yjLyi * yjLyi + zjLzi * zjLzi; + if (rr < rr_cutoff) { + img_cache[count*THREADS+thread_id] = i; + count++; + } + } + if (count > 0) { + int off = atomicAdd(head, count); + for (int n = 0; n < count; ++n) { + int64_t img = img0 + img_cache[n*THREADS+thread_id]; + // the jsh Id in bas_ij is updated to img*bvk_nbas+jsh + supmol_bas_ij[off+n] = img * bvk_nbas + bas_ij; + } + } + } +} + +extern "C" { +int gaussian_prod_grid_ranges(float2 *grid_frac_ranges, float *pair_ke, + float *Ecut_by_shell, PBCIntEnvVars *envs, + int64_t *bas_ij_idx, int npairs, + int li_inc, int lj_inc, float log_threshold, + float undressed_threshold, float ke_max) +{ + int batches = (npairs + THREADS-1) / THREADS; + grid_ranges_kernel<<>>( + grid_frac_ranges, pair_ke, Ecut_by_shell, *envs, bas_ij_idx, + li_inc, lj_inc, npairs, log_threshold, undressed_threshold, ke_max); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in gaussian_prod_grid_ranges: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} + +int grid_range_to_tiles(int *grid_tile_idx, int64_t *dressed_bas_ij, + int64_t *bas_ij_idx, float2 *grid_frac_ranges, + int *nimgs, int *mesh, int npairs, int nbas, int *head) +{ + cudaMemset(head, 0, sizeof(int)); + int nimgs_x = nimgs[0]; + int nimgs_y = nimgs[1]; + int nimgs_z = nimgs[2]; + int mesh_x = mesh[0]; + int mesh_y = mesh[1]; + int mesh_z = mesh[2]; + int batches = (npairs + THREADS-1) / THREADS; + grid_range_to_tiles_kernel<<>>( + grid_tile_idx, dressed_bas_ij, bas_ij_idx, grid_frac_ranges, + nimgs_x, nimgs_y, nimgs_z, mesh_x, mesh_y, mesh_z, npairs, nbas, head); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in grid_range_to_tiles: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} + +int bvk_ovlp_mask_estimation(int8_t *ovlp_mask, PBCIntEnvVars *envs, + double *img_coords, int nimgs, float log_cutoff) +{ + int nbas = envs->nbas; + int bvk_nbas = nbas * envs->bvk_ncells; + dim3 threads(16, 16); + dim3 blocks((bvk_nbas + 15) / 16, (nbas + 15) / 16); + ovlp_mask_estimation_kernel<<>>( + ovlp_mask, *envs, img_coords, nimgs, log_cutoff); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in bvk_ovlp_mask_estimation: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} +int supmol_non_trivial_pairs(int64_t *supmol_bas_ij, int64_t *bas_ij_idx, + PBCIntEnvVars *envs, int npairs, float log_cutoff, + int is_mgga, int *head) +{ + cudaMemset(head, 0, sizeof(int)); + int blocks = (npairs + THREADS-1)/THREADS; + supmol_non_trivial_pairs_kernel<<>>( + supmol_bas_ij, bas_ij_idx, *envs, npairs, log_cutoff, is_mgga, head); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in bvk_ovlp_mask_estimation: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} + +int estimate_aft_Ecut(float *Ecut, int64_t *bas_ij_idx, PBCIntEnvVars *envs, + double *img_coords, int nimgs, int npairs, + float log_cutoff, float Ecut_max, int is_mgga) +{ + int blocks = (npairs + THREADS-1)/THREADS; + estimate_aft_Ecut_kernel<<>>( + Ecut, bas_ij_idx, *envs, img_coords, nimgs, npairs, log_cutoff, + Ecut_max, is_mgga); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in raw_ovlp_mask: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} +} diff --git a/gpu4pyscf/lib/multigrid/multigrid_v3/utils.cuh b/gpu4pyscf/lib/multigrid/multigrid_v3/utils.cuh new file mode 100644 index 000000000..808904ea1 --- /dev/null +++ b/gpu4pyscf/lib/multigrid/multigrid_v3/utils.cuh @@ -0,0 +1,53 @@ +/* + * Copyright 2026 The PySCF Developers. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +template +__host__ __device__ T distance_squared(const T x, const T y, const T z) { + return x * x + y * y + z * z; +} + +__device__ __forceinline__ +void multiply(double aR, double aI, double bR, double bI, double &cR, double &cI) +{ + double outR = aR * bR - aI * bI; + double outI = aR * bI + aI * bR; + cR = outR; + cI = outI; +} + +__device__ __forceinline__ +double reduce(double val, double *swap, int thread_id) +{ + constexpr int WARP_SIZE = 32; + for (int offset = 16; offset > 0; offset >>= 1) { + val += __shfl_down_sync(0xffffffff, val, offset); + } + int lane = thread_id % WARP_SIZE; + int warp = thread_id / WARP_SIZE; + if (lane == 0) { + swap[warp] = val; + } + __syncthreads(); + if (warp == 0) { + val = (thread_id < 8) ? swap[lane] : 0.; + for (int offset = 4; offset > 0; offset >>= 1) { + val += __shfl_down_sync(0xffffffff, val, offset); + } + } + return val; +} diff --git a/gpu4pyscf/lib/pbc/overlap.cu b/gpu4pyscf/lib/pbc/overlap.cu index 52a05d85d..296bba364 100644 --- a/gpu4pyscf/lib/pbc/overlap.cu +++ b/gpu4pyscf/lib/pbc/overlap.cu @@ -2198,22 +2198,24 @@ void ovlp_mask_estimation_kernel(int8_t *ovlp_mask, float *exps, float *log_coef float xjxi = xj - xi; float yjyi = yj - yi; float zjzi = zj - zi; - float fac_norm = log_coeff[ish] + log_coeff[jsh] + 1.717f - 1.5f * logf(aij); + + float log_cicj = log_coeff[ish] + log_coeff[jsh]; + float log_fac = log_cicj + 1.717f - 1.5f * logf(aij) - log_cutoff; + log_fac = max(log_fac, 1e-9f); + float rr_raw = log_fac / theta; + float dri_fac = .5f * logf(.5f*li/aij + fi*fi*rr_raw); + float drj_fac = .5f * logf(.5f*lj/aij + fj*fj*rr_raw); + // An approximate penalty for the polynomial part of the gaussian product + log_fac += li * dri_fac + lj * drj_fac; + log_fac = max(log_fac, 0.f); + float rr_cutoff = log_fac / theta; + for (int img = 0; img < nimgs; ++img) { float xjLxi = xjxi + img_coords[img*3+0]; float yjLyi = yjyi + img_coords[img*3+1]; float zjLzi = zjzi + img_coords[img*3+2]; float rr_ij = xjLxi * xjLxi + yjLyi * yjLyi + zjLzi * zjLzi; - if (theta*rr_ij > REMOTE_THRESHOLD) { - continue; - } - float dr = sqrtf(rr_ij); - float dri = fj * dr; - float drj = fi * dr; - float dri_fac = .5f*li * logf(.5f*li/aij + dri*dri + 1e-9f); - float drj_fac = .5f*lj * logf(.5f*lj/aij + drj*drj + 1e-9f); - float log_ovlp = fac_norm - theta*rr_ij + dri_fac + drj_fac; - if (log_ovlp > log_cutoff) { + if (rr_ij < rr_cutoff) { ovlp_mask[img*nbas+bas_ij] = 1; if (hermi) { ovlp_mask[img*nbas+bas_ji] = 1; diff --git a/gpu4pyscf/pbc/df/aft.py b/gpu4pyscf/pbc/df/aft.py index aea94bb3f..0b1b27901 100644 --- a/gpu4pyscf/pbc/df/aft.py +++ b/gpu4pyscf/pbc/df/aft.py @@ -38,7 +38,7 @@ from gpu4pyscf.gto.mole import SortedGTO from gpu4pyscf.lib import logger, utils from gpu4pyscf.lib.cupy_helper import (return_cupy_array, contract, unpack_tril, - get_avail_mem, asarray) + get_avail_mem, asarray, ndarray) KE_SCALING = aft_cpu.KE_SCALING @@ -404,7 +404,7 @@ def _fake_nuc(cell, with_pseudo=True): fakenuc.rcut = 0.1 return fakenuc -def _get_ZSI(cell, mesh=None): +def _get_ZSI(cell, mesh=None, out=None): ''' Calculate the product of nuclear charges and structure factor @@ -422,5 +422,7 @@ def _get_ZSI(cell, mesh=None): SIy = cp.exp(-1j*rb[:,1,None] * basey) SIz = cp.exp(-1j*rb[:,2,None] * basez) SIx *= Z[:,None] - ZG = cp.einsum('qx,qy,qz->xyz', SIx, SIy, SIz).ravel() + rho_xy = SIx[:,:,None] * SIy[:,None,:] + out = ndarray(mesh, dtype=np.complex128, buffer=out) + ZG = contract('qxy,qz->xyz', rho_xy, SIz, out=out).ravel() return ZG diff --git a/gpu4pyscf/pbc/df/fft_jk.py b/gpu4pyscf/pbc/df/fft_jk.py index 461021763..4f568f85e 100644 --- a/gpu4pyscf/pbc/df/fft_jk.py +++ b/gpu4pyscf/pbc/df/fft_jk.py @@ -402,31 +402,22 @@ def _format_jks(v_kpts, dm_kpts, kpts_band, kpts): if kpts_band is kpts or kpts_band is None: return v_kpts.reshape(dm_kpts.shape) else: - assert v_kpts.ndim == 4 # (Ndm,Nk,Nao,Nao) - # dm_kpts.shape kpts.shape nset - # (Nao,Nao) (1 ,3) None - # (Ndm,Nao,Nao) (1 ,3) Ndm - # (Nk,Nao,Nao) (Nk,3) None - # (Ndm,Nk,Nao,Nao) (Nk,3) Ndm - if kpts_band.ndim == 1: - assert dm_kpts.ndim <= 3 - v_kpts = v_kpts[:,0] - if dm_kpts.ndim < 3: # RHF dm - v_kpts = v_kpts[0] + nkpts_band = len(kpts_band) + nao = dm_kpts.shape[-1] + if dm_kpts.ndim == v_kpts.ndim: + assert v_kpts.shape[-3:] == (nkpts_band, nao, nao) + return v_kpts + + is_single_kpt = kpts is not None and kpts.ndim == 1 + if is_single_kpt: + is_rhf = dm_kpts.ndim == 2 else: - if kpts is None or kpts.ndim == 1 or (kpts.ndim == 2 and kpts.shape[0] == 1): - nkpts = 1 - if dm_kpts.ndim == 2: - dm_kpts = dm_kpts[None,:,:] - else: - nkpts = len(kpts) - assert kpts.ndim == 2 assert dm_kpts.ndim >= 3 - if dm_kpts.ndim == 3: # KRHF dms - assert len(dm_kpts) == nkpts - v_kpts = v_kpts[0] - else: # KUHF dms - assert v_kpts.shape[1] == len(kpts_band) + is_rhf = dm_kpts.ndim == 3 + if is_rhf: + v_kpts = v_kpts.reshape(nkpts_band, nao, nao) + else: # KUHF dms + v_kpts = v_kpts.reshape(2, nkpts_band, nao, nao) return v_kpts def _factorize_dm(dm_kpts, kpts=None): diff --git a/gpu4pyscf/pbc/df/ft_ao.py b/gpu4pyscf/pbc/df/ft_ao.py index 399d82aa4..6c23e5a97 100644 --- a/gpu4pyscf/pbc/df/ft_ao.py +++ b/gpu4pyscf/pbc/df/ft_ao.py @@ -396,13 +396,13 @@ def evaluate_ft(Gv, batch_id=0, out=None): nao_pair = ao_pair_offsets[batch_id+1] - ao_pair_offset out = ndarray((nao_pair, nGv), dtype=np.complex128, buffer=out) if not cart: - out[:] = 0. + out.fill(0) else: pair_blocks = len(shl_pair_offsets) - 1 _shl_pair_offsets = shl_pair_offsets ao_pair_offset = 0 out = ndarray((nao, bvk_ncells, nao, nGv), dtype=np.complex128, buffer=out) - out[:] = 0. + out.fill(0) err = kern( ctypes.cast(out.data.ptr, ctypes.c_void_p), ctypes.byref(aft_envs), diff --git a/gpu4pyscf/pbc/df/tests/test_pbc_ft_ao.py b/gpu4pyscf/pbc/df/tests/test_pbc_ft_ao.py index 8cbe6fd72..cc6e0a8d4 100644 --- a/gpu4pyscf/pbc/df/tests/test_pbc_ft_ao.py +++ b/gpu4pyscf/pbc/df/tests/test_pbc_ft_ao.py @@ -25,6 +25,8 @@ from gpu4pyscf.pbc.df.ft_ao import ft_aopair, ft_aopair_kpts from gpu4pyscf.pbc.lib.kpts_helper import conj_images_in_bvk_cell from gpu4pyscf.pbc.df.ft_ao import libpbc +from gpu4pyscf.pbc.df.aft import AFTDF +from gpu4pyscf.pbc.tools import get_coulG def setUpModule(): global cell @@ -177,6 +179,42 @@ def test_contract_dm(self): ref = cp.einsum('Gpq,qp->G', ft_aopair(cell, Gv), dm).get() self.assertAlmostEqual(abs(ref-rhoG).max(), 0, 10) + def test_j(self): + cp.random.seed(10) + nao = cell.nao + dm = cp.random.rand(nao, nao) + dm = dm + dm.T + mesh = [9,11,3] + Gv = cell.get_Gv(mesh=mesh) + ft_opt = ft_ao_gpu.FTOpt(cell) + rhoG = ft_opt.contract_dm(ft_opt.cell.apply_C_mat_CT(dm), Gv) + coulG = rhoG * cp.array(get_coulG(cell, Gv=Gv)) / cell.vol + vj = ft_opt.contract_rhoG(coulG, -Gv, sort_output=True) + + mydf = AFTDF(cell).set(mesh=mesh) + jref = mydf.get_jk(dm, with_k=False)[0] + assert abs(vj - jref).max().get() < 1e-9 + + @unittest.skip('FTOpt.contract_dm with kpts not supported') + def test_j_kpts(self): + cp.random.seed(10) + kmesh = [3,2,1] + kpts = cell.make_kpts(kmesh) + nk = len(kpts) + nao = cell.nao + dm = cp.random.rand(nk, nao, nao) + dm = dm + dm.transpose(0, 2, 1) + mesh = [9,11,3] + Gv = cell.get_Gv(mesh=mesh) + ft_opt = ft_ao_gpu.FTOpt(cell) + rhoG = ft_opt.contract_dm(ft_opt.cell.apply_C_mat_CT(dm), Gv, kpts=kpts) + coulG = rhoG * cp.array(get_coulG(cell, Gv=Gv)) / cell.vol + vj = ft_opt.contract_rhoG(coulG, -Gv, kpts=kpts, sort_output=True) + + mydf = AFTDF(cell).set(mesh=mesh) + jref = mydf.get_jk(dm, kpts=kpts, with_k=False)[0] + assert abs(vj - jref).max().get() < 1e-9 + def test_gen_ft_kernel(self): kmesh = [2, 3, 2] kpts = cell.make_kpts(kmesh) diff --git a/gpu4pyscf/pbc/dft/krks.py b/gpu4pyscf/pbc/dft/krks.py index e10b262d2..af4ffbe4e 100644 --- a/gpu4pyscf/pbc/dft/krks.py +++ b/gpu4pyscf/pbc/dft/krks.py @@ -30,7 +30,7 @@ from gpu4pyscf.pbc.scf.rsjk import PBCJKMatrixOpt from gpu4pyscf.pbc.scf.j_engine import PBCJMatrixOpt from gpu4pyscf.pbc.dft import rks -from gpu4pyscf.pbc.dft import multigrid, multigrid_v2 +from gpu4pyscf.pbc.dft import multigrid def get_veff(ks, cell=None, dm=None, dm_last=None, vhf_last=None, hermi=1, kpts=None, kpts_band=None): @@ -47,7 +47,7 @@ def get_veff(ks, cell=None, dm=None, dm_last=None, vhf_last=None, hermi=1, nkpts = len(kpts) weight = 1. / nkpts - if isinstance(ni, (multigrid_v2.MultiGridNumInt, multigrid.MultiGridNumInt)): + if isinstance(ni, multigrid.MultiGridNumIntBase): if ks.do_nlc(): raise NotImplementedError(f'MultiGrid for NLC functional {ks.xc} + {ks.nlc}') n, exc, vxc = ni.nr_rks( @@ -240,8 +240,7 @@ def gen_response(self, mo_coeff=None, mo_occ=None, dm0 = None with_j = (singlet is None or singlet) and hermi != 2 - j_in_xc = isinstance(ni, (multigrid_v2.MultiGridNumInt, - multigrid.MultiGridNumInt)) + j_in_xc = isinstance(ni, multigrid.MultiGridNumIntBase) def vind(dm1, kshift=0): assert kshift == 0 diff --git a/gpu4pyscf/pbc/dft/kuks.py b/gpu4pyscf/pbc/dft/kuks.py index c7c71929e..528c29afc 100644 --- a/gpu4pyscf/pbc/dft/kuks.py +++ b/gpu4pyscf/pbc/dft/kuks.py @@ -28,7 +28,7 @@ from gpu4pyscf.lib.cupy_helper import tag_array, get_avail_mem from gpu4pyscf.pbc.scf import khf, kuhf from gpu4pyscf.pbc.dft import rks, krks -from gpu4pyscf.pbc.dft import multigrid, multigrid_v2 +from gpu4pyscf.pbc.dft import multigrid def get_veff(ks, cell=None, dm=None, dm_last=None, vhf_last=None, hermi=1, kpts=None, kpts_band=None): @@ -47,7 +47,7 @@ def get_veff(ks, cell=None, dm=None, dm_last=None, vhf_last=None, hermi=1, nkpts = len(kpts) weight = 1. / nkpts - if isinstance(ni, (multigrid_v2.MultiGridNumInt, multigrid.MultiGridNumInt)): + if isinstance(ni, multigrid.MultiGridNumIntBase): if ks.do_nlc(): raise NotImplementedError(f'MultiGrid for NLC functional {ks.xc} + {ks.nlc}') n, exc, vxc = ni.nr_uks( @@ -159,8 +159,7 @@ def gen_response(self, mo_coeff=None, mo_occ=None, nkpts = len(kpts) with_j = with_j and hermi != 2 - j_in_xc = isinstance(ni, (multigrid_v2.MultiGridNumInt, - multigrid.MultiGridNumInt)) + j_in_xc = isinstance(ni, multigrid.MultiGridNumIntBase) def vind(dm1, kshift=0): assert kshift == 0 diff --git a/gpu4pyscf/pbc/dft/multigrid.py b/gpu4pyscf/pbc/dft/multigrid.py index 16b46fa95..0507f4cea 100644 --- a/gpu4pyscf/pbc/dft/multigrid.py +++ b/gpu4pyscf/pbc/dft/multigrid.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import warnings import itertools import ctypes from dataclasses import dataclass @@ -26,7 +27,7 @@ from gpu4pyscf.lib import utils from gpu4pyscf.lib.cupy_helper import ( load_library, tag_array, contract, sandwich_dot, block_diag, transpose_sum, - dist_matrix, batched_vec_norm2) + dist_matrix, ndarray) from gpu4pyscf.gto.mole import cart2sph_by_l from gpu4pyscf.dft import numint from gpu4pyscf.pbc import tools @@ -38,13 +39,16 @@ __all__ = ['MultiGridNumInt'] -libmgrid = load_library('libmgrid') -libmgrid.MG_eval_rho_orth.restype = ctypes.c_int -libmgrid.MG_eval_mat_lda_orth.restype = ctypes.c_int -libmgrid.MG_eval_mat_gga_orth.restype = ctypes.c_int -libmgrid.MG_init_constant.restype = ctypes.c_int -libmgrid.ovlp_mask_estimation.restype = ctypes.c_int -libmgrid.filter_supmol_bas.restype = ctypes.c_int +try: + libmgrid = load_library('libmgrid') + libmgrid.MG_eval_rho_orth.restype = ctypes.c_int + libmgrid.MG_eval_mat_lda_orth.restype = ctypes.c_int + libmgrid.MG_eval_mat_gga_orth.restype = ctypes.c_int + libmgrid.MG_init_constant.restype = ctypes.c_int + libmgrid.ovlp_mask_estimation.restype = ctypes.c_int + libmgrid.filter_supmol_bas.restype = ctypes.c_int +except OSError: + libmgrid = None PRIMBAS_ANG = 0 PRIMBAS_EXP = 1 @@ -695,36 +699,6 @@ def eval_nucG(cell, mesh): nucG *= tools.get_coulG(cell, Gv=Gv) return nucG -def eval_nucG_SI_gradient(cell, mesh, rho_g): - ngrids = np.prod(mesh) - assert rho_g.shape == (ngrids,) - - assert cell.dimension == 3 - Gv, (basex, basey, basez) = get_Gv_weights(cell, mesh)[:2] - b = cell.reciprocal_vectors() - coords = cell.atom_coords() - rb = cp.asarray(coords.dot(b.T)) - SIx = cp.exp(-1j*rb[:,0,None] * basex) - SIy = cp.exp(-1j*rb[:,1,None] * basey) - SIz = cp.exp(-1j*rb[:,2,None] * basez) - dSI_prefactor = -1j * Gv.T * rho_g.conj() - charges = -cell.atom_charges() - coulG = tools.get_coulG(cell, Gv=Gv) - - de = cp.empty([cell.natm, 3], dtype = cp.complex128) - - for i_atom in range(cell.natm): - SI = (SIx[i_atom,:,None,None] * SIy[i_atom,:,None] * SIz[i_atom]).ravel() - de[i_atom, :] = charges[i_atom] * (dSI_prefactor @ (coulG * SI)) - - grad_max_imag = cp.max(cp.abs(de.imag)) - if grad_max_imag >= 1e-8: - logger.warn(cell, f"Large imaginary part ({grad_max_imag:e}) from nuclear repulsion term structure factor gradient") - - de = de.real - de /= cell.vol - return de - def get_nuc(ni, kpts=None): assert kpts is None or is_zero(kpts) if kpts is None or kpts.ndim == 1: @@ -741,84 +715,85 @@ def get_nuc(ni, kpts=None): vne = vne[0] return vne +def _get_Gv_bases(mesh, b): + Gx = cp.array(np.fft.fftfreq(mesh[0], 1./mesh[0]) * b[0,:,None]) + Gy = cp.array(np.fft.fftfreq(mesh[1], 1./mesh[1]) * b[1,:,None]) + Gz = cp.array(np.fft.fftfreq(mesh[2], 1./mesh[2]) * b[2,:,None]) + return (Gx, Gy, Gz) + _kernel_registery = {} def _append_vpplocG_one_atom_without_gamma(i_atom, natm, rloc, nexp, cexp, charge, - mesh, G2, coulG, SIx, SIy, SIz, vlocG): + mesh, Gv_bases, SIx, SIy, SIz, vlocG): # Result will be appended to vlocG - fn_name = f"gth_loc_reciporcal_nexp_{nexp}_kernel" + fn_name = "gth_loc_reciporcal_kernel" if fn_name not in _kernel_registery: - C_declaration = '' - C_contribution = '' - if nexp >= 1: - C_declaration += ', const double cexp0' - C_contribution += 'cfacs += cexp0;' - if nexp >= 2: - C_declaration += ', const double cexp1' - C_contribution += 'cfacs += cexp1 * (3 - G2_red);' - if nexp >= 3: - C_declaration += ', const double cexp2' - C_contribution += 'cfacs += cexp2 * (15 - 10 * G2_red + G2_red * G2_red);' - if nexp >= 4: - C_declaration += ', const double cexp3' - C_contribution += 'cfacs += cexp3 * (105 - 105 * G2_red + 21 * G2_red * G2_red - G2_red * G2_red * G2_red);' kernel_code = r''' #include extern "C" __global__ void ''' + fn_name + '''( - const double* __restrict__ grids_G2, const double* __restrict__ grids_coulG, + const double* __restrict__ Gx, const double* __restrict__ Gy, const double* __restrict__ Gz, const complex* __restrict__ grids_SIx, const complex* __restrict__ grids_SIy, const complex* __restrict__ grids_SIz, complex* __restrict__ grids_vlocG, - const int n_mesh_x, const int n_mesh_y, const int n_mesh_z, const int i_atom, - const double charge, const double rloc''' + C_declaration + r''') + const int nx, const int ny, const int nz, const int i_atom, + const double charge, const double rloc, + double cexp0, double cexp1, double cexp2, double cexp3, int nexp) { - const int i_grid = blockDim.x * blockIdx.x + threadIdx.x; - const int ngrids = n_mesh_x * n_mesh_y * n_mesh_z; - if (i_grid >= ngrids) return; - - const double G2 = grids_G2[i_grid]; - const double coulG = grids_coulG[i_grid]; - const double G2_red = G2 * rloc * rloc; - const int i_grid_x = i_grid / (n_mesh_y * n_mesh_z); - const int i_grid_y = (i_grid - i_grid_x * (n_mesh_y * n_mesh_z)) / n_mesh_z; - const int i_grid_z = i_grid - i_grid_x * (n_mesh_y * n_mesh_z) - i_grid_y * n_mesh_z; - const complex SIx = grids_SIx[i_atom * n_mesh_x + i_grid_x]; - const complex SIy = grids_SIy[i_atom * n_mesh_y + i_grid_y]; - const complex SIz = grids_SIz[i_atom * n_mesh_z + i_grid_z]; - const complex SI = SIx * SIy * SIz * exp(-0.5 * G2_red); - complex vlocG = -charge * coulG * SI; - - double cfacs = 0; - ''' + C_contribution + r''' - vlocG += 15.749609945722419 * rloc * rloc * rloc * cfacs * SI; - - grids_vlocG[i_grid] += vlocG; + int idx = blockDim.x * blockIdx.x + threadIdx.x; + int stride = gridDim.x * blockDim.x; + size_t nyz = ny * nz; + size_t ng = nx * nyz; + for (size_t i_grid = idx; i_grid < ng; i_grid += stride) { + int ix = i_grid / nyz; + int iyz = i_grid - nyz * ix; + int iy = iyz / nz; + int iz = iyz - nz * iy; + double G2 = 0.; + for (int n = 0; n < 3; ++n) { + double Gv = Gx[nx*n+ix] + Gy[ny*n+iy] + Gz[nz*n+iz]; + G2 += Gv * Gv; + } + double coulG = 0.; + if (G2 != 0) coulG = 12.566370614359172 / G2; + const double G2_red = G2 * rloc * rloc; + const complex SIx = grids_SIx[i_atom * nx + ix]; + const complex SIy = grids_SIy[i_atom * ny + iy]; + const complex SIz = grids_SIz[i_atom * nz + iz]; + const complex SI = SIx * SIy * SIz * exp(-0.5 * G2_red); + complex vlocG = -charge * coulG * SI; + + double cfacs = 0; + if (nexp >= 1) cfacs += cexp0; + if (nexp >= 2) cfacs += cexp1 * (3 - G2_red); + if (nexp >= 3) cfacs += cexp2 * (15 + G2_red * (G2_red - 10)); + if (nexp >= 4) cfacs += cexp3 * (105 + G2_red * (G2_red * (21 - G2_red) - 105)); + vlocG += 15.749609945722419 * rloc * rloc * rloc * cfacs * SI; + + grids_vlocG[i_grid] += vlocG; + } } ''' _kernel_registery[fn_name] = cp.RawKernel(kernel_code, fn_name) kernel = _kernel_registery[fn_name] - ngrids = G2.shape[0] - assert G2.shape == (ngrids,) and G2.dtype == cp.float64 - assert coulG.shape == (ngrids,) and coulG.dtype == cp.float64 + ngrids = np.prod(mesh) assert SIx.shape == (natm, mesh[0]) and SIx.dtype == cp.complex128 and SIx.flags.c_contiguous assert SIy.shape == (natm, mesh[1]) and SIy.dtype == cp.complex128 and SIy.flags.c_contiguous assert SIz.shape == (natm, mesh[2]) and SIz.dtype == cp.complex128 and SIz.flags.c_contiguous assert vlocG.shape == (ngrids,) and vlocG.dtype == cp.complex128 assert ngrids < np.iinfo(np.int32).max - kernel_parameters = [G2, coulG, SIx, SIy, SIz, vlocG, cp.int32(mesh[0]), cp.int32(mesh[1]), cp.int32(mesh[2]), + kernel_parameters = [Gv_bases[0], Gv_bases[1], Gv_bases[2], + SIx, SIy, SIz, vlocG, cp.int32(mesh[0]), cp.int32(mesh[1]), cp.int32(mesh[2]), cp.int32(i_atom), cp.float64(charge), cp.float64(rloc)] - if nexp >= 1: - kernel_parameters.append(cp.float64(cexp[0])) - if nexp >= 2: - kernel_parameters.append(cp.float64(cexp[1])) - if nexp >= 3: - kernel_parameters.append(cp.float64(cexp[2])) - if nexp >= 4: - kernel_parameters.append(cp.float64(cexp[3])) - kernel(((ngrids + 1024 - 1) // 1024, ), (1024, ), kernel_parameters) + + cexp = [cp.float64(x) for x in cexp] + [cp.float64(0.)] * 4 + kernel_parameters.extend(cexp[:4]) + kernel_parameters.append(cp.int32(nexp)) + + workers = gpu_specs['multiProcessorCount'] + kernel((workers, ), (1024, ), kernel_parameters) # SI = (SIx[i_atom,:,None,None] * SIy[i_atom,:,None] * SIz[i_atom]).ravel() # G2_red = G2 * rloc**2 @@ -839,23 +814,21 @@ def _append_vpplocG_one_atom_without_gamma(i_atom, natm, rloc, nexp, cexp, charg return vlocG -def eval_vpplocG(cell, mesh): +def eval_vpplocG(cell, mesh, out=None): '''PRB, 58, 3641 Eq (5) ''' assert cell.dimension == 3 - Gv, (basex, basey, basez) = get_Gv_weights(cell, mesh)[:2] - b = cell.reciprocal_vectors() - coords = cell.atom_coords() - rb = cp.asarray(coords.dot(b.T)) - SIx = cp.exp(-1j*rb[:,0,None] * basex) - SIy = cp.exp(-1j*rb[:,1,None] * basey) - SIz = cp.exp(-1j*rb[:,2,None] * basez) - # G2 = contract('px,px->p', Gv, Gv) - G2 = batched_vec_norm2(Gv) + Gv_bases = _get_Gv_bases(mesh, cell.reciprocal_vectors()) + coords = cp.asarray(cell.atom_coords()) + SIx = cp.exp(-1j * coords.dot(Gv_bases[0])) + SIy = cp.exp(-1j * coords.dot(Gv_bases[1])) + SIz = cp.exp(-1j * coords.dot(Gv_bases[2])) + charges = cell.atom_charges() - coulG = tools.get_coulG(cell, Gv=Gv) - vlocG = cp.zeros(len(G2), dtype=np.complex128) + ngrids = np.prod(mesh) + vlocG = ndarray(ngrids, dtype=np.complex128, buffer=out) + vlocG.fill(0) vlocG0 = 0 for ia in range(cell.natm): @@ -868,55 +841,12 @@ def eval_vpplocG(cell, mesh): vlocG0 += 2*np.pi*charges[ia]*rloc**2 - _append_vpplocG_one_atom_without_gamma(ia, cell.natm, rloc, nexp, cexp, charges[ia], mesh, G2, coulG, SIx, SIy, SIz, vlocG) + _append_vpplocG_one_atom_without_gamma( + ia, cell.natm, rloc, nexp, cexp, charges[ia], mesh, Gv_bases, SIx, SIy, SIz, vlocG) vlocG[0] += vlocG0 return vlocG -def eval_vpplocG_SI_gradient(cell, mesh, rho_g): - ngrids = np.prod(mesh) - assert rho_g.shape == (ngrids,) - - Gv, (basex, basey, basez) = get_Gv_weights(cell, mesh)[:2] - b = cell.reciprocal_vectors() - coords = cell.atom_coords() - rb = cp.asarray(coords.dot(b.T)) - SIx = cp.exp(-1j*rb[:,0,None] * basex) - SIy = cp.exp(-1j*rb[:,1,None] * basey) - SIz = cp.exp(-1j*rb[:,2,None] * basez) - dSI_prefactor = -1j * Gv.T * rho_g.conj() - G2 = batched_vec_norm2(Gv) - charges = cell.atom_charges() - - coulG = tools.get_coulG(cell, Gv=Gv) - vlocG = cp.zeros(len(G2), dtype=np.complex128) - - de = cp.zeros([cell.natm, 3], dtype = cp.complex128) - - for ia in range(cell.natm): - symb = cell.atom_symbol(ia) - if symb not in cell._pseudo: - continue - - pp = cell._pseudo[symb] - rloc, nexp, cexp = pp[1:3+1] - - vlocG.fill(0) - _append_vpplocG_one_atom_without_gamma(ia, cell.natm, rloc, nexp, cexp, charges[ia], mesh, G2, coulG, SIx, SIy, SIz, vlocG) - - vlocG0 = 2*np.pi*charges[ia]*rloc**2 - vlocG[0] += vlocG0 - - de[ia, :] = dSI_prefactor @ vlocG - - grad_max_imag = cp.max(cp.abs(de.imag)) - if grad_max_imag >= 1e-8: - logger.warn(cell, f"Large imaginary part ({grad_max_imag:e}) from pseudopotential local term structure factor gradient") - - de = de.real - de /= cell.vol - return de - def get_pp(ni, kpts=None): '''Get the periodic pseudopotential nuc-el AO matrix, with G=0 removed. ''' @@ -1398,9 +1328,14 @@ class MGridEnvVars(ctypes.Structure): ('lattice_params', ctypes.c_void_p), ] +class MultiGridNumIntBase(lib.StreamObject, numint.LibXCMixin): + pass -class MultiGridNumInt(lib.StreamObject, numint.LibXCMixin): +class MultiGridNumInt(MultiGridNumIntBase): def __init__(self, cell): + warnings.warn( + 'multigrid is deprecated. Use multigrid_v3 instead', + DeprecationWarning, stacklevel=2) self.mesh = cell.mesh self.reset(cell) @@ -1479,6 +1414,18 @@ def get_vxc(self, cell, grids, xc_code, dm_kpts, spin=0, hermi=1, cache_xc_kernel = NotImplemented cache_xc_kernel1 = NotImplemented + def energy_nuclear_gradient(self, xc_code, dm_kpts, kpts=None, spin=None, + with_j=False, with_nuc=False): + raise NotImplementedError('Please use multigrid_v3.MultiGridNumInt instead.') + + def energy_strain_gradient(self, xc_code, dm_kpts, kpts=None, spin=None, + with_j=False, with_nuc=False): + raise NotImplementedError + + def energy_derivatives(self, xc_code, dm_kpts, kpts=None, spin=None, + with_j=False, with_nuc=False): + raise NotImplementedError('Please use multigrid_v3.MultiGridNumInt instead.') + to_gpu = utils.to_gpu device = utils.device diff --git a/gpu4pyscf/pbc/dft/multigrid_v2.py b/gpu4pyscf/pbc/dft/multigrid_v2.py index dc6c1420c..4c4df6f18 100644 --- a/gpu4pyscf/pbc/dft/multigrid_v2.py +++ b/gpu4pyscf/pbc/dft/multigrid_v2.py @@ -43,13 +43,16 @@ __all__ = ['MultiGridNumInt'] -libgpbc = load_library("libmgrid_v2") -libgpbc.evaluate_density_driver.restype = ctypes.c_int -libgpbc.evaluate_xc_driver.restype = ctypes.c_int -libgpbc.evaluate_xc_gradient_driver.restype = ctypes.c_int -libgpbc.count_non_trivial_pairs.restype = ctypes.c_int -libgpbc.screen_gaussian_pairs.restype = ctypes.c_int -libgpbc.count_pairs_on_blocks.restype = ctypes.c_int +try: + libgpbc = load_library("libmgrid_v2") + libgpbc.evaluate_density_driver.restype = ctypes.c_int + libgpbc.evaluate_xc_driver.restype = ctypes.c_int + libgpbc.evaluate_xc_gradient_driver.restype = ctypes.c_int + libgpbc.count_non_trivial_pairs.restype = ctypes.c_int + libgpbc.screen_gaussian_pairs.restype = ctypes.c_int + libgpbc.count_pairs_on_blocks.restype = ctypes.c_int +except OSError: + libgpbc = None def complex_type(dtype): @@ -1468,10 +1471,7 @@ def get_j_kpts(ni, dm_kpts, hermi=1, kpts=None, kpts_band=None): log = logger.new_logger(cell) t0 = log.init_timer() dm_kpts = cp.asarray(dm_kpts, order="C") - dms = _format_dms(dm_kpts, kpts) - nset = dms.shape[0] mesh = ni.mesh - ngrids = np.prod(mesh) density = evaluate_density_on_g_mesh(ni, dm_kpts, kpts) Gv = get_Gv(cell, mesh) @@ -1480,13 +1480,6 @@ def get_j_kpts(ni, dm_kpts, hermi=1, kpts=None, kpts_band=None): coulomb_on_g_mesh = cp.einsum( "ng, g -> ng", density[:, 0], coulomb_kernel_on_g_mesh ) - weight = cell.vol / ngrids - - density = density.reshape(-1, *mesh) - # *(1./weight) because rhoR is scaled by weight in _eval_rhoG. When - # computing rhoR with IFFT, the weight factor is not needed. - density = ifft_in_place(density).real.reshape(nset, -1, ngrids) - density /= weight #if kpts_band is not None: # ni = ni.copy().reset().build() @@ -1524,6 +1517,7 @@ def nr_rks(ni, cell, grids, xc_code, dm_kpts, relativity=0, hermi=1, if ni.sorted_gaussian_pairs is None: ni.build(xc_type) + input_kpts = kpts if kpts is None: kpts = np.zeros((1, 3)) else: @@ -1586,9 +1580,9 @@ def nr_rks(ni, cell, grids, xc_code, dm_kpts, relativity=0, hermi=1, if with_j: xc_for_fock[0] += coulomb_on_g_mesh - kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band + kpts_band, input_band = _format_kpts_band(kpts_band, input_kpts), kpts_band veff = convert_xc_on_g_mesh_to_fock(ni, xc_for_fock, hermi, kpts_band, with_tau = (xc_type == "MGGA")) - veff = _format_jks(veff, dm_kpts, input_band, kpts) + veff = _format_jks(veff, dm_kpts, input_band, input_kpts) veff = tag_array(veff, ecoul=coulomb_energy, exc=xc_energy_sum) t0 = log.timer("xc", *t0) return n_electrons, xc_energy_sum, veff @@ -1624,6 +1618,7 @@ def nr_uks(ni, cell, grids, xc_code, dm_kpts, relativity=0, hermi=1, if ni.sorted_gaussian_pairs is None: ni.build(xc_type) + input_kpts = kpts if kpts is None: kpts = np.zeros((1, 3)) else: @@ -1690,9 +1685,9 @@ def nr_uks(ni, cell, grids, xc_code, dm_kpts, relativity=0, hermi=1, if with_j: xc_for_fock[:, 0] += coulomb_on_g_mesh - kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band + kpts_band, input_band = _format_kpts_band(kpts_band, input_kpts), kpts_band veff = convert_xc_on_g_mesh_to_fock(ni, xc_for_fock, hermi, kpts_band, with_tau = (xc_type == "MGGA")) - veff = _format_jks(veff, dm_kpts, input_band, kpts) + veff = _format_jks(veff, dm_kpts, input_band, input_kpts) veff = tag_array(veff, ecoul=coulomb_energy, exc=xc_energy_sum) t0 = log.timer("xc", *t0) return n_electrons, xc_energy_sum, veff @@ -1755,6 +1750,20 @@ def get_veff_ip1( ngrids = np.prod(mesh) density = evaluate_density_on_g_mesh(ni, dm_kpts, kpts, xc_type) + if with_pseudo_vloc_orbital_derivative: + from gpu4pyscf.pbc.dft.multigrid_v3 import ( + _get_Gv_bases, _pploc_derivatives, _ne_derivatives) + Gv_bases = _get_Gv_bases(mesh, cell.reciprocal_vectors()) + if nset == 1: + rhoG = density[0,0] + else: + rhoG = density[:,0].sum(axis=0) + if cell._pseudo: + grad_pp = _pploc_derivatives(cell, mesh, rhoG, Gv_bases)[0] + else: + grad_pp = _ne_derivatives(cell, mesh, rhoG, Gv_bases)[0] + rhoG = None + Gv = get_Gv(cell, mesh) coulomb_kernel_on_g_mesh = pbc_tools.get_coulG(cell, Gv=Gv) coulomb_on_g_mesh = cp.einsum( @@ -1815,9 +1824,13 @@ def get_veff_ip1( ni, xc_for_fock, dm_kpts, hermi, kpts, with_tau = (xc_type == "MGGA") ) + veff_gradient /= len(kpts) + if with_pseudo_vloc_orbital_derivative: + veff_gradient += grad_pp + t0 = log.timer("veff_gradient", *t0) - return veff_gradient + return veff_gradient.get() def _rks_exc_strain_deriv(ni, xc_code, dm_kpts, kpts=None, with_j=False, with_nuc=False): '''Strain derivatives for Coulomb and Exc with k-point samples @@ -1987,7 +2000,7 @@ def update_pairs_info(cell1): rho1_sf, grids, with_j, with_nuc) return out -class MultiGridNumInt(lib.StreamObject, numint.LibXCMixin): +class MultiGridNumInt(multigrid_v1.MultiGridNumIntBase): def __init__(self, cell): self.cell = cell self.mesh = cell.mesh @@ -2170,6 +2183,59 @@ def cache_xc_kernel1(self, cell, grids, xc_code, dm, spin=0, kpts=None, is_rhf=N cache_xc_kernel = NotImplemented + def energy_nuclear_gradient(self, xc_code, dm_kpts, kpts=None, spin=None, + with_j=False, with_nuc=False): + '''Computes the nuclear gradients of Exc along with additional + contributions from the Coulomb and pseudopotential terms. + + Kwargs: + with_j : + Whether to include the electron-electron Coulomb interactions + with_nuc : + Whether to include the contribution from the local part of + pseudo-potential or electron-nuclear Coulomb interactions + ''' + hermi = 1 + return get_veff_ip1(self, xc_code, dm_kpts, hermi, kpts, with_j, with_nuc) + + def energy_strain_gradient(self, xc_code, dm_kpts, kpts=None, spin=None, + with_j=False, with_nuc=False): + '''Computes the strain derivatives of Exc along with additional + contributions from the Coulomb and pseudopotential terms. + + Kwargs: + with_j : + Whether to include the electron-electron Coulomb interactions + with_nuc : + Whether to include the contribution from the local part of + pseudo-potential or electron-nuclear Coulomb interactions + ''' + if spin is None: + dms = _format_dms(dm_kpts, kpts) + spin = 0 if len(dms) == 1 else 1 + if spin == 0: + sigma = _rks_exc_strain_deriv(self, xc_code, dm_kpts, kpts, with_j, with_nuc) + else: + sigma = _uks_exc_strain_deriv(self, xc_code, dm_kpts, kpts, with_j, with_nuc) + return sigma + + def energy_derivatives(self, xc_code, dm_kpts, kpts=None, spin=None, + with_j=False, with_nuc=False): + '''Computes the nuclear gradients and strain derivatives of Exc + along with additional contributions from the Coulomb and pseudopotential + terms. + + Kwargs: + with_j : + Whether to include the electron-electron Coulomb interactions + with_nuc : + Whether to include the contribution from the local part of + pseudo-potential or electron-nuclear Coulomb interactions + ''' + grad = self.energy_nuclear_gradient(xc_code, dm_kpts, kpts, spin, with_j, with_nuc) + sigma = self.energy_strain_gradient(xc_code, dm_kpts, kpts, spin, with_j, with_nuc) + return grad, sigma + to_gpu = utils.to_gpu device = utils.device diff --git a/gpu4pyscf/pbc/dft/multigrid_v3.py b/gpu4pyscf/pbc/dft/multigrid_v3.py new file mode 100644 index 000000000..537f30154 --- /dev/null +++ b/gpu4pyscf/pbc/dft/multigrid_v3.py @@ -0,0 +1,2399 @@ +#!/usr/bin/env python +# Copyright 2024-2025 The PySCF Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +''' +MultiGrid numerical integration for computing XC matrices + +In the current implementation, the memory required by Vxc integrals are + RKS UKS +LDA 5*N^3 8*N^3 +GGA 8*N^3 14*N^3 +MGGA 9*N^3 18*N^3 + +For 80 GB memory, the upper limits of mesh are approximately + RKS UKS +LDA 1200 1000 +GGA 1000 850 +MGGA 1000 800 +''' + +import math +import ctypes +import numpy as np +import cupy as cp +import cupyx.scipy.fft as fft +from pyscf import lib +from pyscf.gto import ANG_OF, ATOM_OF, PTR_COORD, PTR_EXP, PTR_COEFF, gto_norm +from pyscf.pbc.lib.kpts import KPoints +from pyscf.pbc.df.df_jk import _format_kpts_band +from pyscf.pbc.lib.kpts_helper import is_zero +from pyscf.pbc.tools.pbc import super_cell +from pyscf.pbc.tools.k2gamma import translation_vectors_for_kmesh +from gpu4pyscf.lib import logger +from gpu4pyscf.lib.cupy_helper import ( + contract, transpose_sum, ndarray, asarray, tag_array, load_library, absmax, + get_avail_mem, vec_dot) +from gpu4pyscf.lib.utils import nearest_power2 +from gpu4pyscf.__config__ import props as gpu_specs +from gpu4pyscf.lib import multi_gpu +from gpu4pyscf.dft import numint +from gpu4pyscf.pbc import tools +from gpu4pyscf.pbc.tools import k2gamma, get_coulG +from gpu4pyscf.pbc.lib.kpts_helper import fft_matrix +from gpu4pyscf.pbc.df.fft_jk import _format_dms, _format_jks +from gpu4pyscf.pbc.df.aft import _get_ZSI +from gpu4pyscf.gto.mole import ( + PTR_BAS_COORD, SortedGTO, SortedCell, PBCIntEnvVars, _scale_sp_ctr_coeff) +from gpu4pyscf.pbc.gto.pseudo.pp_int import get_pp_nl_gpu +from gpu4pyscf.pbc.dft import multigrid as multigrid_v1 + +libmgrid = load_library('libmgrid_v3') +NBAS_MAX = 16777216 +LMAX = 4 + +_kernel_registery = {} + +def _aft_eval_density(ni, dm_sc, kpts=None, with_tau=False): + cell = ni.sorted_cell + bvkcell = ni.bvkcell + + a = bvkcell.lattice_vectors() + assert abs(a - np.diag(a.diagonal())).max() < 1e-5, 'Must be orthogonal lattice' + b = cell.reciprocal_vectors() + + nkpts = len(ni.bvkmesh_Ls) + weight = 1./nkpts + + rhoG = cp.zeros(ni.mesh, dtype=np.complex128) + kern = libmgrid.orth_contract_aopair_dm + tauG = None + if with_tau: + tauG = cp.zeros(ni.mesh, dtype=np.complex128) + kern = libmgrid.orth_contract_ft_tau_dm + + for bucket in ni.aft_buckets: + mesh = bucket['mesh'] + mesh_cum = cp.array(np.append(0, np.cumsum(mesh)), dtype=np.int32) + nimgs = bucket['nimgs'] + nimgs_cum = cp.array(np.append(0, np.cumsum(nimgs*2+1)), dtype=np.int32) + Gx, Gy, Gz = _get_Gv_bases(mesh, b) + G_bases = cp.hstack([Gx[0], Gy[1], Gz[2]]) + L_bases = _get_L_bases(nimgs, a) + + # To reduce the overhead of atomicAdd, process multiple pairs for each + # cuda block. + pairs_per_block = 100 + shl_pair_offsets = bucket['shl_pair_offsets'] + offsets = [] + for p0, p1 in zip(shl_pair_offsets[:-1], shl_pair_offsets[1:]): + offsets.append(cp.arange(p0, p1, pairs_per_block, dtype=np.int32)) + offsets.append(np.int32(shl_pair_offsets[-1])) + shl_pair_offsets = cp.hstack(offsets, dtype=np.int32) + nbatches_shl_pair = len(shl_pair_offsets) - 1 + + rhoR = cp.zeros(mesh) + rhoI = cp.zeros(mesh) + tauR = tauI = rhoR + if with_tau: + tauR = cp.zeros(mesh) + tauI = cp.zeros(mesh) + err = kern( + ctypes.cast(rhoR.data.ptr, ctypes.c_void_p), + ctypes.cast(rhoI.data.ptr, ctypes.c_void_p), + ctypes.cast(tauR.data.ptr, ctypes.c_void_p), + ctypes.cast(tauI.data.ptr, ctypes.c_void_p), + ctypes.cast(dm_sc.data.ptr, ctypes.c_void_p), + ctypes.byref(ni.mg_envs), + ctypes.cast(shl_pair_offsets.data.ptr, ctypes.c_void_p), + ctypes.cast(bucket['bas_ij_idx'].data.ptr, ctypes.c_void_p), + ctypes.cast(G_bases.data.ptr, ctypes.c_void_p), + ctypes.cast(L_bases.data.ptr, ctypes.c_void_p), + ctypes.cast(mesh_cum.data.ptr, ctypes.c_void_p), + ctypes.cast(nimgs_cum.data.ptr, ctypes.c_void_p), + (ctypes.c_int*3)(*mesh), + ctypes.c_int(nbatches_shl_pair), + ctypes.c_double(weight)) + if err != 0: + raise RuntimeError('contract_orth_aopair_dm kernel failed') + _inv_take_fft_submesh(rhoG, (rhoR, rhoI), mesh) + if with_tau: + _inv_take_fft_submesh(tauG, (tauR, tauI), mesh) + return rhoG, tauG + +def _aft_eval_xc_matrix(ni, vxcG, out=None): + cell = ni.sorted_cell + bvkcell = ni.bvkcell + + a = bvkcell.lattice_vectors() + b = cell.reciprocal_vectors() + + if isinstance(vxcG, cp.ndarray): + vrhoG = vxcG.reshape(ni.mesh) + vtauG = None + kern = libmgrid.orth_aft_lda_mat + else: + vrhoG, vtauG = vxcG + vrhoG = vrhoG.reshape(ni.mesh) + vtauG = vtauG.reshape(ni.mesh) + kern = libmgrid.orth_aft_mgga_mat + + nao = cell.nao + nkpts = len(ni.bvkmesh_Ls) + vxc_mat = ndarray((nkpts, nao, nao), dtype=np.float64, buffer=out) + vxc_mat.fill(0.) + + for bucket in ni.aft_buckets: + mesh = bucket['mesh'] + mesh_cum = cp.array(np.append(0, np.cumsum(mesh)), dtype=np.int32) + nimgs = bucket['nimgs'] + nimgs_cum = cp.array(np.append(0, np.cumsum(nimgs*2+1)), dtype=np.int32) + # In real space formula, VxcG in reciprocal space is first IFFT to real + # space. Here, AFT integrals for -G are identical to the inverse FT. + Gx, Gy, Gz = _get_Gv_bases(mesh, b) + G_bases = -cp.hstack([Gx[0], Gy[1], Gz[2]]) + L_bases = _get_L_bases(nimgs, a) + + bas_ij_idx = bucket['bas_ij_idx'] + + sub_vrhoG = _take_fft_submesh(vrhoG, mesh) + sub_vtauG = sub_vrhoG + if vtauG is not None: + sub_vtauG = _take_fft_submesh(vtauG, mesh) + err = kern( + ctypes.cast(vxc_mat.data.ptr, ctypes.c_void_p), + ctypes.cast(sub_vrhoG.data.ptr, ctypes.c_void_p), + ctypes.cast(sub_vtauG.data.ptr, ctypes.c_void_p), + ctypes.byref(ni.mg_envs), + ctypes.cast(bas_ij_idx.data.ptr, ctypes.c_void_p), + ctypes.cast(G_bases.data.ptr, ctypes.c_void_p), + ctypes.cast(L_bases.data.ptr, ctypes.c_void_p), + ctypes.cast(mesh_cum.data.ptr, ctypes.c_void_p), + ctypes.cast(nimgs_cum.data.ptr, ctypes.c_void_p), + (ctypes.c_int*3)(*mesh), + ctypes.c_int(len(bas_ij_idx))) + if err != 0: + raise RuntimeError('contract_orth_aopair_coulG kernel failed') + + weight = abs(np.linalg.det(b)) / (2*np.pi)**3 # = get_Gv_weights + vxc_mat *= weight + return vxc_mat + +def _eval_density(ni, dm_sc, kpts=None, with_tau=False): + cell = ni.sorted_cell + if ni.aft_buckets is not None: + rhoG, tauG = _aft_eval_density(ni, dm_sc, kpts, with_tau) + else: + rhoG = cp.zeros(ni.mesh, dtype=np.complex128) + tauG = None + if with_tau: + tauG = cp.zeros(ni.mesh, dtype=np.complex128) + + a = cell.lattice_vectors() + vol = np.linalg.det(a) + nkpts = len(ni.bvkmesh_Ls) + + work = cp.empty_like(rhoG) + if not with_tau: + kern = libmgrid.evaluate_density + else: + kern = libmgrid.evaluate_tau + work1 = cp.empty_like(rhoG) + + mg_envs = ni.mg_envs + + tile_info = None + fft_buckets = ni.fft_buckets or [] + if fft_buckets and fft_buckets[0]['grid_tile_cache'] is None: + tile_info = _grid_range_to_tile_info_converter(fft_buckets, cell) + + for bucket in fft_buckets: + mesh = bucket['mesh'] + ngrids = np.prod(mesh) + weight = vol / ngrids / nkpts + + dxyz_dabc = a / mesh[:,None] + libmgrid.update_dxyz_dabc(dxyz_dabc.ctypes) + + rhoR = ndarray(mesh, dtype=np.complex128, buffer=work) + rhoR.fill(0) + tauR = rhoR # placeholder + if with_tau: + tauR = ndarray(mesh, dtype=np.complex128, buffer=work1) + tauR.fill(0) + + for n, (li, lj) in enumerate(bucket['lij_patterns']): + if tile_info is None: + grid_tile_idx, dressed_bas_ij, shl_pair_offsets = \ + bucket['grid_tile_cache'][n] + else: + grid_tile_idx, dressed_bas_ij, shl_pair_offsets = tile_info( + bucket['bas_ij_cache'][n], bucket['grid_ranges_cache'][n], mesh) + + if len(dressed_bas_ij) == 0: continue + ntiles = len(grid_tile_idx) + tiles_per_block = min(100, max(1, ntiles // 1000)) + err = kern( + ctypes.cast(rhoR.data.ptr, ctypes.c_void_p), + ctypes.cast(tauR.data.ptr, ctypes.c_void_p), + ctypes.cast(dm_sc.data.ptr, ctypes.c_void_p), + ctypes.byref(mg_envs), + dxyz_dabc.ctypes, + ctypes.cast(ni.supmol_img_coords.data.ptr, ctypes.c_void_p), + ctypes.c_int(li), ctypes.c_int(lj), + ctypes.c_int(tiles_per_block), + ctypes.cast(shl_pair_offsets.data.ptr, ctypes.c_void_p), + ctypes.cast(dressed_bas_ij.data.ptr, ctypes.c_void_p), + ctypes.cast(grid_tile_idx.data.ptr, ctypes.c_void_p), + ctypes.c_int(ntiles), + (ctypes.c_int*3)(*mesh), + ctypes.c_double(weight), + ctypes.c_double(bucket['negligible'])) + if err != 0: + raise RuntimeError('evaluate_density kernel failed') + + _inv_take_fft_submesh(rhoG, fft_in_place(rhoR).reshape(mesh), mesh) + if with_tau: + _inv_take_fft_submesh(tauG, fft_in_place(tauR).reshape(mesh), mesh) + + return rhoG, tauG + +# This function is not used by the current implementation. It is provided +# to maintain compatibility with the multigrid_v2 implementation. +def _eval_rhoG(ni, dm_kpts, hermi=1, kpts=None, xctype='LDA'): + assert xctype == 'LDA' + assert not ni.allow_mesh_reduction + dm_sc = _wannier_transform_dm(ni, dm_kpts, kpts, hermi=1) + n_dm, nkpts, nao = dm_sc.shape[:3] + rhoG = cp.array([_eval_density(ni, dm_sc[i])[0] for i in range(n_dm)]) + return rhoG + +def _eval_xc_mat(ni, vxcG, out=None, work=None): + '''Note, contents of vxcG will be destroyed in this function + ''' + cell = ni.sorted_cell + if ni.aft_buckets is not None: + vxc_mat = _aft_eval_xc_matrix(ni, vxcG, out) + else: + nkpts = len(ni.bvkmesh_Ls) + nao = cell.nao + vxc_mat = ndarray((nkpts, nao, nao), dtype=np.float64, buffer=out) + vxc_mat.fill(0.) + + if isinstance(vxcG, cp.ndarray): + vrhoG = vxcG.reshape(ni.mesh) + vtauG = None + work = ndarray((3,vrhoG.size), dtype=np.float64, buffer=work) + kern = libmgrid.evaluate_lda_mat_v2 + else: + vrhoG, vtauG = vxcG + vrhoG = vrhoG.reshape(ni.mesh) + vtauG = vtauG.reshape(ni.mesh) + work = ndarray((4,vrhoG.size), dtype=np.float64, buffer=work) + kern = libmgrid.evaluate_mgga_mat_v2 + + a = cell.lattice_vectors() + + mg_envs = ni.mg_envs + + fft_buckets = ni.fft_buckets or [] + for bucket in fft_buckets: + mesh = bucket['mesh'] + dxyz_dabc = a / mesh[:,None] + libmgrid.update_dxyz_dabc(dxyz_dabc.ctypes) + + # _take_fft_submesh does not always make a copy. In the last bucket, the contents + # of vrhoG will be overwritten by ifft_in_place + sub_vrhoG = _take_fft_submesh(vrhoG, mesh, work[:2]) + sub_vrhoR = ndarray(mesh, dtype=np.float64, buffer=work[2]) + sub_vrhoR[:] = ifft_in_place(sub_vrhoG).real + sub_vtauR = sub_vrhoR # placeholder + + if vtauG is not None: + sub_vtauG = _take_fft_submesh(vtauG, mesh, work[:2]) + sub_vtauR = ndarray(mesh, dtype=np.float64, buffer=work[3]) + sub_vtauR[:] = ifft_in_place(sub_vtauG).real + + for (li, lj), bas_ij_idx, grid_frac_ranges in zip( + bucket['lij_patterns'], bucket['bas_ij_cache'], + bucket['grid_ranges_cache']): + if len(bas_ij_idx) == 0: continue + err = kern( + ctypes.cast(vxc_mat.data.ptr, ctypes.c_void_p), + ctypes.cast(sub_vrhoR.data.ptr, ctypes.c_void_p), + ctypes.cast(sub_vtauR.data.ptr, ctypes.c_void_p), + ctypes.byref(mg_envs), + dxyz_dabc.ctypes, + ctypes.c_int(li), + ctypes.c_int(lj), + ctypes.cast(bas_ij_idx.data.ptr, ctypes.c_void_p), + ctypes.cast(grid_frac_ranges.data.ptr, ctypes.c_void_p), + (ctypes.c_int*3)(*mesh), + ctypes.c_int(len(bas_ij_idx)), + ctypes.c_double(bucket['negligible'])) + if err != 0: + raise RuntimeError('evaluate_xc_mat kernel failed') + return vxc_mat + +# This version is slower than _eval_xc_mat in most scenarios +def _eval_xc_mat_v1(ni, vxcG, out=None, work=None): + '''Note, contents of vxcG will be destroyed in this function + ''' + cell = ni.sorted_cell + if ni.aft_buckets is not None: + vxc_mat = _aft_eval_xc_matrix(ni, vxcG, out) + else: + nkpts = len(ni.bvkmesh_Ls) + nao = cell.nao + vxc_mat = ndarray((nkpts, nao, nao), dtype=np.float64, buffer=out) + vxc_mat.fill(0.) + + a = cell.lattice_vectors() + + if isinstance(vxcG, cp.ndarray): + vrhoG = vxcG.reshape(ni.mesh) + vtauG = None + work = ndarray((3,vrhoG.size), dtype=np.float64, buffer=work) + kern = libmgrid.evaluate_lda_mat + else: + vrhoG, vtauG = vxcG + vrhoG = vrhoG.reshape(ni.mesh) + vtauG = vtauG.reshape(ni.mesh) + work = ndarray((4,vrhoG.size), dtype=np.float64, buffer=work) + kern = libmgrid.evaluate_mgga_mat + + mg_envs = ni.mg_envs + + fft_buckets = ni.fft_buckets or [] + for bucket in fft_buckets: + assert bucket['grid_tile_cache'] is not None + mesh = bucket['mesh'] + dxyz_dabc = a / mesh[:,None] + libmgrid.update_dxyz_dabc(dxyz_dabc.ctypes) + + # _take_fft_submesh does not always make a copy. In the last bucket, the + # contents of vrhoG will be overwritten by ifft_in_place + sub_vrhoG = _take_fft_submesh(vrhoG, mesh, work[:2]) + sub_vrhoR = ndarray(mesh, dtype=np.float64, buffer=work[2]) + sub_vrhoR[:] = ifft_in_place(sub_vrhoG).real + sub_vtauR = sub_vrhoR # placeholder + + if vtauG is not None: + sub_vtauG = _take_fft_submesh(vtauG, mesh, work[:2]) + sub_vtauR = ndarray(mesh, dtype=np.float64, buffer=work[3]) + sub_vtauR[:] = ifft_in_place(sub_vtauG).real + + for ((li, lj), (grid_tile_idx, dressed_bas_ij, shl_pair_offsets)) \ + in zip(bucket['lij_patterns'], bucket['grid_tile_cache']): + if len(dressed_bas_ij) == 0: continue + ntiles = len(grid_tile_idx) + tiles_per_block = min(100, max(1, ntiles // 10000)) + err = kern( + ctypes.cast(vxc_mat.data.ptr, ctypes.c_void_p), + ctypes.cast(sub_vrhoR.data.ptr, ctypes.c_void_p), + ctypes.cast(sub_vtauR.data.ptr, ctypes.c_void_p), + ctypes.byref(mg_envs), + dxyz_dabc.ctypes, + ctypes.cast(ni.supmol_img_coords.data.ptr, ctypes.c_void_p), + ctypes.c_int(li), ctypes.c_int(lj), + ctypes.c_int(tiles_per_block), + ctypes.cast(shl_pair_offsets.data.ptr, ctypes.c_void_p), + ctypes.cast(dressed_bas_ij.data.ptr, ctypes.c_void_p), + ctypes.cast(grid_tile_idx.data.ptr, ctypes.c_void_p), + ctypes.c_int(len(grid_tile_idx)), + (ctypes.c_int*3)(*mesh), + ctypes.c_double(bucket['negligible'])) + if err != 0: + raise RuntimeError('evaluate_xc_mat kernel failed') + return vxc_mat + +def _eval_gradients(ni, dm_sc, vxcG, fft_buckets, work=None): + ''' + Evaluate nuclear gradients and strain gradients together. + + Note, contents of vxcG will be destroyed in this function + ''' + cell = ni.sorted_cell + gradient = cp.zeros((cell.natm, 3)) + sigma = cp.zeros((3, 3)) + + if isinstance(vxcG, cp.ndarray): + vrhoG = vxcG.reshape(ni.mesh) + vtauG = None + work = ndarray((3,vrhoG.size), dtype=np.float64, buffer=work) + kern = libmgrid.evaluate_lda_grad + else: + vrhoG, vtauG = vxcG + vrhoG = vrhoG.reshape(ni.mesh) + vtauG = vtauG.reshape(ni.mesh) + work = ndarray((4,vrhoG.size), dtype=np.float64, buffer=work) + kern = libmgrid.evaluate_mgga_grad + + a = cell.lattice_vectors() + nkpts = len(ni.bvkmesh_Ls) + weight = 1./nkpts + + mg_envs = ni.mg_envs + + for bucket in fft_buckets: + mesh = bucket['mesh'] + dxyz_dabc = a / mesh[:,None] + libmgrid.update_dxyz_dabc(dxyz_dabc.ctypes) + + # _take_fft_submesh does not always make a copy. In the last bucket, the + # contents of vrhoG will be overwritten by ifft_in_place + sub_vrhoG = _take_fft_submesh(vrhoG, mesh, work[:2]) + sub_vrhoR = ndarray(mesh, dtype=np.float64, buffer=work[2]) + sub_vrhoR[:] = ifft_in_place(sub_vrhoG).real + sub_vtauR = sub_vrhoR # placeholder + + if vtauG is not None: + sub_vtauG = _take_fft_submesh(vtauG, mesh, work[:2]) + sub_vtauR = ndarray(mesh, dtype=np.float64, buffer=work[3]) + sub_vtauR[:] = ifft_in_place(sub_vtauG).real + + for (li, lj), bas_ij_idx, grid_frac_ranges in zip( + bucket['lij_patterns'], bucket['bas_ij_cache'], + bucket['grid_ranges_cache']): + if len(bas_ij_idx) == 0: continue + err = kern( + ctypes.cast(gradient.data.ptr, ctypes.c_void_p), + ctypes.cast(sigma.data.ptr, ctypes.c_void_p), + ctypes.cast(dm_sc.data.ptr, ctypes.c_void_p), + ctypes.cast(sub_vrhoR.data.ptr, ctypes.c_void_p), + ctypes.cast(sub_vtauR.data.ptr, ctypes.c_void_p), + ctypes.byref(mg_envs), + dxyz_dabc.ctypes, + ctypes.c_int(li), + ctypes.c_int(lj), + ctypes.cast(bas_ij_idx.data.ptr, ctypes.c_void_p), + ctypes.cast(grid_frac_ranges.data.ptr, ctypes.c_void_p), + (ctypes.c_int*3)(*mesh), + ctypes.c_int(len(bas_ij_idx)), + ctypes.c_double(weight), + ctypes.c_double(bucket['negligible'])) + if err != 0: + raise RuntimeError('evaluate_xc_grad kernel failed') + return gradient, sigma + +def _get_Gv_bases(mesh, b): + Gx = cp.array(np.fft.fftfreq(mesh[0], 1./mesh[0]) * b[0,:,None]) + Gy = cp.array(np.fft.fftfreq(mesh[1], 1./mesh[1]) * b[1,:,None]) + Gz = cp.array(np.fft.fftfreq(mesh[2], 1./mesh[2]) * b[2,:,None]) + return (Gx, Gy, Gz) + +def _get_L_bases(nimgs, a): + Tx = np.arange(-nimgs[0], nimgs[0]+1) * a[0,0] + Ty = np.arange(-nimgs[1], nimgs[1]+1) * a[1,1] + Tz = np.arange(-nimgs[2], nimgs[2]+1) * a[2,2] + L_bases = cp.array(np.hstack([Tx, Ty, Tz])) + return L_bases + +def _estimate_Ecut_and_grid_ranges(ni, bas_ij_idx, ke_max, precision, xctype): + '''Estimate the FFT energy cutoff and the spread of each orbital pair + in real space''' + cell = ni.sorted_cell + # Some orbitals may require high Ecut, sometimes higher than ke_max. + # Use ke_max to limit the highest Ecut. This ensures that these orbital + # pairs are included in the last bucket in _partition_ke_for_fft. + Ecut_by_shell = _estimate_fft_Ecut_per_shell(cell, precision) + Ecut_by_shell[Ecut_by_shell > ke_max] = ke_max + Ecut_by_shell = cp.asarray(Ecut_by_shell, dtype=np.float32) + + npairs = len(bas_ij_idx) + pair_ke = cp.empty(npairs, dtype=np.float32) + grid_frac_ranges = cp.empty((3,npairs,2), dtype=np.float32) + + li_inc = lj_inc = 0 + if xctype == 'MGGA': + li_inc = lj_inc = 1 + + # Higher resolution is required for orbitals centered on the same atom. + # The ke_cut estimated from cell.precision is sufficient for converging the + # Coulomb integrals. However, XC potential is not as smooth near the core + # region. Apply an additional penalty to improve the accuracy of XC integrals. + undressed_threshold = cell.precision * 1e-1 + + err = libmgrid.gaussian_prod_grid_ranges( + ctypes.cast(grid_frac_ranges.data.ptr, ctypes.c_void_p), + ctypes.cast(pair_ke.data.ptr, ctypes.c_void_p), + ctypes.cast(Ecut_by_shell.data.ptr, ctypes.c_void_p), + ctypes.byref(ni.mg_envs), + ctypes.cast(bas_ij_idx.data.ptr, ctypes.c_void_p), + ctypes.c_int(npairs), + ctypes.c_int(li_inc), ctypes.c_int(lj_inc), + ctypes.c_float(math.log(precision)), + ctypes.c_float(undressed_threshold), + ctypes.c_float(ke_max)) + if err != 0: + raise RuntimeError('grid range kernel failed') + return pair_ke, grid_frac_ranges + +def _estimate_fft_Ecut_per_shell(cell, precision): + # To accurately describe the orbital in real space, the resolution for + # real-space grid cannot be reduced, even a small normalized function is + # associated with the orbital. The resolution is estimated in terms of the + # energy cutoff for orbitals with standard normalization. + ai = cell._env[cell._bas[:,PTR_EXP]] + li = cell._bas[:,ANG_OF] + ci = gto_norm(li, ai) + + # Ecut ~ ci * (Ecut/2/ai**2)**(li/2) * exp(-Ecut/(2*ai)) + # = ci / ai**li * E2**(li/2) * exp(-E2/ai), where E2 = Ecut/2 + log_fac = np.log(ci) + 1.717 - (li+1.5)*np.log(ai) - np.log(precision) + log_fac[log_fac <= 0] = 1e-9 + E2 = log_fac * ai + E2 = (log_fac + .5 * li * np.log(E2)) * ai + Ecut = E2 * 2 + return Ecut + +def ke_to_mesh(a, cutoff): + ''' + Based on pyscf.pbc.tools.pbc.cutoff_to_mesh + ''' + b = 2 * np.pi * np.linalg.inv(a.T) + rx = np.linalg.qr(b[[1,2,0]].T)[1][2,2] + ry = np.linalg.qr(b[[2,0,1]].T)[1][2,2] + rz = np.linalg.qr(b.T)[1][2,2] + + Gmax = (2*cutoff)**.5 / np.abs([rx, ry, rz]) + mesh = np.ceil(Gmax * 2).astype(np.int32) + return mesh + +def mesh_to_ke(a, mesh): + ''' + Based on pyscf.pbc.tools.pbc.mesh_to_cutoff + ''' + b = 2 * np.pi * np.linalg.inv(a.T) + rx = np.linalg.qr(b[[1,2,0]].T)[1][2,2] + ry = np.linalg.qr(b[[2,0,1]].T)[1][2,2] + rz = np.linalg.qr(b.T)[1][2,2] + + gs = np.asarray(mesh) / 2 + Gmax = gs * np.array([rx, ry, rz]) + ke_cutoff = Gmax**2 / 2 + return ke_cutoff + +def _partition_ke_for_aft(ni, pair_idx, pair_ke, init_ke, ke_max, xctype, log): + cell = ni.sorted_cell + bvkcell = ni.bvkcell + a = cell.lattice_vectors() + + mesh_max = np.asarray(ni.mesh, dtype=np.int32) + mesh = ke_to_mesh(a, init_ke) + + ang_per_shell = cp.array(bvkcell._bas[:,ANG_OF]) + nimgs = np.asarray(bvkcell.nimgs, dtype=np.int32) + + buckets = [] + + ke_lower, ke_upper = 0, init_ke + while ke_lower <= ke_max: + mesh = np.minimum(mesh, mesh_max) + filtered_pairs = pair_idx[(ke_lower < pair_ke) & (pair_ke <= ke_upper)] + if len(filtered_pairs) > 0: + ish, jsh = divmod(filtered_pairs, NBAS_MAX) + lij = ang_per_shell[ish] * 5 + ang_per_shell[jsh] + idx = cp.argsort(lij) + filtered_pairs = filtered_pairs[idx] + lij = lij[idx] + shl_pair_offsets = _segment_offsets(lij).get() + + # TODO: nimgs can be reduced for large Ecut + # filtered_pairs -> rcut_for_each_pair -> max_rcut -> nimgs + + buckets.append({ + 'ke_cutoff': ke_upper, + 'mesh': np.asarray(mesh, dtype=np.int32), + 'nimgs': nimgs, + 'bas_ij_idx': filtered_pairs, + 'shl_pair_offsets': shl_pair_offsets, + }) + log.debug('Add aft bucket: ke=%g mesh=%s, shl_pairs=%d', ke_upper, + mesh, len(filtered_pairs)) + + mesh = (mesh * 0.75).astype(np.int32) * 2 + mesh[mesh < 8] = 8 + ke_lower, ke_upper = ke_upper, mesh_to_ke(a, mesh).min() + return buckets + +def _partition_ke_for_fft(ni, pair_idx, init_ke, ke_max, precision, xctype, log): + cell = ni.sorted_cell + bvkcell = ni.bvkcell + + a = cell.lattice_vectors() + mesh = ke_to_mesh(a, init_ke) + mesh_final = ni.mesh + + ang_per_shell = cp.array(bvkcell._bas[:,ANG_OF]) + + supmol_bas_ij_idx = _bvk_pairs_to_supmol_pairs( + ni, pair_idx, precision, xctype) + + pair_ke, grid_frac_ranges = _estimate_Ecut_and_grid_ranges( + ni, supmol_bas_ij_idx, ke_max, precision, xctype) + + buckets = [] + + ke_lower, ke_upper = 0, init_ke + while ke_lower < ke_max: + mesh = np.minimum(mesh, mesh_final) + idx = cp.where((ke_lower < pair_ke) & (pair_ke <= ke_upper))[0] + if len(idx) > 0: + filtered_pairs = supmol_bas_ij_idx[idx] + filtered_grid_ranges = grid_frac_ranges[:,idx] + + ish, jsh = divmod(filtered_pairs, NBAS_MAX) + lij = ang_per_shell[ish] * 5 + ang_per_shell[jsh % bvkcell.nbas] + idx = cp.argsort(lij) + lij = lij[idx] + split_points = (cp.where(lij[1:] != lij[:-1])[0] + 1).get() + + #TODO: to avoid too many bas_ij in each sub-bucket, Add more + # split_points and divide idx into more segments. + + # Group bas_ij_idx and grid_frac_ranges by (li, lj) patterns + idx_by_pattern = cp.split(idx, split_points) + lilj_patterns = np.append(lij[0].get(), lij[split_points].get()) + lilj_patterns = [divmod(x, 5) for x in lilj_patterns.tolist()] + + bas_ij_cache = [filtered_pairs[idx] for idx in idx_by_pattern] + grid_ranges_cache = [filtered_grid_ranges[:,idx] for idx in idx_by_pattern] + + # * bas_ij_cache[key] are shell-pairs (one shell in the unit cell, + # the other in supmol) + # * grid_ranges_cache[key] = grid_frac_ranges[3,N,2] + # For each shell pair in bas_ij_idx, stores the fractional-coordinate + # bounds of the real-space grids that are not negligible. + # * grid_tile_cache[key] = (grid_tile_idx, supmol_pair_idx, shl_pair_offsets) + # - grid_tile_idx: + # Unique grid tile indices that contributes to the density. + # - dressed_pair_idx: + # Shell-pair indices contributing to the tiles in grid_tile_idx. + # - shl_pair_offsets: + # Partition the shell pairs in supmol_pair_idx by grid tile. + ao_val_threshold = precision * 1e-5 + buckets.append({ + 'ke_cutoff': ke_upper, + 'mesh': np.asarray(mesh, dtype=np.int32), + 'lij_patterns': lilj_patterns, + 'bas_ij_cache': bas_ij_cache, + 'grid_ranges_cache': grid_ranges_cache, + 'grid_tile_cache': None, + 'negligible': ao_val_threshold + }) + log.debug('Add fft bucket: ke=%g mesh=%s, shl_pairs=%d', + ke_upper, mesh, len(filtered_pairs)) + + mesh = (mesh * 1.2).astype(np.int32) + # For very small initial mesh, such as [2,2,2], mesh*1.2 may not + # increase the mesh, causing the loop stuck + mesh[mesh < 8] = 8 + ke_lower, ke_upper = ke_upper, mesh_to_ke(a, mesh).min() + return buckets + +def _non_trivial_bvk_pairs(ni, precision): + '''Search non-negligible pairs for overlaps''' + cell = ni.sorted_cell + bvkcell = ni.bvkcell + if isinstance(cell, SortedCell): + a = bvkcell.lattice_vectors() + Ls = cp.asarray(lib.cartesian_prod([np.array([0., -1., 1.])]*3).dot(a)) + else: + Ls = cp.zeros((1, 3)) + nimgs = len(Ls) + + nbas = cell.nbas + bvk_nbas = bvkcell.nbas + ovlp_mask = cp.zeros((nbas, bvk_nbas), dtype=bool) + err = libmgrid.bvk_ovlp_mask_estimation( + ctypes.cast(ovlp_mask.data.ptr, ctypes.c_void_p), + ctypes.byref(ni.mg_envs), + ctypes.cast(Ls.data.ptr, ctypes.c_void_p), + ctypes.c_int(nimgs), + ctypes.c_float(math.log(precision))) + if err != 0: + raise RuntimeError('bvk_ovlp_mask_estimation kernel failed') + + ish, jsh = cp.where(ovlp_mask) + bas_ij = ish * NBAS_MAX + jsh + return bas_ij + +def _bvk_pairs_to_supmol_pairs(ni, bas_ij_idx, precision, xctype): + # The bas_ij_idx stores the effective shells in bvkcell. Each of these + # shells involve multiple primitive shells in supmol. Unpack the bvk-shells + # and provide the primitive pair indices in supmol. + nimgs = ni.mg_envs.nimgs + npairs = len(bas_ij_idx) + supmol_bas_ij_idx = cp.empty(npairs * nimgs, dtype=np.int64) + is_mgga = 1 if xctype == 'MGGA' else 0 + counts = cp.empty(1, dtype=np.int32) + err = libmgrid.supmol_non_trivial_pairs( + ctypes.cast(supmol_bas_ij_idx.data.ptr, ctypes.c_void_p), + ctypes.cast(bas_ij_idx.data.ptr, ctypes.c_void_p), + ctypes.byref(ni.mg_envs), + ctypes.c_int(npairs), + ctypes.c_float(math.log(precision)), + ctypes.c_int(is_mgga), + ctypes.cast(counts.data.ptr, ctypes.c_void_p)) + if err != 0: + raise RuntimeError('bvk_ovlp_mask_estimation kernel failed') + supmol_bas_ij_idx = supmol_bas_ij_idx[:int(counts[0].get())] + return supmol_bas_ij_idx + +def _aft_Ecut_estimation(ni, bas_ij_idx, ke_max, precision, xctype='LDA'): + bvkcell = ni.bvkcell + if isinstance(bvkcell, SortedCell): + a = bvkcell.lattice_vectors() + Ls = cp.asarray(lib.cartesian_prod([np.array([0., -1., 1.])]*3).dot(a)) + else: + Ls = cp.zeros((1, 3)) + nimgs = len(Ls) + npairs = len(bas_ij_idx) + is_mgga = 1 if xctype == 'MGGA' else 0 + + Ecut = cp.empty(npairs, dtype=np.float32) + err = libmgrid.estimate_aft_Ecut( + ctypes.cast(Ecut.data.ptr, ctypes.c_void_p), + ctypes.cast(bas_ij_idx.data.ptr, ctypes.c_void_p), + ctypes.byref(ni.mg_envs), + ctypes.cast(Ls.data.ptr, ctypes.c_void_p), + ctypes.c_int(nimgs), + ctypes.c_int(npairs), + ctypes.c_float(math.log(precision)), + # Set the upper limit of Ecut. This ensures all high-Ecut orbital pairs + # are handled by the last bucket in fft_buckets + ctypes.c_float(ke_max), + ctypes.c_int(is_mgga)) + if err != 0: + raise RuntimeError('Ecut kernel failed') + return Ecut + +def _grid_range_to_tile_info_converter(fft_buckets, cell): + buf_size = 0 + for bucket in fft_buckets: + tiles_per_cell = cp.asarray((bucket['mesh']+3) / 4, dtype=np.float32) + for (bas_ij, grid_range) in zip( + bucket['bas_ij_cache'], bucket['grid_ranges_cache']): + raw_tiles = grid_range[:,:,1] - grid_range[:,:,0] + raw_tiles *= tiles_per_cell[:,None] + raw_tiles = cp.ceil(raw_tiles) + raw_tiles += 2 # penalty for rounding on the boundary + n = (raw_tiles[0] * raw_tiles[1] * raw_tiles[2]).sum().get() + buf_size = max(buf_size, int(n)) + + # temporary space to store grid_tile_idx + work = cp.empty(buf_size+1, dtype=np.int32) + # temporary space to store dressed_bas_ij + work1 = cp.empty(buf_size, dtype=np.int64) + + nimgs = cell.nimgs + nbas = cell.nbas + + def tile_info(bas_ij_idx, grid_range, mesh): + npairs = len(bas_ij_idx) + assert npairs > 0 + tile_counts = work[-1:] + err = libmgrid.grid_range_to_tiles( + ctypes.cast(work.data.ptr, ctypes.c_void_p), + ctypes.cast(work1.data.ptr, ctypes.c_void_p), + ctypes.cast(bas_ij_idx.data.ptr, ctypes.c_void_p), + ctypes.cast(grid_range.data.ptr, ctypes.c_void_p), + (ctypes.c_int*3)(*nimgs), + (ctypes.c_int*3)(*mesh), + ctypes.c_int(npairs), + ctypes.c_int(nbas), + ctypes.cast(tile_counts.data.ptr, ctypes.c_void_p)) + if err != 0: + raise RuntimeError('grid_range_to_tiles failed') + n = int(tile_counts[0].get()) + assert n < 2**31, 'int32 indexing in shl_pair_offsets' + + sorted_idx = cp.argsort(work[:n]) + grid_tile_idx = work[sorted_idx] + shl_pair_offsets = _segment_offsets(grid_tile_idx) + + # TODO: Partition large segments in shell_pair_offsets for better load + # balance. + + # Store only the unique grid tile ids. + grid_tile_idx = grid_tile_idx[shl_pair_offsets[:-1]] + + dressed_bas_ij = work1[sorted_idx] + return grid_tile_idx, dressed_bas_ij, shl_pair_offsets + + return tile_info + +def fft_in_place(x): + return fft.fftn(x, axes=(-3, -2, -1), overwrite_x=True) + +def ifft_in_place(x): + return fft.ifftn(x, axes=(-3, -2, -1), overwrite_x=True) + +def _take_fft_submesh(a, mesh, out=None): + assert a.dtype == np.complex128 + assert a.ndim >= 3 + out_shape = mesh = tuple(mesh) + inp_shape = a.shape + if inp_shape[-3:] == out_shape: + return a + + counts = 1 + if a.ndim == 4: + counts, inp_shape = inp_shape[0], inp_shape[1:] + out_shape = (counts,) + mesh + out = ndarray(out_shape, dtype=np.complex128, buffer=out) + err = libmgrid.fft_take( + ctypes.cast(out.data.ptr, ctypes.c_void_p), + ctypes.cast(a.data.ptr, ctypes.c_void_p), + (ctypes.c_int*3)(*mesh), + (ctypes.c_int*3)(*inp_shape), + ctypes.c_int(counts)) + if err != 0: + raise RuntimeError('fft_take kernel failed') + return out + +def _inv_take_fft_submesh(out, a, mesh): + if isinstance(a, cp.ndarray): + assert a.dtype == np.complex128 + else: + aR, aI = a + a = cp.empty(aR.shape, dtype=np.complex128) + a.real = aR + a.imag = aI + assert out.dtype == np.complex128 + assert out.ndim == a.ndim + mesh = tuple(mesh) + assert a.shape[-3:] == mesh + out_shape = out.shape + counts = 1 + if out.ndim == 4: + counts, out_shape = out_shape[0], out_shape[1:] + assert all(x <= y for x, y in zip(mesh, out_shape)), \ + 'folding frequency down unsupported' + err = libmgrid.fft_takebak( + ctypes.cast(out.data.ptr, ctypes.c_void_p), + ctypes.cast(a.data.ptr, ctypes.c_void_p), + (ctypes.c_int*3)(*out_shape), + (ctypes.c_int*3)(*mesh), + ctypes.c_int(counts)) + if err != 0: + raise RuntimeError('fft_takebak kernel failed') + return out + +def _segment_offsets(label, dtype=np.int32): + split_points = cp.nonzero(label[:-1] != label[1:])[0] + 1 + offsets = cp.empty(len(split_points)+2, dtype=dtype) + offsets[0] = 0 + offsets[1:-1] = split_points + offsets[-1] = len(label) + return offsets + +def _conj_dot(a, b): + '''a.conj().dot(b).real''' + return vec_dot(a.view(np.float64), b.view(np.float64)) + +def _apply_Gv_1j(rhoG, Gx, Gy, Gz, out=None): + '''einsum('g,g->g', rhoG, Gv[:,n]*1j), n is 0, 1 or 2''' + fn_name = 'apply_Gv_1j' + if fn_name not in _kernel_registery: + kernel_code = ('''\ +#include +extern "C" __global__ +void ''' + fn_name + r'''(cuDoubleComplex* __restrict__ out, cuDoubleComplex *rhoG, + double *Gx, double *Gy, double *Gz, long long nx, long long ny, long long nz) { + int idx = blockDim.x * blockIdx.x + threadIdx.x; + int stride = gridDim.x * blockDim.x; + size_t nyz = ny * nz; + size_t ng = nx * nyz; + for (size_t g = idx; g < ng; g += stride) { + int ix = g / nyz; + int iyz = g - nyz * ix; + int iy = iyz / nz; + int iz = iyz - nz * iy; + cuDoubleComplex rho = rhoG[g]; + double Gv = Gx[ix] + Gy[iy] + Gz[iz]; + out[g] = make_cuDoubleComplex(-Gv * cuCimag(rho), Gv * cuCreal(rho)); + } +}''') + _kernel_registery[fn_name] = cp.RawKernel(kernel_code, fn_name) + + kernel = _kernel_registery[fn_name] + out = ndarray(rhoG.shape, buffer=out, dtype=np.complex128) + workers = gpu_specs['multiProcessorCount'] + kernel((workers*2,), (1024,), (out, rhoG, Gx, Gy, Gz, len(Gx), len(Gy), len(Gz))) + return out + +def _contract_Gv_1j(out, xc, Gx, Gy, Gz): + '''out += einsum('g,g->g', xc[n], Gv[:,n]*-1j), n is 0, 1 or 2''' + fn_name = 'contract_Gv_1j' + if fn_name not in _kernel_registery: + kernel_code = ('''\ +#include +extern "C" __global__ +void ''' + fn_name + r'''(cuDoubleComplex* __restrict__ out, cuDoubleComplex *vxcG, + double *Gx, double *Gy, double *Gz, long long nx, long long ny, long long nz) { + int idx = blockDim.x * blockIdx.x + threadIdx.x; + int stride = gridDim.x * blockDim.x; + size_t nyz = ny * nz; + size_t ng = nx * nyz; + for (size_t g = idx; g < ng; g += stride) { + int ix = g / nyz; + int iyz = g - nyz * ix; + int iy = iyz / nz; + int iz = iyz - nz * iy; + // (-i Gv) * v + double Gv = Gx[ix] + Gy[iy] + Gz[iz]; + cuDoubleComplex res = out[g]; + cuDoubleComplex v = vxcG[g]; + res.x += Gv * cuCimag(v); + res.y -= Gv * cuCreal(v); + out[g] = res; + } +}''') + _kernel_registery[fn_name] = cp.RawKernel(kernel_code, fn_name) + + kernel = _kernel_registery[fn_name] + workers = gpu_specs['multiProcessorCount'] + kernel((workers*2,), (1024,), (out, xc, Gx, Gy, Gz, len(Gx), len(Gy), len(Gz))) + return out + +def _get_coulomb_in_place(rhoG, Gv_bases): + ''' + Computes + Ecoul = rhoG.conj().dot(rhoG * 4pi/G^2) + rhoG *= 4pi/G^2 + ''' + fn_name = 'get_coulG' + if fn_name not in _kernel_registery: + kernel_code = ('''\ +extern "C" __global__ +void ''' + fn_name + r'''(double *energy, double2* __restrict__ rhoG, + double *Gx, double *Gy, double *Gz, long long nx, long long ny, long long nz) { + int tid = threadIdx.x; + int idx = blockIdx.x * blockDim.x + tid; + int stride = gridDim.x * blockDim.x; + size_t nyz = ny * nz; + size_t ng = nx * nyz; + double Ecoul = 0; + for (size_t g = idx; g < ng; g += stride) { + int ix = g / nyz; + int iyz = g - nyz * ix; + int iy = iyz / nz; + int iz = iyz - nz * iy; + double GG = 0; + for (int n = 0; n < 3; ++n) { + double Gv = Gx[n*nx+ix] + Gy[n*ny+iy] + Gz[n*nz+iz]; + GG += Gv * Gv; + } + double2 coul = {0., 0.}; + if (GG != 0) { + double fac = 12.566370614359172 / GG; + double2 rho = rhoG[g]; + coul.x = fac * rho.x; + coul.y = fac * rho.y; + Ecoul += coul.x * rho.x + coul.y * rho.y; + } + rhoG[g] = coul; + } + for (int offset = 16; offset > 0; offset >>= 1) { + Ecoul += __shfl_down_sync(0xffffffff, Ecoul, offset); + } + __shared__ double swap[32]; + int lane = tid % 32; + int warp = tid / 32; + if (lane == 0) swap[warp] = Ecoul; + __syncthreads(); + int num_warps = blockDim.x / 32; + if (warp == 0) { + Ecoul = (lane < num_warps) ? swap[lane] : 0.; + for (int offset = 16; offset > 0; offset >>= 1) { + Ecoul += __shfl_down_sync(0xffffffff, Ecoul, offset); + } + } + if (tid == 0) atomicAdd(energy, Ecoul); +}''') + _kernel_registery[fn_name] = cp.RawKernel(kernel_code, fn_name) + + kernel = _kernel_registery[fn_name] + nx, ny, nz = [x.shape[1] for x in Gv_bases] + ng = nx * ny * nz + assert rhoG.size == ng + coul_energy = cp.zeros(1) + workers = gpu_specs['multiProcessorCount'] + kernel((workers*2,), (1024,), + (coul_energy, rhoG, Gv_bases[0], Gv_bases[1], Gv_bases[2], nx, ny, nz)) + return coul_energy[0], rhoG + +def _coulomb_strain_derivatives(cell, mesh, rhoG, Gv_bases): + assert cell.dimension == 3 + fn_name = 'coulomb_strain_derivatives' + if fn_name not in _kernel_registery: + kernel_code = r''' +#include +extern "C" __global__ +void ''' + fn_name + '''( +double *out, double *Gx, double *Gy, double *Gz, +complex *rhoG, +int nx, int ny, int nz) +{ + int tid = threadIdx.x; + int idx = blockDim.x * blockIdx.x + tid; + int stride = gridDim.x * blockDim.x; + size_t nyz = ny * nz; + size_t ng = nx * nyz; + double sigma[9] = {}; + for (size_t i_grid = idx; i_grid < ng; i_grid += stride) { + int ix = i_grid / nyz; + int iyz = i_grid - nyz * ix; + int iy = iyz / nz; + int iz = iyz - nz * iy; + double Gv[3]; + double G2 = 0; + for (int n = 0; n < 3; n++) { + Gv[n] = Gx[nx*n+ix] + Gy[ny*n+iy] + Gz[nz*n+iz]; + G2 += Gv[n] * Gv[n]; + } + double prod = 0.; + if (G2 != 0) prod = 12.566370614359172 * 2 / (G2 * G2); + complex r = rhoG[i_grid]; + prod *= r.real() * r.real() + r.imag() * r.imag(); + for (int i = 0; i < 3; i++) { + for (int j = 0; j < 3; j++) { + sigma[i*3+j] += prod * Gv[i] * Gv[j]; + } } + } + __shared__ double swap[32]; + int lane = tid % 32; + int warp = tid / 32; + int num_warps = blockDim.x / 32; + for (int n = 0; n < 9; n++) { + for (int offset = 16; offset > 0; offset >>= 1) { + sigma[n] += __shfl_down_sync(0xffffffff, sigma[n], offset); + } + if (lane == 0) swap[warp] = sigma[n]; + __syncthreads(); + if (warp == 0) { + sigma[n] = (lane < num_warps) ? swap[lane] : 0.; + for (int offset = num_warps/2; offset > 0; offset >>= 1) { + sigma[n] += __shfl_down_sync(0xffffffff, sigma[n], offset); + } + } + if (tid == 0) atomicAdd(out+n, sigma[n]); + } +}''' + _kernel_registery[fn_name] = cp.RawKernel(kernel_code, fn_name) + kernel = _kernel_registery[fn_name] + + nx, ny, nz = [x.shape[1] for x in Gv_bases] + assert rhoG.size == nx * ny * nz + + out = cp.zeros((3, 3)) + + workers = gpu_specs['multiProcessorCount'] + kernel((workers*4,), (512,), + [out, Gv_bases[0], Gv_bases[1], Gv_bases[2], rhoG, + cp.int32(mesh[0]), cp.int32(mesh[1]), cp.int32(mesh[2])]) + + out /= cell.vol + return out + +def _pploc_derivatives(cell, mesh, rhoG, Gv_bases): + assert cell.dimension == 3 + fn_name = 'pploc_strain_derivatives' + if fn_name not in _kernel_registery: + kernel_code = r''' +#include +extern "C" __global__ +void ''' + fn_name + '''( +double *grad, double *strain, double *Gx, double *Gy, double *Gz, +complex *SIx, complex *SIy, complex *SIz, +complex *rhoG, +int nx, int ny, int nz, +int i_atom, double charge, double rloc, +int nexp, double cexp0, double cexp1, double cexp2, double cexp3) +{ + int tid = threadIdx.x; + int idx = blockDim.x * blockIdx.x + tid; + int stride = gridDim.x * blockDim.x; + size_t nyz = ny * nz; + size_t ng = nx * nyz; + double de[3] = {}; + double sigma[9] = {}; + for (size_t i_grid = idx; i_grid < ng; i_grid += stride) { + int ix = i_grid / nyz; + int iyz = i_grid - nyz * ix; + int iy = iyz / nz; + int iz = iyz - nz * iy; + double Gv[3]; + double G2 = 0; + for (int n = 0; n < 3; n++) { + Gv[n] = Gx[nx*n+ix] + Gy[ny*n+iy] + Gz[nz*n+iz]; + G2 += Gv[n] * Gv[n]; + } + double rloc2 = rloc * rloc; + double rloc3 = rloc2 * rloc; + double G2_red = G2 * rloc2; + double expx = exp(-0.5 * G2_red); + double coef1 = 0.; + double coulG = 0; + if (G2 != 0) { + coulG = 12.566370614359172 / G2 * -charge; + coef1 = coulG * (2 / G2 + rloc2); + } + + double cfacs = 0; + double dcfacs = 0; + if (nexp >= 1) cfacs += cexp0; + if (nexp >= 2) { + cfacs += cexp1 * (3 - G2_red); + dcfacs -= cexp1; + } + if (nexp >= 3) { + cfacs += cexp2 * (15 + G2_red * (G2_red - 10)); + dcfacs += cexp2 * (-10 + 2*G2_red); + } + if (nexp >= 4) { + cfacs += cexp3 * (105 + G2_red * (G2_red * (21 - G2_red) - 105)); + dcfacs += cexp3 * (-105 + 42*G2_red - 3*G2_red*G2_red); + } + double coef2 = 15.749609945722419 * rloc2 * rloc3 * (cfacs - 2 * dcfacs); + + complex SI_x = SIx[i_atom * nx + ix]; + complex SI_y = SIy[i_atom * ny + iy]; + complex SI_z = SIz[i_atom * nz + iz]; + complex SI = SI_x * SI_y * SI_z * expx; + complex density = rhoG[i_grid]; + + double prod = density.real() * SI.real() + density.imag() * SI.imag(); + prod *= coef1 + coef2; + for (int i = 0; i < 3; i++) { + for (int j = 0; j < 3; j++) { + sigma[i*3+j] += prod * Gv[i] * Gv[j]; + } } + + // -1j*Gv.T*rhoG.conj().dot(coulG*SI) + prod = density.real() * SI.imag() - density.imag() * SI.real(); + if (G2 == 0) { + prod *= 15.749609945722419 * rloc3 * cfacs; + double vlocG0 = 2 * 3.141592653589793 * charge * rloc2; + prod -= density.imag() * vlocG0; + } else { + prod *= coulG + 15.749609945722419 * rloc3 * cfacs; + } + de[0] += prod * Gv[0]; + de[1] += prod * Gv[1]; + de[2] += prod * Gv[2]; + } + __shared__ double swap[32]; + int lane = tid % 32; + int warp = tid / 32; + int num_warps = blockDim.x / 32; + for (int offset = 16; offset > 0; offset >>= 1) { + for (int n = 0; n < 9; n++) { + sigma[n] += __shfl_down_sync(0xffffffff, sigma[n], offset); + } + for (int n = 0; n < 3; n++) { + de[n] += __shfl_down_sync(0xffffffff, de[n], offset); + } + } + if (lane == 0) { + for (int n = 0; n < 9; n++) atomicAdd(strain+n, sigma[n]); + for (int n = 0; n < 3; n++) atomicAdd(grad+i_atom*3+n, de[n]); + } +}''' + _kernel_registery[fn_name] = cp.RawKernel(kernel_code, fn_name) + kernel = _kernel_registery[fn_name] + + nx, ny, nz = [x.shape[1] for x in Gv_bases] + assert rhoG.size == nx * ny * nz + + coords = cp.asarray(cell.atom_coords()) + SIx = cp.exp(-1j * coords.dot(Gv_bases[0])) + SIy = cp.exp(-1j * coords.dot(Gv_bases[1])) + SIz = cp.exp(-1j * coords.dot(Gv_bases[2])) + + charges = cell.atom_charges() + + grad = cp.zeros((cell.natm, 3)) + sigma = cp.zeros((3, 3)) + + for ia in range(cell.natm): + symb = cell.atom_symbol(ia) + assert symb in cell._pseudo + + pp = cell._pseudo[symb] + rloc, nexp, cexp = pp[1:3+1] + + cexp = [cp.float64(x) for x in cexp] + [cp.float64(0.)] * 4 + kernel_parameters = [ + grad, sigma, Gv_bases[0], Gv_bases[1], Gv_bases[2], + SIx, SIy, SIz, rhoG, + cp.int32(mesh[0]), cp.int32(mesh[1]), cp.int32(mesh[2]), + cp.int32(ia), cp.float64(charges[ia]), cp.float64(rloc), + cp.int32(nexp)] + cexp[:4] + workers = gpu_specs['multiProcessorCount'] + kernel((workers*4,), (512,), kernel_parameters) + + vol = cell.vol + grad /= vol + sigma /= vol + return grad, sigma + +def _ne_derivatives(cell, mesh, rhoG, Gv_bases): + '''Contributions of nuclus-electron interactions''' + assert cell.dimension == 3 + fn_name = 'ne_derivatives' + if fn_name not in _kernel_registery: + kernel_code = r''' +#include +extern "C" __global__ +void ''' + fn_name + '''( +double *grad, double *strain, double *Gx, double *Gy, double *Gz, +complex *SIx, complex *SIy, complex *SIz, +complex *rhoG, double *charges, +int nx, int ny, int nz, int natm) +{ + int tid = threadIdx.x; + int lane = tid % 32; + int warp = tid / 32; + int num_warps = blockDim.x / 32; + int idx = 32 * blockIdx.x + lane; + int stride = gridDim.x * 32; + size_t nyz = ny * nz; + size_t ng = nx * nyz; + double sigma[9] = {}; + for (int i_atom = warp; i_atom < natm; i_atom += num_warps) { + double charge = charges[i_atom]; + double de[3] = {}; + for (size_t i_grid = idx; i_grid < ng; i_grid += stride) { + int ix = i_grid / nyz; + int iyz = i_grid - nyz * ix; + int iy = iyz / nz; + int iz = iyz - nz * iy; + double Gv[3]; + double G2 = 0; + for (int n = 0; n < 3; n++) { + Gv[n] = Gx[nx*n+ix] + Gy[ny*n+iy] + Gz[nz*n+iz]; + G2 += Gv[n] * Gv[n]; + } + double coulG = 0; + if (G2 != 0) { coulG = 12.566370614359172 / G2 * -charge; } + complex SI_x = SIx[i_atom * nx + ix]; + complex SI_y = SIy[i_atom * ny + iy]; + complex SI_z = SIz[i_atom * nz + iz]; + complex SI = SI_x * SI_y * SI_z; + complex density = rhoG[i_grid]; + double prod = density.real() * SI.real() + density.imag() * SI.imag(); + prod *= coulG; + if (G2 != 0) prod *= 2 / G2; + for (int i = 0; i < 3; i++) { + for (int j = 0; j < 3; j++) { + sigma[i*3+j] += prod * Gv[i] * Gv[j]; + } } + // -1j*Gv.T*rhoG.conj().dot(coulG*SI) + prod = density.real() * SI.imag() - density.imag() * SI.real(); + prod *= coulG; + de[0] += prod * Gv[0]; + de[1] += prod * Gv[1]; + de[2] += prod * Gv[2]; + } + for (int offset = 16; offset > 0; offset >>= 1) { + for (int n = 0; n < 3; n++) { + de[n] += __shfl_down_sync(0xffffffff, de[n], offset); + } + } + if (lane == 0) { + for (int n = 0; n < 3; n++) atomicAdd(grad+i_atom*3+n, de[n]); + } + } + for (int offset = 16; offset > 0; offset >>= 1) { + for (int n = 0; n < 9; n++) { + sigma[n] += __shfl_down_sync(0xffffffff, sigma[n], offset); + } + } + if (lane == 0) { + for (int n = 0; n < 9; n++) atomicAdd(strain+n, sigma[n]); + } +}''' + _kernel_registery[fn_name] = cp.RawKernel(kernel_code, fn_name) + kernel = _kernel_registery[fn_name] + + nx, ny, nz = [x.shape[1] for x in Gv_bases] + assert rhoG.size == nx * ny * nz + + coords = cp.asarray(cell.atom_coords()) + SIx = cp.exp(-1j * coords.dot(Gv_bases[0])) + SIy = cp.exp(-1j * coords.dot(Gv_bases[1])) + SIz = cp.exp(-1j * coords.dot(Gv_bases[2])) + + charges = cp.asarray(cell.atom_charges(), dtype=np.float64) + natm = len(charges) + + grad = cp.zeros((natm, 3)) + sigma = cp.zeros((3, 3)) + + workers = gpu_specs['multiProcessorCount'] + kernel((workers*4,), (256,), + [grad, sigma, Gv_bases[0], Gv_bases[1], Gv_bases[2], + SIx, SIy, SIz, rhoG, charges, + cp.int32(mesh[0]), cp.int32(mesh[1]), cp.int32(mesh[2]), + cp.int32(natm)]) + + vol = cell.vol + grad /= vol + sigma /= vol + return grad, sigma + +def _xc_var_length(xctype): + if xctype == 'LDA' or xctype == 'HF': + nvar = 1 + elif xctype == 'GGA': + nvar = 4 + elif xctype == 'MGGA': + nvar = 5 + else: + raise RuntimeError(f'{xctype} not supported') + return nvar + +def _density_to_real_space(rhoG, tauG, Gv_bases, xctype, out=None): + ''' + Perform + stack(ifft(rhoG), cp.einsum('g,gx->xg', rhoG, 1j*Gv), ifft(tauG)).real + with reduced memory footprint. + + Note, this function will use tauG as workspace and the contents of tauG will + be destroyed + ''' + assert rhoG.ndim == 3 + mesh = rhoG.shape + + nvar = _xc_var_length(xctype) + out = ndarray((nvar, *mesh), dtype=np.float64, buffer=out) + if xctype == 'MGGA': + assert tauG is not None + tauR = ifft_in_place(tauG.reshape(mesh)) + out[4] = tauR.real + work = tauG + else: + work = ndarray(mesh, dtype=np.complex128, buffer=tauG) + + if xctype == 'GGA' or xctype == 'MGGA': + Gx, Gy, Gz = Gv_bases + for n in range(3): + work = _apply_Gv_1j(rhoG, Gx[n], Gy[n], Gz[n], work) + out[n+1] = ifft_in_place(work).real + + work[:] = rhoG + out[0] = ifft_in_place(work).real + return out.reshape(nvar, -1) + +def _vxc_to_reciprocal_space(vxc, out, Gv_bases=None, work=None): + ''' + Perform + out += fft(vxc[0]) + cp.einsum('xg,gx->g', fft(vxc[0]), -1j*Gv) + with reduced memory footprint + ''' + assert vxc.ndim == 4 + mesh = vxc.shape[1:] + + work = ndarray(mesh, dtype=np.complex128, buffer=work) + + work.real = vxc[0] + work.imag.fill(0.) + fft_in_place(work) + out += work + + if len(vxc) >= 4: # GGA or MGGA + Gx, Gy, Gz = Gv_bases + for n in range(3): + work.real = vxc[n+1] + work.imag.fill(0.) + _contract_Gv_1j(out, fft_in_place(work), Gx[n], Gy[n], Gz[n]) + + if len(vxc) == 5: # MGGA + work.real = vxc[4] + work.imag.fill(0.) + vxcG_tau = fft_in_place(work) + return out, vxcG_tau + else: + return out + +def _wannier_transform_dm(ni, dm_kpts, kpts, hermi=1, xctype='LDA', out=None): + if kpts is None: + kpts = np.zeros((1, 3)) + else: + kpts = kpts.reshape(-1, 3) + + ni._ensure_initialized(kpts, xctype) + + cell = ni.sorted_cell + dm_kpts = cp.asarray(dm_kpts, order='C') + dms = _format_dms(dm_kpts, kpts) + n_dm, nkpts, nao = dms.shape[:3] + + if hermi != 1: + # the integral kernel only processes tril part of orbital-pairs. + # Due to the symmetry in integrals, the triu contributions can be folded + # into the tril part. + dms = cp.array(dms, copy=True).reshape(n_dm*nkpts,nao,nao) + dms = transpose_sum(dms).reshape(n_dm, nkpts, nao, nao) + + bvk_ncells = len(ni.bvkmesh_Ls) + if bvk_ncells == 1: + dm_sc = dms + else: + if bvk_ncells != nkpts: + expLk = cp.exp(1j*cp.asarray(ni.bvkmesh_Ls).dot(cp.asarray(kpts).T)) + else: + expLk = fft_matrix(ni.kmesh) + dm_sc = contract('nkpq,Lk->nLqp', dms, expLk) + assert absmax(dm_sc.imag) < cell.precision*5e2 + dm_sc = cp.asarray(dm_sc.real, order='C') + + dm_sc = cell.apply_C_mat_CT(dm_sc.reshape(-1,nao,nao), out=out) + + if hermi == 1: + dm_sc *= 2 + + nao = dm_sc.shape[-1] + dm_sc = dm_sc.reshape(n_dm, -1, nao, nao) + return dm_sc + +def _inverse_wannier_transform_fock(ni, veff, kpts): + veff = ni.sorted_cell.apply_CT_mat_C(veff) + + bvk_ncells = len(ni.bvkmesh_Ls) + if bvk_ncells != 1: + if kpts is not None and len(kpts) != bvk_ncells: + expLk = cp.exp(1j*cp.asarray(ni.bvkmesh_Ls).dot(cp.asarray(kpts).T)) + else: + expLk = fft_matrix(ni.kmesh) + nkpts = expLk.shape[1] + expLkz = expLk.view(np.float64).reshape(bvk_ncells, nkpts, 2) + veff = contract('Lpq,Lkz->kpqz', veff, expLkz) + veff = veff.view(np.complex128)[:,:,:,0] + + veff = transpose_sum(veff) + return veff + +def get_rho(ni, dm_kpts, kpts=None): + '''Density in real space + + Args: + ni: + MultiGridNumInt instance + dm: + density matrix at a single k-point or density matrices for k-sampling + + Kwargs: + kpts: (N, 3) ndarray + k points. If not specified, gamma point is assumed + ''' + assert dm_kpts.ndim < 4 + cell = ni.cell + mesh = ni.mesh + + dm_sc = _wannier_transform_dm(ni, dm_kpts, kpts, hermi=1) + n_dm, nkpts, nao = dm_sc.shape[:3] + assert n_dm == 1 + dm_sc = dm_sc[0] + + rhoG = _eval_density(ni, dm_sc)[0] + rhoR = ifft_in_place(rhoG.reshape(mesh)).real.ravel() + weight = cell.vol / np.prod(mesh) + rhoR *= 1./weight + return rhoR + +def get_nuc(ni, kpts=None): + cell = ni.cell + is_single_kpt = kpts is not None and kpts.ndim == 1 + if kpts is None: + kpts = np.zeros((1, 3)) + else: + kpts = kpts.reshape(-1, 3) + + ni._ensure_initialized(kpts, 'LDA') + + mesh = ni.mesh + ZSI = _get_ZSI(cell, mesh) + Gv_bases = _get_Gv_bases(mesh, cell.reciprocal_vectors()) + vneG = _get_coulomb_in_place(ZSI, Gv_bases)[1] + + vne = _eval_xc_mat(ni, vneG) + vne = _inverse_wannier_transform_fock(ni, vne, kpts) + if is_single_kpt: + vne = vne[0] + return vne + +def get_pp(ni, kpts=None): + """Get the periodic pseudopotential nuc-el AO matrix, with G=0 removed. + """ + cell = ni.cell + log = logger.new_logger(cell) + t0 = log.init_timer() + + is_single_kpt = kpts is not None and kpts.ndim == 1 + if kpts is None: + kpts = np.zeros((1, 3)) + else: + kpts = kpts.reshape(-1, 3) + + ni._ensure_initialized(kpts, 'LDA') + + mesh = ni.mesh + # Compute the vpplocG as + # -einsum('ij,ij->j', pseudo.get_vlocG(cell, Gv), cell.get_SI(Gv)) + vpplocG = multigrid_v1.eval_vpplocG(cell, mesh) + vpp = _eval_xc_mat(ni, vpplocG) + vpp = _inverse_wannier_transform_fock(ni, vpp, kpts) + t1 = log.timer_debug1("vpploc", *t0) + + vppnl = get_pp_nl_gpu(cell, kpts) + if kpts is None or is_zero(kpts): + vpp += vppnl[0].real + else: + vpp += vppnl + + if is_single_kpt: + vpp = vpp[0] + log.timer_debug1("vppnl", *t1) + log.timer("get_pp", *t0) + return vpp + +def get_j_kpts(ni, dm_kpts, hermi=1, kpts=None, kpts_band=None): + '''Get the Coulomb (J) AO matrix at sampled k-points. + + Args: + dm_kpts : (*, nkpts, nao, nao) ndarray or a list of (nkpts,nao,nao) ndarray + Density matrix at each k-point. If a list of k-point DMs, eg, + UHF alpha and beta DM, the alpha and beta DMs are contracted + separately. + kpts : (nkpts, 3) ndarray + + Kwargs: + kpts_band : ``(3,)`` ndarray or ``(*,3)`` ndarray + A list of arbitrary "band" k-points at which to evalute the matrix. + + Returns: + vj : (*, nkpts, nao, nao) ndarray + or list of vj if the input dm_kpts is a list of DMs + ''' + assert dm_kpts.ndim < 4 + return nr_rks(ni, ni.cell, None, 'HF', dm_kpts, hermi=hermi, + kpts=kpts, kpts_band=kpts_band, with_j=True)[2] + +def nr_rks(ni, cell, grids, xc_code, dm_kpts, relativity=0, hermi=1, + kpts=None, kpts_band=None, with_j=False, verbose=None): + '''Compute the XC energy and RKS XC matrix at sampled k-points. + multigrid version of function pbc.dft.numint.nr_rks. + + Args: + dm_kpts : (nkpts, nao, nao) ndarray or a list of (nkpts,nao,nao) ndarray + Density matrix at each k-point. + kpts : (nkpts, 3) ndarray + + Kwargs: + kpts_band : ``(3,)`` ndarray or ``(*,3)`` ndarray + A list of arbitrary "band" k-points at which to evalute the matrix. + with_j : bool + Whether to add the Coulomb matrix into the XC matrix. + + Returns: + exc : XC energy + nelec : number of electrons obtained from the numerical integration + veff : (nkpts, nao, nao) ndarray + or list of veff if the input dm_kpts is a list of DMs + ''' + log = logger.new_logger(cell, verbose) + t0 = log.init_timer() + + xctype = ni._xc_type(xc_code) + nvar = _xc_var_length(xctype) + + dm_sc = _wannier_transform_dm(ni, dm_kpts, kpts, hermi, xctype) + assert len(dm_sc) == 1 + dm_sc = dm_sc[0] + + cell = ni.cell + mesh = ni.mesh + ngrids = np.prod(mesh) + vol = cell.vol + weight = vol / ngrids + Gv_bases = _get_Gv_bases(mesh, cell.reciprocal_vectors()) + + rhoG, tauG = _eval_density(ni, dm_sc, with_tau=xctype=='MGGA') + n_electrons = float(rhoG[0,0,0].real.get()) + + if xctype == 'HF': + assert with_j + ecoul, coulomb_on_g_mesh = _get_coulomb_in_place(rhoG, Gv_bases) + ecoul = (.5 / vol) * float(ecoul.get()) + xc_for_fock = coulomb_on_g_mesh + log.debug('Multigrid Coulomb energy %s', ecoul) + rhoG = coulomb_on_g_mesh = None + xc_energy_sum = None + + else: + # dm_sc is represented in primitive bases (by sorted_cell). Its size can be + # much larger than the input dm_kpts. Release its memory if remaining memory + # is insufficient. + if (nvar+4)*ngrids*8 > get_avail_mem(): + dm_sc = None + + density = cp.empty((nvar, ngrids)) + _density_to_real_space(rhoG, tauG, Gv_bases, xctype, out=density) + # *(1./weight) because rhoR is scaled by weight in _eval_density. If + # computing rhoR with IFFT, the weight factor is not needed. + density *= 1/weight + t0 = log.timer_debug1("density", *t0) + + rho_sf = ndarray(ngrids, dtype=np.float64, buffer=tauG) + rho_sf[:] = density[0].real + + # eval_xc_eff supports float64 only + xc_for_energy, xc_for_fock = ni.eval_xc_eff( + xc_code, density, deriv=1, xctype=xctype, spin=0, inplace=True)[:2] + + xc_for_fock = xc_for_fock.reshape(nvar, *mesh) + + xc_energy_sum = float(vec_dot(rho_sf, xc_for_energy).get()) * weight + xc_for_energy = density = rho_sf = None + log.debug("Multigrid exc %s nelec %s", xc_energy_sum, n_electrons) + t0 = log.timer_debug1("eval_xc_eff", *t0) + + if with_j: + ecoul, coulomb_on_g_mesh = _get_coulomb_in_place(rhoG, Gv_bases) + ecoul = (.5 / vol) * float(ecoul.get()) + log.debug('Multigrid Coulomb energy %s', ecoul) + else: + ecoul = None + coulomb_on_g_mesh = rhoG + coulomb_on_g_mesh.fill(0) + rhoG = None + + xc_for_fock *= weight + # Now xc_for_fock represents xc on G space + xc_for_fock = _vxc_to_reciprocal_space( + xc_for_fock, coulomb_on_g_mesh, Gv_bases, work=tauG) + coulomb_on_g_mesh = tauG = None + + if kpts_band is None: + veff = _eval_xc_mat(ni, xc_for_fock, out=dm_sc) + veff = _inverse_wannier_transform_fock(ni, veff, kpts) + else: + kpts_band = kpts_band.reshape(-1, 3) + kmesh = k2gamma.kpts_to_kmesh(cell, kpts_band) + ni = ni.copy().reset().build(kmesh=kmesh, xctype=xctype) + # ni.build may alter the mesh. vxc was created with mesh different to + # this new mesh. + ni.mesh = mesh + veff = _eval_xc_mat(ni, xc_for_fock) + veff = _inverse_wannier_transform_fock(ni, veff, kpts_band) + + veff = _format_jks(veff, dm_kpts, kpts_band, kpts) + veff = tag_array(veff, ecoul=ecoul, exc=xc_energy_sum) + t0 = log.timer_debug1("xc matrix", *t0) + return n_electrons, xc_energy_sum, veff + +def nr_uks(ni, cell, grids, xc_code, dm_kpts, relativity=0, hermi=1, + kpts=None, kpts_band=None, with_j=False, verbose=None): + '''Compute the XC energy and UKS XC matrix at sampled k-points. + multigrid version of function pbc.dft.numint.nr_rks. + + Args: + dm_kpts : (nkpts, nao, nao) ndarray or a list of (nkpts,nao,nao) ndarray + Density matrix at each k-point. + kpts : (nkpts, 3) ndarray + + Kwargs: + kpts_band : ``(3,)`` ndarray or ``(*,3)`` ndarray + A list of arbitrary "band" k-points at which to evalute the matrix. + with_j : bool + Whether to add the Coulomb matrix into the XC matrix. + + Returns: + exc : XC energy + nelec : number of electrons obtained from the numerical integration + veff : (nkpts, nao, nao) ndarray + or list of veff if the input dm_kpts is a list of DMs + ''' + log = logger.new_logger(cell, verbose) + t0 = log.init_timer() + + xctype = ni._xc_type(xc_code) + nvar = _xc_var_length(xctype) + if xctype == 'HF': + if with_j: + vj = ni.get_j(dm_kpts[0]+dm_kpts[1], hermi, kpts, kpts_band) + veff = cp.stack([vj, vj]) + return lib.tag_array(veff, ecoul=vj.ecoul, exc=0) + else: + veff = cp.zeros_like(dm_kpts) + return lib.tag_array(veff, ecoul=0, exc=0) + + dm_sc = _wannier_transform_dm(ni, dm_kpts, kpts, hermi, xctype) + assert len(dm_sc) == 2 + + cell = ni.cell + mesh = ni.mesh + ngrids = np.prod(mesh) + vol = cell.vol + weight = vol / ngrids + + Gv_bases = _get_Gv_bases(mesh, cell.reciprocal_vectors()) + + rhoG , tauG = _eval_density(ni, dm_sc[0], with_tau=xctype=='MGGA') + rhoGb, tauGb = _eval_density(ni, dm_sc[1], with_tau=xctype=='MGGA') + n_electrons_a = rhoG[0,0,0].real.get() + n_electrons_b = rhoGb[0,0,0].real.get() + n_electrons = np.array([n_electrons_a, n_electrons_b]) + + # dm_sc is represented in primitive bases (by sorted_cell). Its size can be + # much larger than the input dm_kpts. Release its memory if remaining memory + # is insufficient. + if (2*nvar+4)*ngrids*8 > get_avail_mem(): + dm_sc = [None, None] + + density = cp.empty((2, nvar, ngrids)) + _density_to_real_space(rhoG , tauG , Gv_bases, xctype, out=density[0]) + _density_to_real_space(rhoGb, tauGb, Gv_bases, xctype, out=density[1]) + rhoG += rhoGb # rhoG now stores the spin-free density + # release tauG's memory, keep rhoG. rhoG will be used as the workspace for + # _get_coulomb_in_place + tauG = tauGb = None + + # *(1./weight) because rhoR is scaled by weight in _eval_density. If + # computing rhoR with IFFT, the weight factor is not needed. + density *= 1./weight + t0 = log.timer_debug1("density", *t0) + + rho_sf = ndarray(ngrids, dtype=np.float64, buffer=rhoGb) + rho_sf[:] = density[0,0].real + rho_sf[:] += density[1,0].real + + # eval_xc_eff supports float64 only + xc_for_energy, xc_for_fock = ni.eval_xc_eff( + xc_code, density, deriv=1, xctype=xctype, spin=1, inplace=True)[:2] + + xc_for_fock = xc_for_fock.reshape(2, nvar, *mesh) + + xc_energy_sum = float(vec_dot(rho_sf, xc_for_energy).get()) * weight + xc_for_energy = density = rho_sf = None + log.debug("Multigrid exc %s nelec %s", xc_energy_sum, n_electrons) + t0 = log.timer_debug1("eval_xc_eff", *t0) + + xc_for_fock *= weight + coulomb_a, coulomb_b = rhoG, rhoGb + if with_j: + ecoul, coulomb_a = _get_coulomb_in_place(rhoG, Gv_bases) + ecoul = (.5 / vol) * float(ecoul.get()) + log.debug('Multigrid Coulomb energy %s', ecoul) + coulomb_b[:] = coulomb_a + else: + ecoul = None + coulomb_a.fill(0) + coulomb_b.fill(0) + rhoG = rhoGb = None + + # maximum memory usage = (2,nvar,ngrids) float64s + 3(or 4 for MGGA)*ngrids complex128s + # The 3*ngrids complex128s consist of coulomb_a, coulomb_b and the + # workspace required by _vxc_to_reciprocal_space. + vxc_a = _vxc_to_reciprocal_space(xc_for_fock[0], coulomb_a, Gv_bases) + vxc_b = _vxc_to_reciprocal_space(xc_for_fock[1], coulomb_b, Gv_bases) + xc_for_fock = coulomb_a = coulomb_b = None # release memory + + if kpts_band is None: + # dm_sc and the output have the shape shape. Reuse its memory. + veff_a = _eval_xc_mat(ni, vxc_a, out=dm_sc[0]) + veff_b = _eval_xc_mat(ni, vxc_b, out=dm_sc[1]) + veff = cp.stack([ + _inverse_wannier_transform_fock(ni, veff_a, kpts), + _inverse_wannier_transform_fock(ni, veff_b, kpts)]) + + else: + kpts_band = kpts_band.reshape(-1, 3) + kmesh = k2gamma.kpts_to_kmesh(cell, kpts_band) + ni = ni.copy().reset().build(kmesh=kmesh, xctype=xctype) + # ni.build may alter the mesh. vxc was created with mesh different to + # this new mesh. + ni.mesh = mesh + veff_a = _eval_xc_mat(ni, vxc_a) + veff_b = _eval_xc_mat(ni, vxc_b) + veff = cp.stack([ + _inverse_wannier_transform_fock(ni, veff_a, kpts_band), + _inverse_wannier_transform_fock(ni, veff_b, kpts_band)]) + + veff = _format_jks(veff, dm_kpts, kpts_band, kpts) + veff = tag_array(veff, ecoul=ecoul, exc=xc_energy_sum) + t0 = log.timer_debug1("xc matrix", *t0) + return n_electrons, xc_energy_sum, veff + +def get_veff_ip1( + ni, + xc_code, + dm_kpts, + hermi=1, + kpts=None, + with_j=True, + with_pseudo_vloc_orbital_derivative=True, + verbose=None, +): + raise DeprecationWarning + nkpts = len(kpts) if kpts is not None else 1 + grad = ni.energy_nuclear_gradient( + xc_code, dm_kpts, kpts, with_j, with_pseudo_vloc_orbital_derivative) + return grad * nkpts + +class MultiGridNumInt(multigrid_v1.MultiGridNumIntBase): + # Enable analytical Fourier transforms (AFT), which are typically more + # efficient for small unit cells. + enable_aft = True + + # Mesh in the final bucket can be smaller the estimated cell.mesh. + # Allow the overall mesh to be reduced to the one in the final bucket. + # - For energy and nuclear gradients computation, this setting can + # potentially introduce errors. + # - Stain derivatives are very sensitive to the integration mesh. During + # geometry or lattice optimization, self.mesh should be fixed throughout the + # optimization to avoid discontinuities in the computed forces and stress. + allow_mesh_reduction = False + + def __init__(self, cell): + self.reset(cell) + self.mesh = cell.mesh + + def reset(self, cell=None): + if cell is not None: + self.cell = cell + # Preferable to preserve the mesh setting during geometry + # optimization + #self.mesh = cell.mesh + self.bvkcell = None + self.mg_envs = None + self.supmol_img_coords = None + self.aft_buckets = None + self.fft_buckets = None + self.xctype = None + return self + + def build(self, kmesh=None, xctype='MGGA'): + log = logger.new_logger(self.cell) + t0 = log.init_timer() + cell = self.sorted_cell = SortedGTO.from_cell( + self.cell, decontract=True, diffuse_cutoff=1e200) + assert cell.uniq_l_ctr[:,0].max() <= LMAX + + self.xctype = xctype + self.kmesh = kmesh + if kmesh is None: + bvkcell = cell + bvkmesh_Ls = np.zeros((1, 3)) + else: + bvkcell = super_cell(cell, kmesh, wrap_around=True) + # PTR_BAS_COORD was not initialized in the super_cell function + bvkcell._bas[:,PTR_BAS_COORD] = bvkcell._atm[bvkcell._bas[:,ATOM_OF],PTR_COORD] + bvkmesh_Ls = translation_vectors_for_kmesh(cell, kmesh, wrap_around=True) + self.bvkcell = bvkcell + self.bvkmesh_Ls = bvkmesh_Ls + bvk_ncells = len(bvkmesh_Ls) + + Ls = cp.asarray(bvkcell.get_lattice_Ls()) + Ls = Ls[cp.linalg.norm(Ls-.5, axis=1).argsort()] + nimgs = len(Ls) + log.debug1('ft_ao bvk_ncells=%d, nimgs=%d', bvk_ncells, nimgs) + _env = _scale_sp_ctr_coeff(bvkcell) + ao_loc = bvkcell.ao_loc + self.mg_envs = PBCIntEnvVars.new( + cell.natm, cell.nbas, bvk_ncells, nimgs, + bvkcell._atm, bvkcell._bas, _env, ao_loc, Ls) + + a = cell.lattice_vectors() + b = cell.reciprocal_vectors(norm_to=1) + libmgrid.update_lattice_vectors(a.ctypes, b.ctypes) + + # a penalty to encounter for lattice sum + rad = cell.rcut / bvkcell.vol**(1./3) + 1 + surface = 4*np.pi * rad**2 + # Consider two layers near the surface + lattice_sum_factor = surface * 2 + log.debug1('lattice_sum_factor = %g', lattice_sum_factor) + precision = cell.precision / lattice_sum_factor + bas_ij_idx = _non_trivial_bvk_pairs(self, precision) + + # Initialize buckets + is_orth_lattice = abs(a - np.diag(a.diagonal())).max() < 1e-5 + self.aft_buckets = None + self.fft_buckets = None + + mesh = self.mesh + ke_cutoff = self.ke_cutoff = max(0.1, mesh_to_ke(a, mesh).min()) + + init_ke = mesh_to_ke(a, [16]*3).max() + log.debug1('initial ke_cutoff = %g', init_ke) + + if self.enable_aft and is_orth_lattice and nimgs > 30: + # Estimate Ecut for AFT integrals. These can be potentially handled by + # aft_eval_* functions. + # Use ke_cutoff to limit the highest Ecut. This ensures to handle + # shell-pairs even if their Ecuts are higher than ke_cutoff. + aft_Ecut = _aft_Ecut_estimation( + self, bas_ij_idx, ke_cutoff, precision, xctype) + + # aft_final_ke based on system size + final_ke_fac = max(nimgs / 40, 1.) + aft_final_ke = min(init_ke * final_ke_fac, ke_cutoff) + log.debug1('aft init_ke_cutoff = %g, final_ke_cutoff = %g (%.2fx)', + init_ke, aft_final_ke, final_ke_fac) + self.aft_buckets = _partition_ke_for_aft( + self, bas_ij_idx, aft_Ecut, init_ke, aft_final_ke, xctype, log) + + # Filter shell pairs that are not handled by AFT. Using FFT code for + # the remaining pairs. + if self.aft_buckets: + aft_ke_max = self.aft_buckets[-1]['ke_cutoff'] + if aft_ke_max < ke_cutoff: + bas_ij_idx = bas_ij_idx[aft_Ecut > aft_ke_max] + else: + bas_ij_idx = None + + init_ke = aft_final_ke * 1.5 + + if bas_ij_idx is not None and len(bas_ij_idx) > 0: + # bas_ij_idx are the effective paris between cell0 and bvkcell. + # The FFT-MultiGrid code operates on cell0-supmol paris. + # Every bvkcell shell in bas_ij_idx needs to be unpacked to several + # primitive shells in supmol. + self.fft_buckets = _partition_ke_for_fft( + self, bas_ij_idx, init_ke, ke_cutoff, precision, xctype, log) + + nimgs = cell.nimgs + Tx = np.arange(-nimgs[0], nimgs[0]+1, dtype=np.float64) + Ty = np.arange(-nimgs[1], nimgs[1]+1, dtype=np.float64) + Tz = np.arange(-nimgs[2], nimgs[2]+1, dtype=np.float64) + self.supmol_img_coords = cp.asarray(lib.cartesian_prod([Tx, Ty, Tz]).dot(a)) + + # If memory is sufficient, cache tile info for each bucket, including: + # effective tile indices, orbital pairs indices, and corresponding offsets + if len(bas_ij_idx) < 3000000 and np.prod(mesh) < 400**3: + mem = get_avail_mem() + t1 = log.timer_debug1('generating orbital pairs', *t0) + tile_info = _grid_range_to_tile_info_converter(self.fft_buckets, cell) + for bucket in self.fft_buckets: + bucket['grid_tile_cache'] = [ + tile_info(bas_ij_idx, grid_range, bucket['mesh']) + for bas_ij_idx, grid_range in zip( + bucket['bas_ij_cache'], bucket['grid_ranges_cache']) + ] + log.timer_debug1('grid_tile_cache', *t1) + tile_cache_mem = mem - get_avail_mem() + log.debug1('grid_tile_cache memory usage = %.2f MB', tile_cache_mem*1e-6) + + if self.allow_mesh_reduction: + mesh = self.mesh + if self.fft_buckets: + self.mesh = self.fft_buckets[-1]['mesh'] + else: + self.mesh = self.aft_buckets[-1]['mesh'] + log.info('Reduce MultiGrid maximum mesh %s to %s', mesh, self.mesh) + t0 = log.timer_debug1('Initialize buckets', *t0) + return self + + def _ensure_initialized(self, kpts, xctype): + kmesh = k2gamma.kpts_to_kmesh(self.cell, kpts) + if (self.bvkcell is None or + any(self.kmesh != kmesh) or + # LDA and GGA share the same initialization parameters. + # MGGA requires a little bit higher energy cutoff and rcut + (xctype == 'MGGA' and self.xctype != xctype)): + self.build(kmesh, xctype) + return self + + get_nuc = get_nuc + get_pp = get_pp + + get_rho = get_rho + + get_j = get_j_kpts + nr_rks = nr_rks + nr_uks = nr_uks + + def get_vxc(self, cell, grids, xc_code, dm_kpts, spin=0, hermi=1, + kpts=None, kpts_band=None, with_j=False, verbose=None): + fn = self.nr_rks if spin == 0 else self.nr_uks + return fn(cell, grids, xc_code, dm_kpts, spin, hermi=hermi, + kpts=kpts, kpts_band=kpts_band, with_j=with_j, verbose=verbose) + nr_vxc = get_vxc + + eval_xc_eff = numint.NumInt.eval_xc_eff + _init_xcfuns = numint.NumInt._init_xcfuns + + def nr_rks_fxc(self, cell, grids, xc_code, dm0, dms, hermi=0, fxc=None, + kpts=None, with_j=False): + if isinstance(kpts, KPoints): + kpts = kpts.kpts_ibz + assert kpts is None or kpts.ndim == 2 + + assert dms.ndim == 4 + n_dm, nkpts, nao = dms.shape[:3] + + xctype = self._xc_type(xc_code) + nvar = _xc_var_length(xctype) + if xctype == 'HF': + return cp.zeros_like(dms) + + if fxc is None: + spin = 0 + fxc = self.cache_xc_kernel1(cell, grids, xc_code, dm0, spin, kpts, is_rhf=True)[2] + + out = cp.empty_like(dms) + + cell = self.cell + mesh = self.mesh + Gv_bases = _get_Gv_bases(mesh, cell.reciprocal_vectors()) + mem = get_avail_mem() + + for i_dm in range(n_dm): + dm_sc = _wannier_transform_dm(self, dms[i_dm], kpts, hermi, xctype) + rhoG, tauG = _eval_density(self, dm_sc, with_tau=xctype=='MGGA') + rho1 = _density_to_real_space(rhoG, tauG, Gv_bases, xctype) + if dm_sc.nbytes * 10 > mem: + dm_sc = None + + # rho1 has been scaled by (ngrids/cell.vol) + wv = cp.einsum('xg,xyg->yg', rho1, fxc).reshape(nvar, *mesh) + rho1 = None + + if with_j: + coulomb = _get_coulomb_in_place(rhoG, Gv_bases)[1] + else: + coulomb = rhoG + coulomb.fill(0.) + wv = _vxc_to_reciprocal_space(wv, coulomb, Gv_bases, work=tauG) + rhoG = tauG = coulomb = None + + veff = _eval_xc_mat(self, wv, out=dm_sc) + out[i_dm] = _inverse_wannier_transform_fock(self, veff, kpts) + veff = dm_sc = wv = None + + return out.reshape(dms.shape) + + def nr_rks_fxc_st(self, cell, grids, xc_code, dm0, dms, hermi=0, singlet=True, + fxc=None, kpts=None, with_j=False): + if fxc is None: + spin = 1 + fxc = self.cache_xc_kernel1(cell, grids, xc_code, dm0, spin, kpts, + is_rhf=True)[2] + if singlet: + fxc = fxc[0,:,0] + fxc[0,:,1] + else: + fxc = fxc[0,:,0] - fxc[0,:,1] + return self.nr_rks_fxc(cell, grids, xc_code, dm0, dms, hermi, fxc, kpts, with_j) + + def nr_uks_fxc(self, cell, grids, xc_code, dm0, dms, hermi=0, fxc=None, + kpts=None, with_j=False): + if isinstance(kpts, KPoints): + kpts = kpts.kpts_ibz + assert kpts is None or kpts.ndim == 2 + + assert dms.ndim == 5 + n_dm, nkpts, nao = dms.shape[1:4] + + xctype = self._xc_type(xc_code) + nvar = _xc_var_length(xctype) + if xctype == 'HF': + return cp.zeros_like(dms) + + if fxc is None: + spin = 1 + fxc = self.cache_xc_kernel1(cell, grids, xc_code, dm0, spin, kpts, is_rhf=False)[2] + + out = cp.empty_like(dms) + + cell = self.cell + mesh = self.mesh + ngrids = np.prod(mesh) + Gv_bases = _get_Gv_bases(mesh, cell.reciprocal_vectors()) + mem = get_avail_mem() + + for i_dm in range(n_dm): + dm_sc = _wannier_transform_dm(self, dms[:,i_dm], kpts, hermi, xctype) + rhoG , tauG = _eval_density(self, dm_sc[0], with_tau=xctype=='MGGA') + rhoGb, tauGb = _eval_density(self, dm_sc[1], with_tau=xctype=='MGGA') + rho1 = cp.empty((2, nvar, ngrids)) + _density_to_real_space(rhoG , tauG , Gv_bases, xctype, out=rho1[0]) + _density_to_real_space(rhoGb, tauGb, Gv_bases, xctype, out=rho1[1]) + rhoG += rhoGb + tauG = tauGb = None # release memory + if dm_sc.nbytes * 10 > mem: + dm_sc = None + + wv = cp.einsum('axg,axbyg->byg', rho1, fxc).reshape(2, nvar, *mesh) + rho1 = None + + coulomb_a, coulomb_b = rhoG, rhoGb + if with_j: + coulomb_a = _get_coulomb_in_place(rhoG, Gv_bases)[1] + coulomb_b[:] = coulomb_a + else: + coulomb_a.fill(0) + coulomb_b.fill(0) + rhoG = rhoGb = None + + wv_a = _vxc_to_reciprocal_space(wv[0], coulomb_a, Gv_bases) + wv_b = _vxc_to_reciprocal_space(wv[1], coulomb_b, Gv_bases) + coulomb_a = coulomb_b = wv = None + + veff = _eval_xc_mat(self, wv_a, out=dm_sc[0]) + out[0,i_dm] = _inverse_wannier_transform_fock(self, veff, kpts) + veff = _eval_xc_mat(self, wv_b, out=dm_sc[1]) + out[1,i_dm] = _inverse_wannier_transform_fock(self, veff, kpts) + veff = dm_sc = wv_a = wv_b = None + + return out.reshape(dms.shape) + + def cache_xc_kernel1(self, cell, grids, xc_code, dm, spin=0, kpts=None, is_rhf=None): + if isinstance(kpts, KPoints): + kpts = kpts.kpts_ibz + assert kpts is None or kpts.ndim == 2 + + xctype = self._xc_type(xc_code) + nvar = _xc_var_length(xctype) + + dm_sc = _wannier_transform_dm(self, dm, kpts, 1, xctype) + + if is_rhf is None: + is_rhf = len(dm_sc) == 1 + elif is_rhf: + assert len(dm_sc) == 1 + else: + assert spin == 1 + assert len(dm_sc) == 2 + + mesh = self.mesh + ngrids = np.prod(mesh) + Gv_bases = _get_Gv_bases(mesh, cell.reciprocal_vectors()) + + if is_rhf: + rhoG, tauG = _eval_density(self, dm_sc, with_tau=xctype=='MGGA') + if spin == 1: + density = cp.empty((2, nvar, ngrids)) + _density_to_real_space(rhoG, tauG, Gv_bases, xctype, out=density[0]) + density[0] *= .5 + density[1] = density[0] + else: + density = _density_to_real_space(rhoG, tauG, Gv_bases, xctype) + else: + density = cp.empty((2, nvar, ngrids)) + rhoG, tauG = _eval_density(self, dm_sc[0], with_tau=xctype=='MGGA') + _density_to_real_space(rhoG, tauG, Gv_bases, xctype, out=density[0]) + rhoG, tauG = _eval_density(self, dm_sc[1], with_tau=xctype=='MGGA') + _density_to_real_space(rhoG, tauG, Gv_bases, xctype, out=density[1]) + rhoG = tauG = None + + weight = cell.vol / ngrids + density *= 1./weight + vxc, fxc = self.eval_xc_eff(xc_code, density, deriv=2, xctype=xctype, + spin=spin, inplace=True)[1:3] + return None, vxc, fxc + + cache_xc_kernel = NotImplemented + + def energy_nuclear_gradient(self, xc_code, dm_kpts, kpts=None, spin=None, + with_j=False, with_nuc=False): + '''Computes the nuclear gradients of Exc along with additional + contributions from the Coulomb and pseudopotential terms. + + Kwargs: + with_j : + Whether to include the electron-electron Coulomb interactions + with_nuc : + Whether to include the contribution from the local part of + pseudo-potential or electron-nuclear Coulomb interactions + ''' + return self.energy_derivatives(xc_code, dm_kpts, kpts, spin, with_j, + with_nuc)[0] + + def energy_strain_gradient(self, xc_code, dm_kpts, kpts=None, spin=None, + with_j=False, with_nuc=False): + '''Computes the strain derivatives of Exc along with additional + contributions from the Coulomb and pseudopotential terms. + + Kwargs: + with_j : + Whether to include the electron-electron Coulomb interactions + with_nuc : + Whether to include the contribution from the local part of + pseudo-potential or electron-nuclear Coulomb interactions + ''' + return self.energy_derivatives(xc_code, dm_kpts, kpts, spin, with_j, + with_nuc)[1] + + def energy_derivatives(self, xc_code, dm_kpts, kpts=None, spin=None, + with_j=False, with_nuc=False): + '''Computes the nuclear gradients and strain derivatives of Exc + along with additional contributions from the Coulomb and pseudopotential + terms. + + Kwargs: + with_j : + Whether to include the electron-electron Coulomb interactions + with_nuc : + Whether to include the contribution from the local part of + pseudo-potential or electron-nuclear Coulomb interactions + ''' + cell = self.cell + log = logger.new_logger(cell) + t0 = log.init_timer() + + xctype = self._xc_type(xc_code) + nvar = _xc_var_length(xctype) + + dm_sc = _wannier_transform_dm(self, dm_kpts, kpts, 1, xctype) + n_dm = len(dm_sc) + + mesh = self.mesh + ngrids = np.prod(mesh) + vol = cell.vol + weight = vol / ngrids + + Gv_bases = _get_Gv_bases(mesh, cell.reciprocal_vectors()) + + if xctype == 'HF': + if n_dm == 2: + # XC contribution = 0. Only needs to consider Coulomb energy + dm_sc = dm_sc[0] + dm_sc[1] + n_dm = 1 + rhoG, tauG = _eval_density(self, dm_sc) + + sigma = cp.zeros((3, 3)) + vxc = cp.zeros(ngrids) + + elif n_dm == 1: # RHF + assert spin is None or spin == 0 + rhoG, tauG = _eval_density(self, dm_sc, with_tau=xctype=='MGGA') + density = _density_to_real_space(rhoG, tauG, Gv_bases, xctype) + spin = 0 + + density *= 1/weight + rho_sf = ndarray(ngrids, dtype=np.float64, buffer=tauG) + rho_sf[:] = density[0].real + + else: # UHF + assert spin is None or spin == 1 + rhoG , tauG = _eval_density(self, dm_sc[0], with_tau=xctype=='MGGA') + rhoGb, tauGb = _eval_density(self, dm_sc[1], with_tau=xctype=='MGGA') + density = cp.empty((2, nvar, ngrids)) + _density_to_real_space(rhoG , tauG , Gv_bases, xctype, out=density[0]) + _density_to_real_space(rhoGb, tauGb, Gv_bases, xctype, out=density[1]) + rhoG += rhoGb + rhoGb = tauGb = None + spin = 1 + + density *= 1/weight + rho_sf = ndarray(ngrids, dtype=np.float64, buffer=tauG) + rho_sf[:] = density[0,0].real + rho_sf[:] += density[1,0].real + + if xctype != 'HF': + exc, vxc = self.eval_xc_eff( + xc_code, density, deriv=1, xctype=xctype, spin=spin)[:2] + vxc *= weight + + # grid weight response + sigma = rho_sf.dot(exc.ravel()) * weight * cp.eye(3) + + if xctype == 'GGA' or xctype == 'MGGA': + # The response of grids wrt the lattice vectors introduces an + # extra term r_t in the Vxc integral: + # integrate[(\nabla_s rho) Vxc[r] r_t] . + # When applying _vxc_to_reciprocal_space to Vxc, the Fourier + # transform also needs to be performed on r_t, leading to an + # additional term in the integral: + # sigma_st ~ \sum_i (nabla_s rho) Vxc[r]_i \nabla_i r_t + # = (nabla_s rho) Vxc[r]_t + vxc = vxc.reshape(n_dm, nvar, ngrids) + density = density.reshape(n_dm, nvar, ngrids) + for s in range(n_dm): + sigma -= cp.einsum('xg,yg->xy', density[s,1:4], vxc[s,1:4]) + + vxc = vxc.reshape(n_dm, nvar, *mesh) + + density = exc = rho_sf = None + + grad = 0 + coulomb_on_g_mesh = cp.zeros_like(rhoG) + if with_nuc: + if cell._pseudo: + coulomb_on_g_mesh = multigrid_v1.eval_vpplocG(cell, mesh, out=tauG).reshape(mesh) + grad, sigma1 = _pploc_derivatives(cell, mesh, rhoG, Gv_bases) + sigma += sigma1 + else: + grad, sigma1 = _ne_derivatives(cell, mesh, rhoG, Gv_bases) + sigma += sigma1 + ZSI = _get_ZSI(cell, mesh, out=tauG) + vneG = _get_coulomb_in_place(ZSI, Gv_bases)[1] + coulomb_on_g_mesh = vneG.reshape(mesh) + ZSI = vneG = None + + if with_j: + sigma += _coulomb_strain_derivatives(cell, mesh, rhoG, Gv_bases) * 0.5 + # rhoG will be overwritten by _get_coulomb_in_place. Must be called + # after other operations. + ecoul, coulomb_on_g_mesh1 = _get_coulomb_in_place(rhoG, Gv_bases) + ecoul = (.5 / vol) * ecoul + # grid weight response + sigma += ecoul * cp.eye(3) + coulomb_on_g_mesh += coulomb_on_g_mesh1 + coulomb_on_g_mesh1 = None + rhoG = tauG = None + + # Reconstruct fft_buckets if aft_buckets is initialized. In this case, + # self.fft_buckets miss orbital pairs with low Ecut. + if self.aft_buckets is None: + fft_buckets = self.fft_buckets + else: + rad = cell.rcut / self.bvkcell.vol**(1./3) + 1 + surface = 4*np.pi * rad**2 + lattice_sum_factor = surface + log.debug1('lattice_sum_factor = %g', lattice_sum_factor) + precision = cell.precision / lattice_sum_factor + bas_ij_idx = _non_trivial_bvk_pairs(self, precision) + init_ke = mesh_to_ke(cell.lattice_vectors(), [16]*3).max() + fft_buckets = _partition_ke_for_fft( + self, bas_ij_idx, init_ke, self.ke_cutoff, precision, xctype, log) + + if n_dm == 1: # RHF + vxc = _vxc_to_reciprocal_space(vxc[0], coulomb_on_g_mesh, Gv_bases) + grad1, sigma1 = _eval_gradients(self, dm_sc, vxc, fft_buckets) + grad += grad1 + sigma += sigma1 + else: + vxc_a, coulomb_on_g_mesh = coulomb_on_g_mesh, None + vxc_b = vxc_a.copy() + vxc_a = _vxc_to_reciprocal_space(vxc[0], vxc_a, Gv_bases) + vxc_b = _vxc_to_reciprocal_space(vxc[1], vxc_b, Gv_bases) + vxc = None + grad1, sigma1 = _eval_gradients(self, dm_sc[0], vxc_a, fft_buckets) + grad += grad1 + sigma += sigma1 + grad1, sigma1 = _eval_gradients(self, dm_sc[1], vxc_b, fft_buckets) + grad += grad1 + sigma += sigma1 + + t0 = log.timer("xc derivatives", *t0) + return grad.get(), sigma.get() + + to_cpu = NotImplemented + to_gpu = NotImplemented diff --git a/gpu4pyscf/pbc/dft/numint.py b/gpu4pyscf/pbc/dft/numint.py index 2a8b27faf..d60c8bb8d 100644 --- a/gpu4pyscf/pbc/dft/numint.py +++ b/gpu4pyscf/pbc/dft/numint.py @@ -401,7 +401,7 @@ def nr_rks(ni, cell, grids, xc_code, dm_kpts, relativity=0, hermi=1, exc, vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype, spin=0)[:2] den = rho[0] * split_grids.weights - nelec += den.sum() + nelec += den.sum().get()[()] excsum += den.dot(exc).get()[()] wv = vxc * split_grids.weights @@ -481,7 +481,7 @@ def nr_uks(ni, cell, grids, xc_code, dm_kpts, relativity=0, hermi=1, exc, vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype, spin=1)[:2] den = rho[:,0] * split_grids.weights - nelec += den.sum(axis=1) + nelec += den.sum(axis=1).get() excsum += den.dot(exc).sum().get()[()] wv = vxc * split_grids.weights diff --git a/gpu4pyscf/pbc/dft/rks.py b/gpu4pyscf/pbc/dft/rks.py index 02efc0447..00443371b 100644 --- a/gpu4pyscf/pbc/dft/rks.py +++ b/gpu4pyscf/pbc/dft/rks.py @@ -31,7 +31,7 @@ from gpu4pyscf.pbc.df.df import GDF from gpu4pyscf.pbc.dft import gen_grid from gpu4pyscf.pbc.dft import numint -from gpu4pyscf.pbc.dft import multigrid, multigrid_v2 +from gpu4pyscf.pbc.dft import multigrid, multigrid_v3 from gpu4pyscf.lib.cupy_helper import tag_array, get_avail_mem from pyscf import __config__ @@ -69,7 +69,7 @@ def get_veff(ks, cell=None, dm=None, dm_last=None, vhf_last=None, hermi=1, ni = ks._numint hybrid = ni.libxc.is_hybrid_xc(ks.xc) - if isinstance(ni, (multigrid_v2.MultiGridNumInt, multigrid.MultiGridNumInt)): + if isinstance(ni, multigrid.MultiGridNumIntBase): if ks.do_nlc(): raise NotImplementedError(f'MultiGrid for NLC functional {ks.xc} + {ks.nlc}') n, exc, vxc = ni.nr_rks( @@ -268,7 +268,7 @@ def dump_flags(self, verbose=None): def multigrid_numint(self, mesh=None): '''Apply the MultiGrid algorithm for XC numerical integartion''' mf = self.copy() - mf._numint = multigrid_v2.MultiGridNumInt(self.cell) + mf._numint = multigrid_v3.MultiGridNumInt(self.cell) if mesh is not None: mf._numint.mesh = mesh return mf @@ -309,8 +309,7 @@ def gen_response(self, mo_coeff=None, mo_occ=None, dm0 = None with_j = (singlet is None or singlet) and hermi != 2 - j_in_xc = isinstance(ni, (multigrid_v2.MultiGridNumInt, - multigrid.MultiGridNumInt)) + j_in_xc = isinstance(ni, multigrid.MultiGridNumIntBase) def vind(dm1): dm1_shape = dm1.shape diff --git a/gpu4pyscf/pbc/dft/tests/test_multigrid.py b/gpu4pyscf/pbc/dft/tests/test_multigrid.py index 3b5ad9cde..8de79c2c1 100644 --- a/gpu4pyscf/pbc/dft/tests/test_multigrid.py +++ b/gpu4pyscf/pbc/dft/tests/test_multigrid.py @@ -29,6 +29,9 @@ from gpu4pyscf.pbc.dft import multigrid from gpu4pyscf.pbc.tools import ifft, fft +if multigrid.libmgrid is None: + raise unittest.SkipTest('multigrid kernels not compiled') + diamond = ''' C 0. 0. 0. C 0.8917 0.8917 0.8917 diff --git a/gpu4pyscf/pbc/dft/tests/test_multigrid_v2.py b/gpu4pyscf/pbc/dft/tests/test_multigrid_v2.py index 30349cc01..c7548e9d8 100644 --- a/gpu4pyscf/pbc/dft/tests/test_multigrid_v2.py +++ b/gpu4pyscf/pbc/dft/tests/test_multigrid_v2.py @@ -33,6 +33,9 @@ import pytest +if multigrid.libgpbc is None: + raise unittest.SkipTest('multigrid v2 kernels not compiled') + def setUpModule(): global cell_orth, cell_nonorth, cell_he global kpts, dm, dm1 @@ -745,9 +748,7 @@ def test_shell_splitting_for_large_fock_in_imagediff_space_gamma(self): mf = KRKS_gpu(cell, xc = 'pbe', kpts = kpts) mf.conv_tol = 1e-10 - # mf = mf.multigrid_numint() - # assert type(mf._numint) is multigrid.MultiGridNumInt - + # mf._numint = multigrid.MultiGridNumInt(cell) # ref_energy = mf.kernel() # assert mf.converged # ref_gradient = mf.Gradients().kernel() @@ -755,8 +756,7 @@ def test_shell_splitting_for_large_fock_in_imagediff_space_gamma(self): # print(repr(ref_gradient)) with lib.temporary_env(multigrid, get_avail_mem=(lambda **kw: 2**28)): - mf = mf.multigrid_numint() - assert type(mf._numint) is multigrid.MultiGridNumInt + mf._numint = multigrid.MultiGridNumInt(cell) test_energy = mf.kernel() assert mf.converged @@ -801,9 +801,7 @@ def test_shell_splitting_for_large_fock_in_imagediff_space_k(self): mf = KRKS_gpu(cell, xc = 'pbe', kpts = kpts) mf.conv_tol = 1e-10 - # mf = mf.multigrid_numint() - # assert type(mf._numint) is multigrid.MultiGridNumInt - + # mf._numint = multigrid.MultiGridNumInt(cell) # ref_energy = mf.kernel() # assert mf.converged # ref_gradient = mf.Gradients().kernel() @@ -811,8 +809,7 @@ def test_shell_splitting_for_large_fock_in_imagediff_space_k(self): # print(repr(ref_gradient)) with lib.temporary_env(multigrid, get_avail_mem=(lambda **kw: 2**28)): - mf = mf.multigrid_numint() - assert type(mf._numint) is multigrid.MultiGridNumInt + mf._numint = multigrid.MultiGridNumInt(cell) test_energy = mf.kernel() assert mf.converged @@ -855,9 +852,7 @@ def test_shell_splitting_for_large_fock_in_imagediff_space_unrestricted(self): mf = KUKS_gpu(cell, xc = 'pbe', kpts = kpts) mf.conv_tol = 1e-10 - # mf = mf.multigrid_numint() - # assert type(mf._numint) is multigrid.MultiGridNumInt - + # mf._numint = multigrid.MultiGridNumInt(cell) # ref_energy = mf.kernel() # assert mf.converged # ref_gradient = mf.Gradients().kernel() @@ -865,8 +860,7 @@ def test_shell_splitting_for_large_fock_in_imagediff_space_unrestricted(self): # print(repr(ref_gradient)) with lib.temporary_env(multigrid, get_avail_mem=(lambda **kw: 2**25)): - mf = mf.multigrid_numint() - assert type(mf._numint) is multigrid.MultiGridNumInt + mf._numint = multigrid.MultiGridNumInt(cell) test_energy = mf.kernel() assert mf.converged diff --git a/gpu4pyscf/pbc/dft/tests/test_pbc_dftu.py b/gpu4pyscf/pbc/dft/tests/test_pbc_dftu.py index 61561623f..92903581d 100644 --- a/gpu4pyscf/pbc/dft/tests/test_pbc_dftu.py +++ b/gpu4pyscf/pbc/dft/tests/test_pbc_dftu.py @@ -50,7 +50,7 @@ def test_KRKSpU(self): minao_ref='gth-szv') mf.conv_tol = 1e-10 e1 = mf.kernel() - self.assertAlmostEqual(e1, -10.694460059491741, 8) + self.assertAlmostEqual(e1, -10.694460059491741, 6) def test_KUKSpU(self): kmesh = [1, 1, 1] @@ -61,7 +61,7 @@ def test_KUKSpU(self): minao_ref='gth-szv') mf.conv_tol = 1e-10 e1 = mf.kernel() - self.assertAlmostEqual(e1, -10.1793267284188, 8) + self.assertAlmostEqual(e1, -10.1793267284188, 6) def test_get_veff(self): kmesh = [2, 1, 1] diff --git a/gpu4pyscf/pbc/dft/tests/test_pbc_krkspu.py b/gpu4pyscf/pbc/dft/tests/test_pbc_krkspu.py index 38c76f9d4..1f3e1b161 100644 --- a/gpu4pyscf/pbc/dft/tests/test_pbc_krkspu.py +++ b/gpu4pyscf/pbc/dft/tests/test_pbc_krkspu.py @@ -55,7 +55,7 @@ def test_KRKSpU(self): minao_ref='gth-szv') mf.conv_tol = 1e-10 e1 = mf.kernel() - self.assertAlmostEqual(e1, -10.694460059491741, 8) + self.assertAlmostEqual(e1, -10.694460059491741, 6) def test_get_veff(self): kmesh = [2, 1, 1] @@ -93,7 +93,7 @@ def test_KRKSpU_linear_response(self): minao_ref='gth-szv') mf.conv_tol = 1e-10 e_tot = mf.kernel() - self.assertAlmostEqual(e_tot, -10.6191452297714, 8) + self.assertAlmostEqual(e_tot, -10.6191452297714, 6) uresp = krkspu.linear_response_u(mf, (0.03, 0.08)) self.assertAlmostEqual(uresp, 6.279179, 2) diff --git a/gpu4pyscf/pbc/dft/tests/test_pbc_kukspu.py b/gpu4pyscf/pbc/dft/tests/test_pbc_kukspu.py index 2e78663f8..5731715cd 100644 --- a/gpu4pyscf/pbc/dft/tests/test_pbc_kukspu.py +++ b/gpu4pyscf/pbc/dft/tests/test_pbc_kukspu.py @@ -55,7 +55,7 @@ def test_KUKSpU(self): minao_ref='gth-szv') mf.conv_tol = 1e-10 e1 = mf.kernel() - self.assertAlmostEqual(e1, -10.694460059491741, 8) + self.assertAlmostEqual(e1, -10.694460059491741, 6) def test_get_veff(self): kmesh = [2, 1, 1] @@ -93,7 +93,7 @@ def test_KUKSpU_linear_response(self): minao_ref='gth-szv') mf.conv_tol = 1e-10 e_tot = mf.kernel() - self.assertAlmostEqual(e_tot, -10.6191452297714, 8) + self.assertAlmostEqual(e_tot, -10.6191452297714, 6) uresp = kukspu.linear_response_u(mf, (0.03, 0.08)) self.assertAlmostEqual(uresp, 6.279179, 2) diff --git a/gpu4pyscf/pbc/dft/tests/test_pbc_numint.py b/gpu4pyscf/pbc/dft/tests/test_pbc_numint.py index 226baa630..02f4bc0e4 100644 --- a/gpu4pyscf/pbc/dft/tests/test_pbc_numint.py +++ b/gpu4pyscf/pbc/dft/tests/test_pbc_numint.py @@ -243,7 +243,7 @@ def test_nr_uks(self): ni = numint.NumInt() ne, exc, vmat = ni.nr_uks(cell, grids, 'lda', dms, hermi=1, kpts=kpts[0]) ref = ni.to_cpu().nr_uks(cell, grids.to_cpu(), 'lda', dms.get(), hermi=1, kpt=kpts[0]) - self.assertAlmostEqual(abs(ne.get() - ref[0]).max(), 0, 9) + self.assertAlmostEqual(abs(ne - ref[0]).max(), 0, 9) self.assertAlmostEqual(float(exc), ref[1], 9) self.assertAlmostEqual(abs(vmat.get() - ref[2]).max(), 0, 9) @@ -283,19 +283,19 @@ def test_knumint_nr_uks(self): ni = numint.KNumInt() ne, exc, vmat = ni.nr_uks(cell, grids, 'm06', dms, hermi=1, kpts=kpts) ref = ni.to_cpu().nr_uks(cell, grids.to_cpu(), 'm06', dms.get(), hermi=1, kpts=kpts) - self.assertAlmostEqual(abs(ne.get() - ref[0]).max(), 0, 9) + self.assertAlmostEqual(abs(ne - ref[0]).max(), 0, 9) self.assertAlmostEqual(exc, ref[1], 9) self.assertAlmostEqual(abs(vmat.get() - ref[2]).max(), 0, 9) ne, exc, vmat = ni.nr_uks(cell, grids, 'blyp', dms, hermi=1, kpts=kpts) ref = ni.to_cpu().nr_uks(cell, grids.to_cpu(), 'blyp', dms.get(), hermi=1, kpts=kpts) - self.assertAlmostEqual(abs(ne.get() - ref[0]).max(), 0, 9) + self.assertAlmostEqual(abs(ne - ref[0]).max(), 0, 9) self.assertAlmostEqual(exc, ref[1], 9) self.assertAlmostEqual(abs(vmat.get() - ref[2]).max(), 0, 9) ne, exc, vmat = ni.nr_uks(cell, grids, 'lda', dms, hermi=1, kpts=kpts) ref = ni.to_cpu().nr_uks(cell, grids.to_cpu(), 'lda', dms.get(), hermi=1, kpts=kpts) - self.assertAlmostEqual(abs(ne.get() - ref[0]).max(), 0, 9) + self.assertAlmostEqual(abs(ne - ref[0]).max(), 0, 9) self.assertAlmostEqual(exc, ref[1], 9) self.assertAlmostEqual(abs(vmat.get() - ref[2]).max(), 0, 9) @@ -418,14 +418,14 @@ def test_nr_uks_division_mode(self): ni = numint.NumInt() ne, exc, vmat = ni.nr_uks(cell, grids, 'lda', dms[:,0], hermi=1, kpts=kpts[0]) ref = ni.to_cpu().nr_uks(cell, grids.to_cpu(), 'lda', dms[:,0].get(), hermi=1, kpt=kpts[0]) - self.assertAlmostEqual(abs(ne.get() - ref[0]).max(), 0, 9) + self.assertAlmostEqual(abs(ne - ref[0]).max(), 0, 9) self.assertAlmostEqual(float(exc), ref[1], 9) self.assertAlmostEqual(abs(vmat.get() - ref[2]).max(), 0, 9) ni = numint.KNumInt() ne, exc, vmat = ni.nr_uks(cell, grids, 'm06', dms, hermi=1, kpts=kpts) ref = ni.to_cpu().nr_uks(cell, grids.to_cpu(), 'm06', dms.get(), hermi=1, kpts=kpts) - self.assertAlmostEqual(abs(ne.get() - ref[0]).max(), 0, 9) + self.assertAlmostEqual(abs(ne - ref[0]).max(), 0, 9) self.assertAlmostEqual(exc, ref[1], 9) self.assertAlmostEqual(abs(vmat.get() - ref[2]).max(), 0, 9) diff --git a/gpu4pyscf/pbc/dft/tests/test_pbc_rks.py b/gpu4pyscf/pbc/dft/tests/test_pbc_rks.py index 2b3b12824..cbecb3568 100644 --- a/gpu4pyscf/pbc/dft/tests/test_pbc_rks.py +++ b/gpu4pyscf/pbc/dft/tests/test_pbc_rks.py @@ -67,8 +67,8 @@ def test_lda_fft(self): kpts_band = np.random.random((2,3)) e0, c0 = mf_ref.get_bands(kpts_band) e1, c1 = mf.get_bands(kpts_band) - self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, 6) - self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, 6) + self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, delta=1e-5) + self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, delta=1e-5) def test_gga_fft(self): cell = self.cell @@ -81,8 +81,8 @@ def test_gga_fft(self): kpts_band = np.random.random((2,3)) e0, c0 = mf_ref.get_bands(kpts_band) e1, c1 = mf.get_bands(kpts_band) - self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, 6) - self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, 6) + self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, delta=1e-5) + self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, delta=1e-5) def test_rsh_fft(self): cell = self.cell @@ -95,8 +95,8 @@ def test_rsh_fft(self): kpts_band = np.random.random((2,3)) e0, c0 = mf_ref.get_bands(kpts_band) e1, c1 = mf.get_bands(kpts_band) - self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, 6) - self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, 6) + self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, delta=1e-5) + self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, delta=1e-5) def test_kpts_mgga(self): cell = self.cell @@ -109,8 +109,8 @@ def test_kpts_mgga(self): kpts_band = np.random.random((2,3)) e0, c0 = mf_ref.get_bands(kpts_band) e1, c1 = mf.get_bands(kpts_band) - self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, 6) - self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, 6) + self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, delta=1e-5) + self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, delta=1e-5) def test_lda_gdf(self): from pyscf.pbc.df.df import _load3c @@ -183,15 +183,15 @@ def test_lda_fft_with_kpt(self): mf.time_reversal_symmetry = False mf.run() mf_ref = mf.to_cpu().run() - self.assertAlmostEqual(mf.e_tot, mf_ref.e_tot, 7) + self.assertAlmostEqual(mf.e_tot, mf_ref.e_tot, delta=1e-6) # test bands np.random.seed(1) kpts_band = np.random.random((2,3)) e0, c0 = mf_ref.get_bands(kpts_band) e1, c1 = mf.get_bands(kpts_band) - self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, 6) - self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, 6) + self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, delta=1e-5) + self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, delta=1e-5) def test_gga_fft_with_kpt(self): cell = self.cell @@ -201,15 +201,15 @@ def test_gga_fft_with_kpt(self): mf.time_reversal_symmetry = False mf.run(conv_tol=1e-10) mf_ref = mf.to_cpu().run(conv_tol=1e-10) - self.assertAlmostEqual(mf.e_tot, mf_ref.e_tot, 7) + self.assertAlmostEqual(mf.e_tot, mf_ref.e_tot, delta=1e-6) # test bands np.random.seed(1) kpts_band = np.random.random((2,3)) e0, c0 = mf_ref.get_bands(kpts_band) e1, c1 = mf.get_bands(kpts_band) - self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, 6) - self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, 6) + self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, delta=1e-5) + self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, delta=1e-5) def test_rsh_fft_with_kpt(self): cell = self.cell @@ -219,15 +219,15 @@ def test_rsh_fft_with_kpt(self): mf.time_reversal_symmetry = False mf.run(conv_tol=1e-10) mf_ref = mf.to_cpu().run(conv_tol=1e-10) - self.assertAlmostEqual(mf.e_tot, mf_ref.e_tot, 7) + self.assertAlmostEqual(mf.e_tot, mf_ref.e_tot, delta=1e-6) # test bands np.random.seed(1) kpts_band = np.random.random((2,3)) e0, c0 = mf_ref.get_bands(kpts_band) e1, c1 = mf.get_bands(kpts_band) - self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, 6) - self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, 6) + self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, delta=1e-5) + self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, delta=1e-5) def test_kpts_lda_fft(self): cell = self.cell @@ -242,8 +242,8 @@ def test_kpts_lda_fft(self): kpts_band = np.random.random((2,3)) e0, c0 = mf_ref.get_bands(kpts_band) e1, c1 = kmf.get_bands(kpts_band) - self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, delta=1e-6) - self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, delta=1e-6) + self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, delta=1e-5) + self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, delta=1e-5) def test_kpts_gga_fft(self): cell = self.cell @@ -501,6 +501,7 @@ def test_cambl3yp_krks_rsjk(self): #ref = cell.KRKS(xc='camb3lyp', kpts=kpts).run() #self.assertAlmostEqual(mf.e_tot, ref.e_tot, 8) + @unittest.skip('Only support Monkhorst-Pack kmesh') def test_unpaired_kpts(self): kpts = cell.make_kpts([1,1,5])[:3] kmf = cell.KRKS(xc='pbe', kpts=kpts).to_gpu() diff --git a/gpu4pyscf/pbc/dft/tests/test_pbc_uks.py b/gpu4pyscf/pbc/dft/tests/test_pbc_uks.py index d6232e2c3..298329e6f 100644 --- a/gpu4pyscf/pbc/dft/tests/test_pbc_uks.py +++ b/gpu4pyscf/pbc/dft/tests/test_pbc_uks.py @@ -52,8 +52,8 @@ def test_lda_fft(self): kpts_band = np.random.random((2,3)) e0, c0 = mf_ref.get_bands(kpts_band) e1, c1 = mf.get_bands(kpts_band) - self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, 6) - self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, 6) + self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, delta=1e-5) + self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, delta=1e-5) def test_gga_fft(self): mf = pbcdft.UKS(cell, xc='pbe0').run(conv_tol=1e-9) @@ -65,8 +65,8 @@ def test_gga_fft(self): kpts_band = np.random.random((2,3)) e0, c0 = mf_ref.get_bands(kpts_band) e1, c1 = mf.get_bands(kpts_band) - self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, 6) - self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, 6) + self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, delta=1e-5) + self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, delta=1e-5) def test_rsh_fft(self): mf = pbcdft.UKS(cell, xc='camb3lyp').run(conv_tol=1e-9) @@ -79,8 +79,8 @@ def test_rsh_fft(self): kpts_band = np.random.random((2,3)) e0, c0 = mf_ref.get_bands(kpts_band) e1, c1 = mf.get_bands(kpts_band) - self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, 6) - self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, 6) + self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, delta=1e-5) + self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, delta=1e-5) def test_mgga_fft(self): mf = pbcdft.UKS(cell, xc='tpss').run(conv_tol=1e-9) @@ -92,8 +92,8 @@ def test_mgga_fft(self): kpts_band = np.random.random((2,3)) e0, c0 = mf_ref.get_bands(kpts_band) e1, c1 = mf.get_bands(kpts_band) - self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, 6) - self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, 6) + self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, delta=1e-5) + self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, delta=1e-5) def test_rsh_gdf(self): mf = pbcdft.UKS(cell, xc='camb3lyp').density_fit().run() @@ -107,15 +107,15 @@ def test_lda_fft_with_kpt(self): mf.time_reversal_symmetry = False mf.run(conv_tol=1e-10) mf_ref = mf.to_cpu().run() - self.assertAlmostEqual(mf.e_tot, mf_ref.e_tot, 7) + self.assertAlmostEqual(mf.e_tot, mf_ref.e_tot, delta=1e-6) # test bands np.random.seed(1) kpts_band = np.random.random((2,3)) e0, c0 = mf_ref.get_bands(kpts_band) e1, c1 = mf.get_bands(kpts_band) - self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, 6) - self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, 6) + self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, delta=1e-5) + self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, delta=1e-5) def test_gga_fft_with_kpt(self): np.random.seed(1) @@ -124,15 +124,15 @@ def test_gga_fft_with_kpt(self): mf.time_reversal_symmetry = False mf.run(conv_tol=1e-10) mf_ref = mf.to_cpu().run(conv_tol=1e-10) - self.assertAlmostEqual(mf.e_tot, mf_ref.e_tot, 7) + self.assertAlmostEqual(mf.e_tot, mf_ref.e_tot, delta=1e-6) # test bands np.random.seed(1) kpts_band = np.random.random((2,3)) e0, c0 = mf_ref.get_bands(kpts_band) e1, c1 = mf.get_bands(kpts_band) - self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, 5) - self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, 5) + self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, delta=1e-5) + self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, delta=1e-5) def test_rsh_fft_with_kpt(self): np.random.seed(1) @@ -141,15 +141,15 @@ def test_rsh_fft_with_kpt(self): mf.time_reversal_symmetry = False mf.run(conv_tol=1e-10) mf_ref = mf.to_cpu().run(conv_tol=1e-10) - self.assertAlmostEqual(mf.e_tot, mf_ref.e_tot, 7) + self.assertAlmostEqual(mf.e_tot, mf_ref.e_tot, delta=1e-6) # test bands np.random.seed(1) kpts_band = np.random.random((2,3)) e0, c0 = mf_ref.get_bands(kpts_band) e1, c1 = mf.get_bands(kpts_band) - self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, 6) - self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, 6) + self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, delta=1e-5) + self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, delta=1e-5) def test_kpts_lda_fft(self): nk = [2, 1, 1] @@ -163,22 +163,22 @@ def test_kpts_lda_fft(self): kpts_band = np.random.random((2,3)) e0, c0 = mf_ref.get_bands(kpts_band) e1, c1 = kmf.get_bands(kpts_band) - self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, delta=1e-6) - self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, delta=1e-6) + self.assertAlmostEqual(abs(e1[0].get() - e0[0]).max(), 0, delta=1e-5) + self.assertAlmostEqual(abs(e1[1].get() - e0[1]).max(), 0, delta=1e-5) def test_kpts_gga_fft(self): nk = [2, 1, 1] kpts = cell.make_kpts(nk) kmf = pbcdft.KUKS(cell, xc='pbe0', kpts=kpts).run(conv_tol=1e-10) mf_ref = kmf.to_cpu().run() - self.assertAlmostEqual(kmf.e_tot, mf_ref.e_tot, 7) + self.assertAlmostEqual(kmf.e_tot, mf_ref.e_tot, delta=1e-6) def test_kpts_rsh_fft(self): nk = [2, 1, 1] kpts = cell.make_kpts(nk) kmf = pbcdft.KUKS(cell, xc='camb3lyp', kpts=kpts).run(conv_tol=1e-9) mf_ref = kmf.to_cpu().run() - self.assertAlmostEqual(kmf.e_tot, mf_ref.e_tot, 7) + self.assertAlmostEqual(kmf.e_tot, mf_ref.e_tot, delta=1e-6) def test_kpts_gga_gdf(self): from gpu4pyscf.pbc.df.df import GDF @@ -213,7 +213,7 @@ def test_kpts_gga_gdf(self): self.assertAlmostEqual(kmf.e_tot, mf_ref.e_tot, 7) def test_rsjk(self): - from gpu4pyscf.pbc.dft.multigrid_v2 import MultiGridNumInt + from gpu4pyscf.pbc.dft.multigrid_v3 import MultiGridNumInt L = 4. cell = pbcgto.Cell() cell.a = np.eye(3)*L @@ -254,7 +254,7 @@ def test_rsjk(self): self.assertAlmostEqual(mf.e_tot, -0.442283740471709, 8) def test_rsjk_kuks(self): - from gpu4pyscf.pbc.dft.multigrid_v2 import MultiGridNumInt + from gpu4pyscf.pbc.dft.multigrid_v3 import MultiGridNumInt L = 4. cell = pbcgto.Cell() cell.a = np.eye(3)*L diff --git a/gpu4pyscf/pbc/dft/uks.py b/gpu4pyscf/pbc/dft/uks.py index d2d479eab..28fa6f37d 100644 --- a/gpu4pyscf/pbc/dft/uks.py +++ b/gpu4pyscf/pbc/dft/uks.py @@ -30,7 +30,7 @@ from gpu4pyscf.lib.cupy_helper import tag_array, get_avail_mem from gpu4pyscf.dft import uks as mol_uks from gpu4pyscf.pbc.dft import rks, krks -from gpu4pyscf.pbc.dft import multigrid, multigrid_v2 +from gpu4pyscf.pbc.dft import multigrid def get_veff(ks, cell=None, dm=None, dm_last=None, vhf_last=None, hermi=1, @@ -55,7 +55,7 @@ def get_veff(ks, cell=None, dm=None, dm_last=None, vhf_last=None, hermi=1, ni = ks._numint hybrid = ni.libxc.is_hybrid_xc(ks.xc) - if isinstance(ni, (multigrid_v2.MultiGridNumInt, multigrid.MultiGridNumInt)): + if isinstance(ni, multigrid.MultiGridNumIntBase): if ks.do_nlc(): raise NotImplementedError(f'MultiGrid for NLC functional {ks.xc} + {ks.nlc}') n, exc, vxc = ni.nr_uks( @@ -160,8 +160,7 @@ def gen_response(self, mo_coeff=None, mo_occ=None, dm0 = None with_j = with_j and hermi != 2 - j_in_xc = isinstance(ni, (multigrid_v2.MultiGridNumInt, - multigrid.MultiGridNumInt)) + j_in_xc = isinstance(ni, multigrid.MultiGridNumIntBase) def vind(dm1): dm1_shape = dm1.shape diff --git a/gpu4pyscf/pbc/grad/krhf.py b/gpu4pyscf/pbc/grad/krhf.py index 19063e58a..64b6db3a7 100644 --- a/gpu4pyscf/pbc/grad/krhf.py +++ b/gpu4pyscf/pbc/grad/krhf.py @@ -36,7 +36,7 @@ from gpu4pyscf.pbc.grad.pp import vppnl_nuc_grad from gpu4pyscf.pbc.grad.rhf import contract_h1e_dm, jk_energy_per_atom from gpu4pyscf.pbc.grad import rhf as pbchf_grad -from gpu4pyscf.pbc.dft import multigrid, multigrid_v2 +from gpu4pyscf.pbc.dft import multigrid __all__ = ['Gradients'] @@ -71,22 +71,10 @@ def grad_elec(mf_grad, mo_energy=None, mo_coeff=None, mo_occ=None): t1 = log.timer('gradients of 2e part', *t0) ni = mf._numint - if isinstance(ni, multigrid.MultiGridNumInt): - raise NotImplementedError( - "Gradient with kpts not implemented with multigrid.MultiGridNumInt. " - "Please use the default KNumInt or multigrid_v2.MultiGridNumInt instead.") - elif isinstance(ni, multigrid_v2.MultiGridNumInt): - # Attention: The orbital derivative of vpploc term is in multigrid_v2.get_veff_ip1() function. - rho_g = multigrid_v2.evaluate_density_on_g_mesh(ni, dm0, kpts) - rho_g = rho_g[0,0] - if cell._pseudo: - dh1e = multigrid.eval_vpplocG_SI_gradient(cell, ni.mesh, rho_g) * nkpts - else: - dh1e = multigrid.eval_nucG_SI_gradient(cell, ni.mesh, rho_g) * nkpts - - dh1e = dh1e.get() + if isinstance(ni, multigrid.MultiGridNumIntBase): + # Vne or pploc contribution is evaluated in energy_ee dh1e_kin = int1e.int1e_ipkin(cell, kpts) - dh1e -= contract_h1e_dm(cell, dh1e_kin, dm0, hermi=1) + dh1e = -contract_h1e_dm(cell, dh1e_kin, dm0, hermi=1) else: hcore_deriv = mf_grad.hcore_generator(cell, kpts) dh1e = cp.empty([natm, 3]) @@ -270,13 +258,11 @@ def energy_ee(self, dm, kpts): # When J is evaluated using mf.j_engine or mf.rsjk, it is identical to # the J from MultiGridNumInt. The contribution from J matrix can be # efficiently evaluated using the MultiGridNumInt integrator. - j_in_xc = ni is not None and isinstance(ni, multigrid_v2.MultiGridNumInt) + j_in_xc = ni is not None and isinstance(ni, multigrid.MultiGridNumIntBase) if j_in_xc: j_factor = 0 - de = multigrid_v2.get_veff_ip1( - ni, 'HF', dm, kpts=kpts, with_j=j_in_xc, - with_pseudo_vloc_orbital_derivative=True).get() - de /= len(kpts) + de = ni.energy_nuclear_gradient( + 'HF', dm, kpts=kpts, spin=0, with_j=j_in_xc, with_nuc=True) else: j_factor = 1 de = 0 diff --git a/gpu4pyscf/pbc/grad/krhf_stress.py b/gpu4pyscf/pbc/grad/krhf_stress.py index eeac7eaa6..975016cb9 100644 --- a/gpu4pyscf/pbc/grad/krhf_stress.py +++ b/gpu4pyscf/pbc/grad/krhf_stress.py @@ -74,6 +74,8 @@ def kernel(mf_grad): sigma -= int1e.ovlp_strain_deriv(cell, dme0, kpts) sigma += int1e.kin_strain_deriv(cell, dm0, kpts) sigma += get_nuc(mf_grad, cell, dm0, kpts) + if cell._pseudo: + sigma += _get_pp_nonloc_strain_derivatives(cell, cell.mesh, dm0, kpts) t0 = log.timer_debug1('hcore derivatives', *t0) sigma += get_veff(mf_grad, cell, dm0, kpts) @@ -171,7 +173,6 @@ def partial_dot(bra, ket): vpplocR = pbctools.ifft(vpplocG_0, mesh).real Ene = contract('xyg,g->xy', rho1, vpplocR).real.get() Ene += contract('g,xyg->xy', rhoG.conj(), vpplocG_1).real.get() * (1./ngrids) - Ene += _get_pp_nonloc_strain_derivatives(cell, mesh, dm, kpts) else: Gv = cell.get_Gv(mesh) coulG_0, coulG_1 = _get_coulG_strain_derivatives(cell, Gv) diff --git a/gpu4pyscf/pbc/grad/krks.py b/gpu4pyscf/pbc/grad/krks.py index ab94973d2..5277b038b 100644 --- a/gpu4pyscf/pbc/grad/krks.py +++ b/gpu4pyscf/pbc/grad/krks.py @@ -24,8 +24,7 @@ from gpu4pyscf.pbc.grad import krhf as krhf_grad from gpu4pyscf.grad import rks as rks_grad from gpu4pyscf.lib.cupy_helper import contract -from gpu4pyscf.pbc.dft import BeckeGrids -from gpu4pyscf.pbc.dft import multigrid, multigrid_v2 +from gpu4pyscf.pbc.dft import multigrid, BeckeGrids __all__ = ['Gradients'] @@ -45,25 +44,13 @@ def energy_ee(ks_grad, dm, kpts): omega, k_lr, k_sr = ni.rsh_and_hybrid_coeff(mf.xc) j_factor = 1 - if isinstance(ni, multigrid.MultiGridNumInt): - raise NotImplementedError( - "Gradient with kpts not implemented with multigrid.MultiGridNumInt. " - "Please use the default KNumInt or multigrid_v2.MultiGridNumInt instead.") - - if isinstance(ni, multigrid_v2.MultiGridNumInt): + if isinstance(ni, multigrid.MultiGridNumIntBase): # Note the j_in_xc treatment here slightly differs from KRHF. In KRHF, # if GDF is enabled, J is evaluated with GDF CDERI. However, in KRKS, # J is evaluated using MultiGridNumInt whenever applicable. See also # the implementation in pbc.scf.krks - if kpts is None: - nkpts = 1 - else: - nkpts = len(kpts) - exc = multigrid_v2.get_veff_ip1( - ni, mf.xc, dm, with_j=True, with_pseudo_vloc_orbital_derivative=True, kpts=kpts).get() - # exc of multigrid_v2 is the full response of dE/dX. However, - # get_veff in grad_elec evaluates the contraction Tr(dm, ). - exc /= nkpts + exc = ni.energy_nuclear_gradient( + mf.xc, dm, kpts=kpts, spin=0, with_j=True, with_nuc=True) j_factor = 0 else: if ks_grad.grids is not None: diff --git a/gpu4pyscf/pbc/grad/krks_stress.py b/gpu4pyscf/pbc/grad/krks_stress.py index d3b10c470..77305ba01 100644 --- a/gpu4pyscf/pbc/grad/krks_stress.py +++ b/gpu4pyscf/pbc/grad/krks_stress.py @@ -56,7 +56,7 @@ from gpu4pyscf.pbc.df.aft import _get_ZSI from gpu4pyscf.pbc.dft.numint import KNumInt, eval_ao_kpts, _GTOvalOpt from gpu4pyscf.pbc.dft.krkspu import _set_U, _make_minao_lo, reference_mol -from gpu4pyscf.pbc.dft.multigrid_v2 import _rks_exc_strain_deriv, MultiGridNumInt +from gpu4pyscf.pbc.dft import multigrid, BeckeGrids from gpu4pyscf.pbc.grad import krks as krks_grad from gpu4pyscf.pbc.gto import int1e from gpu4pyscf.pbc.gto.cell import get_Gv @@ -99,13 +99,17 @@ def get_veff(mf_grad, cell, dm, kpts, with_j=False, with_nuc=False): ni = mf._numint is_hybrid = ni.libxc.is_hybrid_xc(mf.xc) + if isinstance(mf.grids, BeckeGrids): + raise NotImplementedError('gradients for BeckeGrids not supported') + j_factor = 1 if with_j else 0 if is_hybrid and with_rsjk is not None: with_j = False # TODO: with_nuc should be disabled for all-electron calculations - if isinstance(ni, MultiGridNumInt): - sigma = _rks_exc_strain_deriv(ni, mf.xc, dm, kpts, with_j, with_nuc) + if isinstance(ni, multigrid.MultiGridNumIntBase): + sigma = ni.energy_strain_gradient(mf.xc, dm, kpts, spin=0, + with_j=with_j, with_nuc=with_nuc) elif isinstance(ni, KNumInt): sigma = get_vxc(mf_grad, cell, dm, kpts, with_j, with_nuc) else: @@ -277,9 +281,9 @@ def _contract_coulomb_and_nuc(cell, mesh, dm, kpts, rho0, rho1, grids, with_j, w Gv = get_Gv(cell, mesh) coulG_0, coulG_1 = _get_coulG_strain_derivatives(cell, Gv) rhoG = pbctools.fft(rho0, mesh) + weight_0, weight_1 = _get_weight_strain_derivatives(cell, grids) out = 0 if with_j: - weight_0, weight_1 = _get_weight_strain_derivatives(cell, grids) vR = pbctools.ifft(rhoG * coulG_0, mesh) EJ = cp.einsum('xyg,g->xy', rho1, vR).real.get() * weight_0 * 2 EJ += cp.einsum('g,g->', rho0, vR).real.get() * weight_1 @@ -292,9 +296,7 @@ def _contract_coulomb_and_nuc(cell, mesh, dm, kpts, rho0, rho1, grids, with_j, w vpplocR = pbctools.ifft(vpplocG_0, mesh).real Ene = cp.einsum('xyg,g->xy', rho1, vpplocR).real.get() Ene += cp.einsum('g,xyg->xy', rhoG.conj(), vpplocG_1).real.get() * (1./ngrids) - Ene += _get_pp_nonloc_strain_derivatives(cell, mesh, dm, kpts) else: - coulG_0, coulG_1 = _get_coulG_strain_derivatives(cell, Gv) # SI corresponds to Fourier components of the fractional atomic # positions within the cell. It does not respond to the strain # transformation @@ -339,22 +341,10 @@ def kernel(mf_grad): sigma = ewald(cell) kpts = mf.kpts - kmesh = kpts_to_kmesh(cell, kpts, bound_by_supmol=True) sigma -= int1e.ovlp_strain_deriv(cell, dme0, kpts) - - scaled_kpts = kpts.dot(cell.lattice_vectors().T) - nkpts = len(kpts) - disp = 1e-5 - for x in range(3): - for y in range(3): - cell1, cell2 = _finite_diff_cells(cell, x, y, disp) - kpts1 = scaled_kpts.dot(cell1.reciprocal_vectors(norm_to=1)) - kpts2 = scaled_kpts.dot(cell2.reciprocal_vectors(norm_to=1)) - t1 = int1e.int1e_kin(cell1, kpts1, kmesh) - t2 = int1e.int1e_kin(cell2, kpts2, kmesh) - t1 = cp.einsum('kij,kji->', t1, dm0).real - t2 = cp.einsum('kij,kji->', t2, dm0).real - sigma[x,y] += (t1 - t2).get() / (2*disp) / nkpts + sigma += int1e.kin_strain_deriv(cell, dm0, kpts) + if cell._pseudo: + sigma += _get_pp_nonloc_strain_derivatives(cell, cell.mesh, dm0, kpts) t0 = log.timer_debug1('hcore derivatives', *t0) sigma += get_veff(mf_grad, cell, dm0, kpts=kpts, with_j=True, with_nuc=True) diff --git a/gpu4pyscf/pbc/grad/kuhf.py b/gpu4pyscf/pbc/grad/kuhf.py index fa9a35787..985eef018 100644 --- a/gpu4pyscf/pbc/grad/kuhf.py +++ b/gpu4pyscf/pbc/grad/kuhf.py @@ -24,7 +24,7 @@ from gpu4pyscf.pbc.grad import krhf as krhf_grad from gpu4pyscf.lib.cupy_helper import contract, ensure_numpy from gpu4pyscf.pbc.grad.pp import vppnl_nuc_grad -from gpu4pyscf.pbc.dft import multigrid, multigrid_v2 +from gpu4pyscf.pbc.dft import multigrid from gpu4pyscf.pbc.gto import int1e from gpu4pyscf.pbc.grad.uhf import jk_energy_per_atom @@ -57,22 +57,10 @@ def grad_elec(mf_grad, mo_energy=None, mo_coeff=None, mo_occ=None, atmlst=None): dm0_sf = dm0[0] + dm0[1] ni = mf._numint - if isinstance(ni, multigrid.MultiGridNumInt): - raise NotImplementedError( - "Gradient with kpts not implemented with multigrid.MultiGridNumInt. " - "Please use the default KNumInt or multigrid_v2.MultiGridNumInt instead.") - elif isinstance(ni, multigrid_v2.MultiGridNumInt): - # Attention: The orbital derivative of vpploc term is in multigrid_v2.get_veff_ip1() function. - rho_g = multigrid_v2.evaluate_density_on_g_mesh(ni, dm0_sf, kpts) - rho_g = rho_g[0,0] - if cell._pseudo: - dh1e = multigrid.eval_vpplocG_SI_gradient(cell, ni.mesh, rho_g) * nkpts - else: - dh1e = multigrid.eval_nucG_SI_gradient(cell, ni.mesh, rho_g) * nkpts - - dh1e = dh1e.get() + if isinstance(ni, multigrid.MultiGridNumIntBase): + # Vne or pploc contribution is evaluated in energy_ee dh1e_kin = int1e.int1e_ipkin(cell, kpts) - dh1e -= krhf_grad.contract_h1e_dm(cell, dh1e_kin, dm0_sf, hermi=1) + dh1e = -krhf_grad.contract_h1e_dm(cell, dh1e_kin, dm0_sf, hermi=1) else: hcore_deriv = mf_grad.hcore_generator(cell, kpts) dh1e = cp.empty([natm, 3]) @@ -115,13 +103,11 @@ def energy_ee(self, dm, kpts): # When J is evaluated using mf.j_engine or mf.rsjk, it is identical to # the J from MultiGridNumInt. The contribution from J matrix can be # efficiently evaluated using the MultiGridNumInt integrator. - j_in_xc = ni is not None and isinstance(ni, multigrid_v2.MultiGridNumInt) + j_in_xc = ni is not None and isinstance(ni, multigrid.MultiGridNumIntBase) if j_in_xc: j_factor = 0 - de = multigrid_v2.get_veff_ip1( - ni, 'HF', dm[0]+dm[1], kpts=kpts, with_j=j_in_xc, - with_pseudo_vloc_orbital_derivative=True).get() - de /= len(kpts) + de = ni.energy_nuclear_gradient( + 'HF', dm, kpts=kpts, spin=1, with_j=j_in_xc, with_nuc=True) else: j_factor = 1 de = 0 diff --git a/gpu4pyscf/pbc/grad/kuhf_stress.py b/gpu4pyscf/pbc/grad/kuhf_stress.py index 70dee7486..9ad6bbb28 100644 --- a/gpu4pyscf/pbc/grad/kuhf_stress.py +++ b/gpu4pyscf/pbc/grad/kuhf_stress.py @@ -23,7 +23,7 @@ from gpu4pyscf.pbc.grad import kuhf as kuhf_grad from gpu4pyscf.pbc.gto import int1e from gpu4pyscf.pbc.tools.k2gamma import kpts_to_kmesh -from gpu4pyscf.pbc.grad.rks_stress import ewald +from gpu4pyscf.pbc.grad.rks_stress import ewald, _get_pp_nonloc_strain_derivatives from gpu4pyscf.pbc.grad.krhf_stress import get_nuc, get_veff def kernel(mf_grad): @@ -63,6 +63,8 @@ def kernel(mf_grad): sigma -= int1e.ovlp_strain_deriv(cell, dme0_sf, kpts) sigma += int1e.kin_strain_deriv(cell, dm0_sf, kpts) sigma += get_nuc(mf_grad, cell, dm0_sf, kpts) + if cell._pseudo: + sigma += _get_pp_nonloc_strain_derivatives(cell, cell.mesh, dm0_sf, kpts) t0 = log.timer_debug1('hcore derivatives', *t0) sigma += get_veff(mf_grad, cell, dm0, kpts) diff --git a/gpu4pyscf/pbc/grad/kuks.py b/gpu4pyscf/pbc/grad/kuks.py index 3f964dfe7..10f261d23 100644 --- a/gpu4pyscf/pbc/grad/kuks.py +++ b/gpu4pyscf/pbc/grad/kuks.py @@ -25,8 +25,7 @@ from gpu4pyscf.pbc.grad import kuhf as kuhf_grad from gpu4pyscf.pbc.grad import krks as krks_grad from gpu4pyscf.lib.cupy_helper import contract -from gpu4pyscf.pbc.dft import BeckeGrids -from gpu4pyscf.pbc.dft import multigrid, multigrid_v2 +from gpu4pyscf.pbc.dft import multigrid, BeckeGrids __all__ = ['Gradients'] @@ -46,21 +45,9 @@ def energy_ee(ks_grad, dm, kpts): omega, k_lr, k_sr = ni.rsh_and_hybrid_coeff(mf.xc) j_factor = 1 - if isinstance(ni, multigrid.MultiGridNumInt): - raise NotImplementedError( - "Gradient with kpts not implemented with multigrid.MultiGridNumInt. " - "Please use the default KNumInt or multigrid_v2.MultiGridNumInt instead.") - - if isinstance(ni, multigrid_v2.MultiGridNumInt): - if kpts is None: - nkpts = 1 - else: - nkpts = len(kpts) - exc = multigrid_v2.get_veff_ip1( - ni, mf.xc, dm, with_j=True, with_pseudo_vloc_orbital_derivative=True, kpts=kpts).get() - # exc of multigrid_v2 is the full response of dE/dX. However, - # get_veff in grad_elec evaluates the contraction Tr(dm, ). - exc /= nkpts + if isinstance(ni, multigrid.MultiGridNumIntBase): + exc = ni.energy_nuclear_gradient( + mf.xc, dm, kpts=kpts, spin=1, with_j=True, with_nuc=True) j_factor = 0 else: if ks_grad.grids is not None: diff --git a/gpu4pyscf/pbc/grad/kuks_stress.py b/gpu4pyscf/pbc/grad/kuks_stress.py index 80bf20d92..e0e3f5064 100644 --- a/gpu4pyscf/pbc/grad/kuks_stress.py +++ b/gpu4pyscf/pbc/grad/kuks_stress.py @@ -23,7 +23,7 @@ from gpu4pyscf.pbc.df import FFTDF from gpu4pyscf.pbc.dft.numint import KNumInt, eval_ao_kpts, _GTOvalOpt from gpu4pyscf.pbc.dft.krkspu import _set_U, _make_minao_lo, reference_mol -from gpu4pyscf.pbc.dft.multigrid_v2 import _uks_exc_strain_deriv, MultiGridNumInt +from gpu4pyscf.pbc.dft import multigrid, BeckeGrids from gpu4pyscf.pbc.grad.krks_stress import ( get_ovlp, _get_first_order_local_orbitals, _contract_coulomb_and_nuc) from gpu4pyscf.pbc.grad import kuks as kuks_grad @@ -51,13 +51,17 @@ def get_veff(mf_grad, cell, dm, kpts, with_j=False, with_nuc=False): ni = mf._numint is_hybrid = ni.libxc.is_hybrid_xc(mf.xc) + if isinstance(mf.grids, BeckeGrids): + raise NotImplementedError('gradients for BeckeGrids not supported') + j_factor = 1 if with_j else 0 if is_hybrid and with_rsjk is not None: with_j = False # TODO: with_nuc should be disabled for all-electron calculations - if isinstance(ni, MultiGridNumInt): - sigma = _uks_exc_strain_deriv(ni, mf.xc, dm, kpts, with_j, with_nuc) + if isinstance(ni, multigrid.MultiGridNumIntBase): + sigma = ni.energy_strain_gradient(mf.xc, dm, kpts, spin=1, + with_j=with_j, with_nuc=with_nuc) elif isinstance(ni, KNumInt): sigma = get_vxc(mf_grad, cell, dm, kpts, with_j, with_nuc) else: @@ -258,30 +262,18 @@ def kernel(mf_grad): log.debug('Computing stress tensor') cell = mf.cell - dm0 = mf.make_rdm1().sum(axis=0) + dm0 = mf.make_rdm1() + dm0_sf = dm0[0] + dm0[1] dme0 = mf_grad.make_rdm1e().sum(axis=0) sigma = ewald(cell) kpts = mf.kpts - kmesh = kpts_to_kmesh(cell, kpts, bound_by_supmol=True) sigma -= int1e.ovlp_strain_deriv(cell, dme0, kpts) - - scaled_kpts = kpts.dot(cell.lattice_vectors().T) - nkpts = len(kpts) - disp = 1e-5 - for x in range(3): - for y in range(3): - cell1, cell2 = _finite_diff_cells(cell, x, y, disp) - kpts1 = scaled_kpts.dot(cell1.reciprocal_vectors(norm_to=1)) - kpts2 = scaled_kpts.dot(cell2.reciprocal_vectors(norm_to=1)) - t1 = int1e.int1e_kin(cell1, kpts1, kmesh) - t2 = int1e.int1e_kin(cell2, kpts2, kmesh) - t1 = cp.einsum('kij,kji->', t1, dm0).real - t2 = cp.einsum('kij,kji->', t2, dm0).real - sigma[x,y] += (t1 - t2).get() / (2*disp) / nkpts + sigma += int1e.kin_strain_deriv(cell, dm0_sf, kpts) + if cell._pseudo: + sigma += _get_pp_nonloc_strain_derivatives(cell, cell.mesh, dm0_sf, kpts) t0 = log.timer_debug1('hcore derivatives', *t0) - dm0 = mf.make_rdm1() sigma += get_veff(mf_grad, cell, dm0, kpts=kpts, with_j=True, with_nuc=True) t0 = log.timer_debug1('Vxc and Coulomb derivatives', *t0) diff --git a/gpu4pyscf/pbc/grad/rhf.py b/gpu4pyscf/pbc/grad/rhf.py index eab6c1a48..edb62a12a 100644 --- a/gpu4pyscf/pbc/grad/rhf.py +++ b/gpu4pyscf/pbc/grad/rhf.py @@ -24,8 +24,7 @@ from gpu4pyscf.lib import logger import gpu4pyscf.grad.rhf as mol_rhf from gpu4pyscf.pbc.tools.k2gamma import kpts_to_kmesh -from gpu4pyscf.pbc.dft import multigrid_v2 -import gpu4pyscf.pbc.dft.multigrid as multigrid_v1 +from gpu4pyscf.pbc.dft import multigrid from gpu4pyscf.pbc.scf.rsjk import PBCJKMatrixOpt from gpu4pyscf.pbc.df.df import GDF from gpu4pyscf.pbc.gto import int1e @@ -122,7 +121,7 @@ def energy_ee(self, dm): # the GDF-based J. In this case, j_in_xc must be disabled, and the J # contribution must be evaluated using the GDF jk_energy_per_atom function. ni = mf._numint - j_in_xc = isinstance(ni, multigrid_v2.MultiGridNumInt) + j_in_xc = isinstance(ni, multigrid.MultiGridNumIntBase) de = 0 xc = getattr(mf, 'xc', 'HF') if xc.upper() == 'HF': @@ -144,9 +143,9 @@ def energy_ee(self, dm): # TODO: handle all-electron+GGA and pseudo+GGA differently # pseudo+GGA does not need to evaluate the gradients with PBCJKMatrixOpt if j_in_xc: - de += multigrid_v2.get_veff_ip1( - ni, xc, dm, with_j=j_in_xc, - with_pseudo_vloc_orbital_derivative=True).get() + assert isinstance(ni, multigrid.MultiGridNumIntBase) + de += ni.energy_nuclear_gradient( + xc, dm, spin=0, with_j=j_in_xc, with_nuc=True) j_factor = 0 elif xc.upper() != 'HF': from gpu4pyscf.pbc.grad.krks import get_vxc @@ -184,14 +183,8 @@ def grad_elec( de = self.energy_ee(dm0) ni = mf._numint - if isinstance(ni, multigrid_v2.MultiGridNumInt): - rhoG = multigrid_v2.evaluate_density_on_g_mesh(ni, dm0) - rhoG = rhoG[0,0] - if cell._pseudo: - de += multigrid_v1.eval_vpplocG_SI_gradient(cell, ni.mesh, rhoG).get() - else: - de += multigrid_v1.eval_nucG_SI_gradient(cell, ni.mesh, rhoG).get() - + if isinstance(ni, multigrid.MultiGridNumIntBase): + # Vne or pploc contribution is evaluated in energy_ee dh1e_kin = int1e.int1e_ipkin(cell) de -= contract_h1e_dm(cell, dh1e_kin, dm0, hermi=1) else: diff --git a/gpu4pyscf/pbc/grad/rhf_stress.py b/gpu4pyscf/pbc/grad/rhf_stress.py index 772c9581c..2508677c9 100644 --- a/gpu4pyscf/pbc/grad/rhf_stress.py +++ b/gpu4pyscf/pbc/grad/rhf_stress.py @@ -22,9 +22,10 @@ from gpu4pyscf.lib import logger from gpu4pyscf.pbc.grad import rhf as rhf_grad from gpu4pyscf.pbc.gto import int1e -from gpu4pyscf.pbc.grad.rks_stress import ewald from gpu4pyscf.pbc.scf.rsjk import PBCJKMatrixOpt from gpu4pyscf.pbc.df import aft, aft_jk +from gpu4pyscf.pbc.grad.rks_stress import ( + _get_pp_nonloc_strain_derivatives, ewald) ALIGNED = 256 @@ -61,6 +62,8 @@ def kernel(mf_grad): sigma -= int1e.ovlp_strain_deriv(cell, dme0) sigma += int1e.kin_strain_deriv(cell, dm0) sigma += get_nuc(mf_grad, cell, dm0) + if cell._pseudo: + sigma += _get_pp_nonloc_strain_derivatives(cell, cell.mesh, dm0) t0 = log.timer_debug1('hcore derivatives', *t0) sigma += get_veff(mf_grad, cell, dm0) diff --git a/gpu4pyscf/pbc/grad/rks_stress.py b/gpu4pyscf/pbc/grad/rks_stress.py index 8a835a290..a5a991dd5 100644 --- a/gpu4pyscf/pbc/grad/rks_stress.py +++ b/gpu4pyscf/pbc/grad/rks_stress.py @@ -57,7 +57,7 @@ from gpu4pyscf.pbc.df import FFTDF, ft_ao from gpu4pyscf.pbc.df.aft import get_SI, _get_ZSI from gpu4pyscf.pbc.dft.numint import NumInt, eval_ao_kpts, _GTOvalOpt -from gpu4pyscf.pbc.dft.multigrid_v2 import _rks_exc_strain_deriv, MultiGridNumInt +from gpu4pyscf.pbc.dft import multigrid from gpu4pyscf.pbc.grad import rks as rks_grad from gpu4pyscf.pbc.gto import int1e from gpu4pyscf.pbc.scf.rsjk import PBCJKMatrixOpt @@ -83,10 +83,12 @@ def _finite_diff_cells(cell, x, y, disp=1e-4, precision=None): e_strain = strain_tensor_dispalcement(x, y, disp) cell1 = cell.set_geom_(r.dot(e_strain.T), inplace=False) cell1.a = a.dot(e_strain.T) + cell1.mesh = cell.mesh e_strain = strain_tensor_dispalcement(x, y, -disp) cell2 = cell.set_geom_(r.dot(e_strain.T), inplace=False) cell2.a = a.dot(e_strain.T) + cell2.mesh = cell.mesh if cell.space_group_symmetry: cell1.build(False, False) @@ -193,8 +195,9 @@ def get_veff(mf_grad, cell, dm, with_j=False, with_nuc=False): with_j = False # TODO: with_nuc should be disabled for all-electron calculations - if isinstance(ni, MultiGridNumInt): - sigma = _rks_exc_strain_deriv(ni, mf.xc, dm[None], None, with_j, with_nuc) + if isinstance(ni, multigrid.MultiGridNumIntBase): + sigma = ni.energy_strain_gradient(mf.xc, dm[None], spin=0, + with_j=with_j, with_nuc=with_nuc) elif isinstance(ni, NumInt): sigma = get_vxc(mf_grad, cell, dm, with_j, with_nuc) else: @@ -364,7 +367,6 @@ def partial_dot(bra, ket): vpplocR = pbctools.ifft(vpplocG_0, mesh).real Ene = cp.einsum('xyg,g->xy', rho1[:,:,0], vpplocR).real.get() Ene += cp.einsum('g,xyg->xy', rhoG.conj(), vpplocG_1).real.get() * (1./ngrids) - Ene += _get_pp_nonloc_strain_derivatives(cell, mesh, dm) else: # SI corresponds to Fourier components of the fractional atomic # positions within the cell. It does not respond to the strain @@ -376,22 +378,81 @@ def partial_dot(bra, ket): out += Ene return out +def _get_Gv_bases(mesh, b): + Gx = cp.array(np.fft.fftfreq(mesh[0], 1./mesh[0]) * b[0,:,None]) + Gy = cp.array(np.fft.fftfreq(mesh[1], 1./mesh[1]) * b[1,:,None]) + Gz = cp.array(np.fft.fftfreq(mesh[2], 1./mesh[2]) * b[2,:,None]) + return (Gx, Gy, Gz) + def _get_vpplocG_strain_derivatives(cell, mesh): - disp = 1e-5 + assert cell.dimension == 3 + Gv_bases = _get_Gv_bases(mesh, cell.reciprocal_vectors()) + coords = cp.asarray(cell.atom_coords()) + SIx = cp.exp(-1j * coords.dot(Gv_bases[0])) + SIy = cp.exp(-1j * coords.dot(Gv_bases[1])) + SIz = cp.exp(-1j * coords.dot(Gv_bases[2])) + ngrids = np.prod(mesh) - v1 = cp.empty((3,3, ngrids), dtype=np.complex128) - SI = get_SI(cell, mesh=mesh) - for x in range(3): - for y in range(3): - cell1, cell2 = _finite_diff_cells(cell, x, y, disp) - vpplocG1 = pseudo.get_vlocG(cell1, cell1.get_Gv(mesh)) - vpplocG2 = pseudo.get_vlocG(cell2, cell2.get_Gv(mesh)) - vpplocG1 = -np.einsum('ij,ij->j', SI, vpplocG1) - vpplocG2 = -np.einsum('ij,ij->j', SI, vpplocG2) - v1[x,y] = asarray((vpplocG1 - vpplocG2) / (2*disp)) - vpplocG = pseudo.get_vlocG(cell, cell.get_Gv(mesh)) - v0 = asarray(-np.einsum('ij,ij->j', SI, vpplocG)) - return v0, v1 + Gx, Gy, Gz = Gv_bases + GvT = Gx[:,:,None,None] + Gy[:,None,:,None] + Gz[:,None,None,:] + GvT = GvT.reshape(3, ngrids) + G2 = cp.einsum('xg,xg->g', GvT, GvT) + coulG = 4 * np.pi / G2 + coulG[0] = 0 + xyG = cp.einsum('xg,yg->xyg', GvT, GvT) + + charges = cell.atom_charges() + + vlocG0 = 0 + vlocG_0 = cp.zeros(ngrids, dtype=np.complex128) + vlocG_1 = cp.zeros((3, 3, ngrids), dtype=np.complex128) + + for ia in range(cell.natm): + symb = cell.atom_symbol(ia) + if symb not in cell._pseudo: + continue + + pp = cell._pseudo[symb] + rloc, nexp, cexp = pp[1:3+1] + + SI = (SIx[ia,:,None,None] * SIy[ia,:,None] * SIz[ia]).ravel() + x = G2 * rloc**2 + expx = cp.exp(-0.5*x) + SI *= expx + Z = charges[ia] + + coef1 = -Z * coulG * SI * (2/G2 + rloc**2) + coef1[0] = 0 + + cfacs = 0 + dcfacs = 0 + if nexp >= 1: + cfacs += cexp[0] + if nexp >= 2: + cfacs += cexp[1] * (3 - x) + dcfacs -= cexp[1] + if nexp >= 3: + cfacs += cexp[2] * (15 - 10*x + x*x) + dcfacs += cexp[2] * (-10 + 2*x) + if nexp >= 4: + cfacs += cexp[3] * (105 - 105*x + 21*x*x - x*x*x) + dcfacs += cexp[3] * (-105 + 42*x - 3*x*x) + + coef2 = ( + (2*np.pi)**1.5 + * rloc**5 + * SI + * (cfacs - 2 * dcfacs) + ) + + vlocG0 += 2*np.pi*Z*rloc**2 + vlocG_0 -= Z * coulG * SI + vlocG_0 += (2*np.pi)**(3/2.)*rloc**3 * cfacs * SI + + vlocG_1 += (coef1 + coef2) * xyG + + vlocG_0[0] += vlocG0 + return vlocG_0, vlocG_1 def _get_pp_nonloc_strain_derivatives(cell, mesh, dm_kpts, kpts=None): if kpts is None: @@ -518,16 +579,10 @@ def kernel(mf_grad): dme0 = mf_grad.make_rdm1e() sigma = ewald(cell) sigma -= int1e.ovlp_strain_deriv(cell, dme0) - - disp = 1e-5 - for x in range(3): - for y in range(3): - cell1, cell2 = _finite_diff_cells(cell, x, y, disp) - t1 = int1e.int1e_kin(cell1) - t2 = int1e.int1e_kin(cell2) - t1 = cp.einsum('ij,ji->', t1, dm0) - t2 = cp.einsum('ij,ji->', t2, dm0) - sigma[x,y] += (t1 - t2) / (2*disp) + sigma += int1e.kin_strain_deriv(cell, dm0) + if cell._pseudo: + # pploc contribution is evaluated in get_veff + sigma += _get_pp_nonloc_strain_derivatives(cell, cell.mesh, dm0) t0 = log.timer_debug1('hcore derivatives', *t0) sigma += get_veff(mf_grad, cell, dm0, with_j=True, with_nuc=True) diff --git a/gpu4pyscf/pbc/grad/tests/test_pbc_grad_kuks.py b/gpu4pyscf/pbc/grad/tests/test_pbc_grad_kuks.py index 68f6a9676..c03d919d9 100644 --- a/gpu4pyscf/pbc/grad/tests/test_pbc_grad_kuks.py +++ b/gpu4pyscf/pbc/grad/tests/test_pbc_grad_kuks.py @@ -139,7 +139,7 @@ def test_rsjk_lda_grad(self): mf.j_engine = PBCJMatrixOpt(cell) g_scan = mf.Gradients().as_scanner() g = g_scan(cell)[1] - np.testing.assert_almost_equal(g, g_ref, 7) + np.testing.assert_almost_equal(g, g_ref, 5) def test_df_lda_grad(self): kpts = cell_be.make_kpts([1,1,3]) diff --git a/gpu4pyscf/pbc/grad/tests/test_pbc_grad_rks.py b/gpu4pyscf/pbc/grad/tests/test_pbc_grad_rks.py index 32bcb6482..1b522e51e 100644 --- a/gpu4pyscf/pbc/grad/tests/test_pbc_grad_rks.py +++ b/gpu4pyscf/pbc/grad/tests/test_pbc_grad_rks.py @@ -182,7 +182,7 @@ def test_gga_grad_nonorth(self): mf = mf.multigrid_numint() g1 = mf.Gradients().kernel() - self.assertAlmostEqual(abs(g - g1).max(), 0, 8) + self.assertAlmostEqual(abs(g - g1).max(), 0, delta=5e-6) @unittest.skipIf(num_devices > 1, '') def test_mgga_grad(self): diff --git a/gpu4pyscf/pbc/grad/tests/test_pbc_grad_uks.py b/gpu4pyscf/pbc/grad/tests/test_pbc_grad_uks.py index 3fe0d8678..e874ed5c0 100644 --- a/gpu4pyscf/pbc/grad/tests/test_pbc_grad_uks.py +++ b/gpu4pyscf/pbc/grad/tests/test_pbc_grad_uks.py @@ -120,7 +120,7 @@ def test_lda_grad_nonorth(self): mf.run() mf = mf.multigrid_numint() g = mf.nuc_grad_method().kernel() - self.assertAlmostEqual(abs(g - ref).max(), 0, 6) + self.assertAlmostEqual(abs(g - ref).max(), 0, delta=5e-6) @unittest.skipIf(num_devices > 1, '') def test_gga_grad(self): @@ -142,7 +142,7 @@ def test_gga_grad_nonorth(self): mf.run() mf = mf.multigrid_numint() g = mf.nuc_grad_method().kernel() - self.assertAlmostEqual(abs(g - ref).max(), 0, 6) + self.assertAlmostEqual(abs(g - ref).max(), 0, delta=5e-6) @unittest.skipIf(num_devices > 1, '') def test_mgga_grad(self): diff --git a/gpu4pyscf/pbc/grad/tests/test_pbc_krks_stress.py b/gpu4pyscf/pbc/grad/tests/test_pbc_krks_stress.py index bc0d595a5..58a49eea3 100644 --- a/gpu4pyscf/pbc/grad/tests/test_pbc_krks_stress.py +++ b/gpu4pyscf/pbc/grad/tests/test_pbc_krks_stress.py @@ -15,6 +15,7 @@ import unittest import numpy as np +import cupy as cp from pyscf.gto import ATOM_OF, intor_cross from pyscf.pbc import dft, gto, grad from pyscf.pbc.tools import pbc @@ -26,8 +27,9 @@ from gpu4pyscf.pbc.grad.krks_stress import _finite_diff_cells from gpu4pyscf.pbc.scf.j_engine import PBCJMatrixOpt from gpu4pyscf.pbc.scf.rsjk import PBCJKMatrixOpt -from gpu4pyscf.pbc.dft.multigrid_v2 import _rks_exc_strain_deriv, MultiGridNumInt +from gpu4pyscf.pbc.dft.multigrid_v3 import MultiGridNumInt from gpu4pyscf.lib.multi_gpu import num_devices +from gpu4pyscf.pbc.lib.kpts_helper import fft_matrix import pytest class KnownValues(unittest.TestCase): @@ -86,24 +88,38 @@ def test_get_vxc_lda(self): kmesh = [3, 1, 1] kpts = cell.make_kpts(kmesh) nao = cell.nao - dm = np.random.rand(np.prod(kmesh), nao, nao) - (.5+.1j) + dm = np.random.rand(np.prod(kmesh), nao, nao) - .5 + phase = fft_matrix(kmesh).get() / np.prod(kmesh) + dm = np.einsum('Lpq,Lk->kpq', dm, phase.conj()) dm = np.einsum('kpi,kqi->kpq', dm, dm.conj()) xc = 'lda,' mf_grad = krks.Gradients(cell.KRKS(xc=xc, kpts=kpts).to_gpu()) - dat = krks_stress.get_vxc(mf_grad, cell, dm, kpts=kpts) + dat1 = krks_stress.get_vxc(mf_grad, cell, dm, kpts=kpts) + + ni = MultiGridNumInt(cell) + ni.allow_mesh_reduction = False + dat2 = ni.energy_strain_gradient(xc, dm, kpts, spin=0, with_j=False, with_nuc=False) + assert abs(dat1 - dat2).max() < 1e-6 + ni = KNumInt() for (i, j) in [(0, 0), (0, 1), (0, 2), (2, 0), (2, 2)]: - cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-5) + cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-4) cell1.precision = 1e-10 cell2.precision = 1e-10 exc1 = ni.nr_rks(cell1, UniformGrids(cell1), xc, dm, kpts=cell1.make_kpts(kmesh))[1] exc2 = ni.nr_rks(cell2, UniformGrids(cell2), xc, dm, kpts=cell2.make_kpts(kmesh))[1] - assert abs(dat[i,j] - (exc1 - exc2)/2e-5) < 1e-8 + assert abs(dat1[i,j] - (exc1 - exc2)/2e-4) < 1e-7 + assert abs(dat2[i,j] - (exc1 - exc2)/2e-4) < 1e-7 - ref = dat - ni = MultiGridNumInt(cell).build() - dat = _rks_exc_strain_deriv(ni, xc, dm, kpts, with_j=False, with_nuc=False) - assert abs(dat - ref).max() < 1e-6 + ni = MultiGridNumInt(cell) + dat2 = ni.energy_strain_gradient(xc, dm, kpts, spin=0, with_j=False, with_nuc=False) + for (i, j) in [(0, 0), (0, 1), (0, 2), (2, 0), (2, 2)]: + cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-4) + cell1.precision = 1e-10 + cell2.precision = 1e-10 + exc1 = MultiGridNumInt(cell1).nr_rks(cell1, None, xc, dm, kpts=cell1.make_kpts(kmesh))[1] + exc2 = MultiGridNumInt(cell2).nr_rks(cell2, None, xc, dm, kpts=cell2.make_kpts(kmesh))[1] + assert abs(dat2[i,j] - (exc1 - exc2)/2e-4) < 3e-7 def test_get_vxc_gga(self): a = np.eye(3) * 5 @@ -114,24 +130,28 @@ def test_get_vxc_gga(self): kmesh = [3, 1, 1] kpts = cell.make_kpts(kmesh) nao = cell.nao - dm = np.random.rand(np.prod(kmesh), nao, nao) - (.5+.1j) + dm = np.random.rand(np.prod(kmesh), nao, nao) - .5 + phase = fft_matrix(kmesh).get() / np.prod(kmesh) + dm = np.einsum('Lpq,Lk->kpq', dm, phase.conj()) dm = np.einsum('kpi,kqi->kpq', dm, dm.conj()) xc = 'pbe,' mf_grad = krks.Gradients(cell.KRKS(xc=xc, kpts=kpts).to_gpu()) - dat = krks_stress.get_vxc(mf_grad, cell, dm, kpts=kpts) + dat1 = krks_stress.get_vxc(mf_grad, cell, dm, kpts=kpts) + + ni = MultiGridNumInt(cell) + ni.allow_mesh_reduction = False + dat2 = ni.energy_strain_gradient(xc, dm, kpts, spin=0, with_j=False, with_nuc=False) + assert abs(dat1 - dat2).max() < 1e-6 + ni = KNumInt() for (i, j) in [(0, 0), (0, 1), (0, 2), (2, 0), (2, 2)]: - cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-5) + cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-4) cell1.precision = 1e-10 cell2.precision = 1e-10 exc1 = ni.nr_rks(cell1, UniformGrids(cell1), xc, dm, kpts=cell1.make_kpts(kmesh))[1] exc2 = ni.nr_rks(cell2, UniformGrids(cell2), xc, dm, kpts=cell2.make_kpts(kmesh))[1] - assert abs(dat[i,j] - (exc1 - exc2)/2e-5) < 1e-8 - - ref = dat - ni = MultiGridNumInt(cell).build() - dat = _rks_exc_strain_deriv(ni, xc, dm, kpts, with_j=False, with_nuc=False) - assert abs(dat - ref).max() < 1e-6 + assert abs(dat1[i,j] - (exc1 - exc2)/2e-4) < 1e-7 + assert abs(dat2[i,j] - (exc1 - exc2)/2e-4) < 1e-7 def test_get_vxc_mgga(self): a = np.eye(3) * 5 @@ -142,24 +162,28 @@ def test_get_vxc_mgga(self): kmesh = [3, 1, 1] kpts = cell.make_kpts(kmesh) nao = cell.nao - dm = np.random.rand(np.prod(kmesh), nao, nao) - (.5+.1j) + dm = np.random.rand(np.prod(kmesh), nao, nao) - .5 + phase = fft_matrix(kmesh).get() / np.prod(kmesh) + dm = np.einsum('Lpq,Lk->kpq', dm, phase.conj()) dm = np.einsum('kpi,kqi->kpq', dm, dm.conj()) xc = 'm06,' mf_grad = krks.Gradients(cell.KRKS(xc=xc, kpts=kpts).to_gpu()) - dat = krks_stress.get_vxc(mf_grad, cell, dm, kpts=kpts) + dat1 = krks_stress.get_vxc(mf_grad, cell, dm, kpts=kpts) + + ni = MultiGridNumInt(cell) + ni.allow_mesh_reduction = False + dat2 = ni.energy_strain_gradient(xc, dm, kpts, spin=0, with_j=False, with_nuc=False) + assert abs(dat1 - dat2).max() < 1e-6 + ni = KNumInt() for (i, j) in [(0, 0), (0, 1), (0, 2), (2, 0), (2, 2)]: - cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-5) + cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-4) cell1.precision = 1e-10 cell2.precision = 1e-10 exc1 = ni.nr_rks(cell1, UniformGrids(cell1), xc, dm, kpts=cell1.make_kpts(kmesh))[1] exc2 = ni.nr_rks(cell2, UniformGrids(cell2), xc, dm, kpts=cell2.make_kpts(kmesh))[1] - assert abs(dat[i,j] - (exc1 - exc2)/2e-5) < 1e-8 - - ref = dat - ni = MultiGridNumInt(cell).build() - dat = _rks_exc_strain_deriv(ni, xc, dm, kpts, with_j=False, with_nuc=False) - assert abs(dat - ref).max() < 1e-6 + assert abs(dat1[i,j] - (exc1 - exc2)/2e-4) < 1e-7 + assert abs(dat2[i,j] - (exc1 - exc2)/2e-4) < 1e-7 def test_get_j(self): a = np.eye(3) * 5 @@ -169,7 +193,9 @@ def test_get_j(self): basis=[[0, [.5, 1]], [1, [.8, 1]], [2, [.6, 1]]], a=a, unit='Bohr') kmesh = [3, 1, 3] nao = cell.nao - dm = np.random.rand(np.prod(kmesh), nao, nao) - (.5+.1j) + dm = np.random.rand(np.prod(kmesh), nao, nao) - .5 + phase = fft_matrix(kmesh).get() / np.prod(kmesh) + dm = np.einsum('Lpq,Lk->kpq', dm, phase.conj()) dm = np.einsum('kpi,kqi->kpq', dm, dm.conj()) xc = 'lda,' kpts = cell.make_kpts(kmesh) @@ -177,7 +203,7 @@ def test_get_j(self): dat = krks_stress.get_vxc(mf_grad, cell, dm, kpts=kpts, with_j=True) ni = KNumInt() for (i, j) in [(0, 0), (0, 1), (0, 2), (2, 1), (2, 2)]: - cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-5) + cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-4) vj1 = FFTDF(cell1).get_jk(dm, kpts=cell1.make_kpts(kmesh), with_k=False)[0] vj1 *= .5 exc1 = ni.nr_rks(cell1, UniformGrids(cell1), xc, dm, kpts=cell1.make_kpts(kmesh))[1] @@ -186,7 +212,7 @@ def test_get_j(self): exc2 = ni.nr_rks(cell2, UniformGrids(cell2), xc, dm, kpts=cell2.make_kpts(kmesh))[1] de = np.einsum('kij,kji', dm, (vj1-vj2)) / len(kpts) de += exc1 - exc2 - assert abs(dat[i,j] - de/2e-5) < 1e-8 + assert abs(dat[i,j] - de/2e-4) < 1e-7 def test_get_nuc(self): a = np.eye(3) * 5 @@ -196,15 +222,19 @@ def test_get_nuc(self): basis=[[0, [.5, 1]], [1, [.8, 1]], [2, [.6, 1]]], a=a, unit='Bohr') kmesh = [3, 1, 1] nao = cell.nao - dm = np.random.rand(np.prod(kmesh), nao, nao) - (.5+.1j) + dm = np.random.rand(np.prod(kmesh), nao, nao) - .5 + phase = fft_matrix(kmesh).get() / np.prod(kmesh) + dm = np.einsum('Lpq,Lk->kpq', dm, phase.conj()) dm = np.einsum('kpi,kqi->kpq', dm, dm.conj()) xc = 'lda,' kpts = cell.make_kpts(kmesh) mf_grad = krks.Gradients(cell.KRKS(xc=xc, kpts=kpts).to_gpu()) dat = krks_stress.get_vxc(mf_grad, cell, dm, kpts=kpts, with_nuc=True) + dat += krks_stress._get_pp_nonloc_strain_derivatives( + cell, cell.mesh, cp.array(dm), kpts) ni = KNumInt() for (i, j) in [(0, 0), (0, 1), (0, 2), (2, 1), (2, 2)]: - cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-5) + cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-4) cell1.precision = 1e-10 cell2.precision = 1e-10 vne1 = FFTDF(cell1).get_nuc(kpts=cell1.make_kpts(kmesh)) @@ -213,7 +243,7 @@ def test_get_nuc(self): exc2 = ni.nr_rks(cell2, UniformGrids(cell2), xc, dm, kpts=cell2.make_kpts(kmesh))[1] de = np.einsum('kij,kji', dm, (vne1-vne2)) / len(kpts) de += exc1 - exc2 - assert abs(dat[i,j] - de/2e-5) < 1e-8 + assert abs(dat[i,j] - de/2e-4) < 1e-7 def test_get_pp(self): a = np.eye(3) * 5 @@ -224,15 +254,19 @@ def test_get_pp(self): pseudo='gth-pade', a=a, unit='Bohr', precision=1e-9) kmesh = [3, 1, 1] nao = cell.nao - dm = np.random.rand(np.prod(kmesh), nao, nao) - (.5+.1j) + dm = np.random.rand(np.prod(kmesh), nao, nao) - .5 + phase = fft_matrix(kmesh).get() / np.prod(kmesh) + dm = np.einsum('Lpq,Lk->kpq', dm, phase.conj()) dm = np.einsum('kpi,kqi->kpq', dm, dm.conj()) xc = 'lda,' kpts = cell.make_kpts(kmesh) mf_grad = krks.Gradients(cell.KRKS(xc=xc, kpts=kpts).to_gpu()) dat = krks_stress.get_vxc(mf_grad, cell, dm, kpts=kpts, with_nuc=True) + dat += krks_stress._get_pp_nonloc_strain_derivatives( + cell, cell.mesh, cp.array(dm), kpts) ni = KNumInt() for (i, j) in [(0, 0), (0, 1), (0, 2), (2, 1), (2, 2)]: - cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-5) + cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-4) cell1.precision = 1e-10 cell2.precision = 1e-10 vne1 = FFTDF(cell1).get_pp(kpts=cell1.make_kpts(kmesh)) @@ -241,7 +275,7 @@ def test_get_pp(self): exc2 = ni.nr_rks(cell2, UniformGrids(cell2), xc, dm, kpts=cell2.make_kpts(kmesh))[1] de = np.einsum('kij,kji', dm, (vne1-vne2)) / len(kpts) de += exc1 - exc2 - assert abs(dat[i,j] - de/2e-5) < 1e-8 + assert abs(dat[i,j] - de/2e-4) < 3e-7 def test_lda_vs_finite_difference(self): a = np.eye(3) * 3 @@ -274,7 +308,14 @@ def test_gga_vs_finite_difference(self): kmesh = [3, 1, 1] mf = cell.KRKS(xc=xc, kpts=cell.make_kpts(kmesh)).to_gpu().run() mf_grad = mf.Gradients() + ref = mf_grad.get_stress() + + mf = mf.multigrid_numint() + mf._numint.allow_mesh_reduction = False + mf_grad = mf.Gradients() dat = mf_grad.get_stress() + assert abs(dat - ref).max() < 1e-6 + mf_scanner = mf.as_scanner() vol = cell.vol for (i, j) in [(0, 0), (0, 1), (0, 2), (1, 0), (2, 2)]: @@ -283,12 +324,6 @@ def test_gga_vs_finite_difference(self): e2 = mf_scanner(cell2) assert abs(dat[i,j] - (e1-e2)/2e-3/vol) < 1e-6 - ref = dat - mf = mf.multigrid_numint() - mf_grad = mf.Gradients() - dat = mf_grad.get_stress() - assert abs(dat - ref).max() < 1e-6 - @unittest.skipIf(num_devices > 1, '') def test_mgga_vs_finite_difference(self): a = np.eye(3) * 3.5 @@ -300,7 +335,9 @@ def test_mgga_vs_finite_difference(self): xc = 'scan' kmesh = [3, 1, 1] mf = cell.KRKS(xc=xc, kpts=cell.make_kpts(kmesh)).to_gpu() - mf = mf.multigrid_numint().run() + mf = mf.multigrid_numint() + mf._numint.allow_mesh_reduction = False + mf.run() mf_grad = mf.Gradients() dat = mf_grad.get_stress() mf_scanner = mf.as_scanner() @@ -384,7 +421,7 @@ def test_hubbard_U(self): e1 = mf.get_veff().E_U.real mf.reset(cell2) e2 = mf.get_veff().E_U.real - assert abs(sigma[i,j] - (e1 - e2) / 2e-4) < 1e-8 + assert abs(sigma[i,j] - (e1 - e2) / 2e-4) < 1e-7 @pytest.mark.slow def test_krkspu_finite_diff(self): diff --git a/gpu4pyscf/pbc/grad/tests/test_pbc_kuks_stress.py b/gpu4pyscf/pbc/grad/tests/test_pbc_kuks_stress.py index 84eed8ab5..dc4b6d347 100644 --- a/gpu4pyscf/pbc/grad/tests/test_pbc_kuks_stress.py +++ b/gpu4pyscf/pbc/grad/tests/test_pbc_kuks_stress.py @@ -26,8 +26,9 @@ from gpu4pyscf.pbc.grad.kuks_stress import _finite_diff_cells from gpu4pyscf.pbc.scf.j_engine import PBCJMatrixOpt from gpu4pyscf.pbc.scf.rsjk import PBCJKMatrixOpt -from gpu4pyscf.pbc.dft.multigrid_v2 import _uks_exc_strain_deriv, MultiGridNumInt +from gpu4pyscf.pbc.dft.multigrid_v3 import MultiGridNumInt from gpu4pyscf.lib.multi_gpu import num_devices +from gpu4pyscf.pbc.lib.kpts_helper import fft_matrix import pytest class KnownValues(unittest.TestCase): @@ -40,24 +41,38 @@ def test_get_vxc_lda(self): kmesh = [3, 1, 1] kpts = cell.make_kpts(kmesh) nao = cell.nao - dm = np.random.rand(2,np.prod(kmesh), nao, nao) - (.5+.1j) + dm = np.random.rand(2,np.prod(kmesh), nao, nao) - .5 + phase = fft_matrix(kmesh).get() / np.prod(kmesh) + dm = np.einsum('sLpq,Lk->skpq', dm, phase.conj()) dm = np.einsum('skpi,skqi->skpq', dm, dm.conj()) xc = 'lda,' mf_grad = kuks.Gradients(cell.KUKS(xc=xc, kpts=kpts).to_gpu()) - dat = kuks_stress.get_vxc(mf_grad, cell, dm, kpts=kpts) + dat1 = kuks_stress.get_vxc(mf_grad, cell, dm, kpts=kpts) + + ni = MultiGridNumInt(cell) + ni.allow_mesh_reduction = False + dat2 = ni.energy_strain_gradient(xc, dm, kpts, spin=1, with_j=False, with_nuc=False) + assert abs(dat1 - dat2).max() < 1e-6 + ni = KNumInt() for (i, j) in [(0, 0), (0, 1), (0, 2), (2, 0), (2, 2)]: - cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-5) + cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-4) cell1.precision = 1e-10 cell2.precision = 1e-10 exc1 = ni.nr_uks(cell1, UniformGrids(cell1), xc, dm, kpts=cell1.make_kpts(kmesh))[1] exc2 = ni.nr_uks(cell2, UniformGrids(cell2), xc, dm, kpts=cell2.make_kpts(kmesh))[1] - assert abs(dat[i,j] - (exc1 - exc2)/2e-5) < 1e-9 + assert abs(dat1[i,j] - (exc1 - exc2)/2e-4) < 1e-7 + assert abs(dat2[i,j] - (exc1 - exc2)/2e-4) < 1e-7 - ref = dat - ni = MultiGridNumInt(cell).build() - dat = _uks_exc_strain_deriv(ni, xc, dm, kpts, with_j=False, with_nuc=False) - assert abs(dat - ref).max() < 1e-6 + ni = MultiGridNumInt(cell) + dat2 = ni.energy_strain_gradient(xc, dm, kpts, spin=1, with_j=False, with_nuc=False) + for (i, j) in [(0, 0), (0, 1), (0, 2), (2, 0), (2, 2)]: + cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-4) + cell1.precision = 1e-10 + cell2.precision = 1e-10 + exc1 = MultiGridNumInt(cell1).nr_uks(cell1, None, xc, dm, kpts=cell1.make_kpts(kmesh))[1] + exc2 = MultiGridNumInt(cell2).nr_uks(cell2, None, xc, dm, kpts=cell2.make_kpts(kmesh))[1] + assert abs(dat2[i,j] - (exc1 - exc2)/2e-4) < 3e-7 def test_get_vxc_gga(self): a = np.eye(3) * 5 @@ -68,24 +83,28 @@ def test_get_vxc_gga(self): kmesh = [3, 1, 1] kpts = cell.make_kpts(kmesh) nao = cell.nao - dm = np.random.rand(2,np.prod(kmesh), nao, nao) - (.5+.1j) + dm = np.random.rand(2,np.prod(kmesh), nao, nao) - .5 + phase = fft_matrix(kmesh).get() / np.prod(kmesh) + dm = np.einsum('sLpq,Lk->skpq', dm, phase.conj()) dm = np.einsum('skpi,skqi->skpq', dm, dm.conj()) xc = 'pbe,' mf_grad = kuks.Gradients(cell.KUKS(xc=xc, kpts=kpts).to_gpu()) - dat = kuks_stress.get_vxc(mf_grad, cell, dm, kpts=kpts) + dat1 = kuks_stress.get_vxc(mf_grad, cell, dm, kpts=kpts) + + ni = MultiGridNumInt(cell) + ni.allow_mesh_reduction = False + dat2 = ni.energy_strain_gradient(xc, dm, kpts, spin=1, with_j=False, with_nuc=False) + assert abs(dat1 - dat2).max() < 1e-6 + ni = KNumInt() for (i, j) in [(0, 0), (0, 1), (0, 2), (2, 0), (2, 2)]: - cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-5) + cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-4) cell1.precision = 1e-10 cell2.precision = 1e-10 exc1 = ni.nr_uks(cell1, UniformGrids(cell1), xc, dm, kpts=cell1.make_kpts(kmesh))[1] exc2 = ni.nr_uks(cell2, UniformGrids(cell2), xc, dm, kpts=cell2.make_kpts(kmesh))[1] - assert abs(dat[i,j] - (exc1 - exc2)/2e-5) < 1e-8 - - ref = dat - ni = MultiGridNumInt(cell).build() - dat = _uks_exc_strain_deriv(ni, xc, dm, kpts, with_j=False, with_nuc=False) - assert abs(dat - ref).max() < 1e-6 + assert abs(dat1[i,j] - (exc1 - exc2)/2e-4) < 1e-7 + assert abs(dat2[i,j] - (exc1 - exc2)/2e-4) < 1e-7 def test_get_vxc_mgga(self): a = np.eye(3) * 5 @@ -96,24 +115,28 @@ def test_get_vxc_mgga(self): kmesh = [3, 1, 1] kpts = cell.make_kpts(kmesh) nao = cell.nao - dm = np.random.rand(2,np.prod(kmesh), nao, nao) - (.5+.1j) + dm = np.random.rand(2,np.prod(kmesh), nao, nao) - .5 + phase = fft_matrix(kmesh).get() / np.prod(kmesh) + dm = np.einsum('sLpq,Lk->skpq', dm, phase.conj()) dm = np.einsum('skpi,skqi->skpq', dm, dm.conj()) xc = 'm06,' mf_grad = kuks.Gradients(cell.KUKS(xc=xc, kpts=kpts).to_gpu()) - dat = kuks_stress.get_vxc(mf_grad, cell, dm, kpts=kpts) + dat1 = kuks_stress.get_vxc(mf_grad, cell, dm, kpts=kpts) + + ni = MultiGridNumInt(cell) + ni.allow_mesh_reduction = False + dat2 = ni.energy_strain_gradient(xc, dm, kpts, spin=1, with_j=False, with_nuc=False) + assert abs(dat1 - dat2).max() < 1e-6 + ni = KNumInt() for (i, j) in [(0, 0), (0, 1), (0, 2), (2, 0), (2, 2)]: - cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-5) + cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-4) cell1.precision = 1e-10 cell2.precision = 1e-10 exc1 = ni.nr_uks(cell1, UniformGrids(cell1), xc, dm, kpts=cell1.make_kpts(kmesh))[1] exc2 = ni.nr_uks(cell2, UniformGrids(cell2), xc, dm, kpts=cell2.make_kpts(kmesh))[1] - assert abs(dat[i,j] - (exc1 - exc2)/2e-5) < 1e-9 - - ref = dat - ni = MultiGridNumInt(cell).build() - dat = _uks_exc_strain_deriv(ni, xc, dm, kpts, with_j=False, with_nuc=False) - assert abs(dat - ref).max() < 1e-6 + assert abs(dat1[i,j] - (exc1 - exc2)/2e-4) < 1e-7 + assert abs(dat2[i,j] - (exc1 - exc2)/2e-4) < 1e-7 def test_get_j(self): a = np.eye(3) * 5 @@ -123,7 +146,9 @@ def test_get_j(self): basis=[[0, [.5, 1]], [1, [.8, 1]], [2, [.6, 1]]], a=a, unit='Bohr') kmesh = [3, 1, 3] nao = cell.nao - dm = np.random.rand(2,np.prod(kmesh), nao, nao) - (.5+.1j) + dm = np.random.rand(2,np.prod(kmesh), nao, nao) - .5 + phase = fft_matrix(kmesh).get() / np.prod(kmesh) + dm = np.einsum('sLpq,Lk->skpq', dm, phase.conj()) dm *= .5 dm = np.einsum('skpi,skqi->skpq', dm, dm.conj()) xc = 'lda,' @@ -174,7 +199,14 @@ def test_gga_vs_finite_difference(self): kmesh = [3, 1, 1] mf = cell.KUKS(xc=xc, kpts=cell.make_kpts(kmesh)).to_gpu().run() mf_grad = mf.Gradients() - dat = mf_grad.get_stress() + dat1 = mf_grad.get_stress() + + mf = mf.multigrid_numint() + mf._numint.allow_mesh_reduction = False + mf_grad = mf.Gradients() + dat2 = mf_grad.get_stress() + assert abs(dat1 - dat2).max() < 1e-6 + mf_scanner = mf.as_scanner() vol = cell.vol for (i, j) in [(0, 0), (0, 1), (0, 2), (1, 0), (2, 2)]: @@ -183,13 +215,8 @@ def test_gga_vs_finite_difference(self): cell2.precision = 1e-10 e1 = mf_scanner(cell1) e2 = mf_scanner(cell2) - assert abs(dat[i,j] - (e1-e2)/2e-3/vol) < 1e-6 - - ref = dat - mf = mf.multigrid_numint() - mf_grad = mf.Gradients() - dat = mf_grad.get_stress() - assert abs(dat - ref).max() < 1e-6 + assert abs(dat1[i,j] - (e1-e2)/2e-3/vol) < 1e-6 + assert abs(dat2[i,j] - (e1-e2)/2e-3/vol) < 1e-6 @unittest.skipIf(num_devices > 1, '') def test_mgga_vs_finite_difference(self): @@ -199,10 +226,12 @@ def test_mgga_vs_finite_difference(self): cell = gto.M(atom='H 1 1 1; H 2 1.5 2.4', basis=[[0, [1.5, 1]], [1, [.8, 1]]], a=a, unit='Bohr', verbose=0) - xc = 'rscan' + xc = 'r2scan' kmesh = [3, 1, 1] mf = cell.KUKS(xc=xc, kpts=cell.make_kpts(kmesh)).to_gpu() - mf = mf.multigrid_numint().run() + mf = mf.multigrid_numint() + mf._numint.allow_mesh_reduction = False + mf.run() mf_grad = mf.Gradients() dat = mf_grad.get_stress() mf_scanner = mf.as_scanner() diff --git a/gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py b/gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py index 0a5c86fb2..8eb47b931 100644 --- a/gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py +++ b/gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py @@ -34,7 +34,7 @@ from pyscf.pbc.gto.pseudo.pp_int import fake_cell_vnl, _int_vnl from pyscf.pbc.lib.kpts_helper import gamma_point import gpu4pyscf.pbc.dft.multigrid as multigrid_v1 -import gpu4pyscf.pbc.dft.multigrid_v2 as multigrid_v2 +from gpu4pyscf.pbc.dft import multigrid_v3 import pytest disp = 1e-4 @@ -318,6 +318,7 @@ def test_iron_fd(self): self._fd_check(cell_fe, atom_id=1, cart_id=0, places=4) def test_pseudo_gradient_term_with_zero_nexp(self): + from gpu4pyscf.pbc.dft.multigrid_v3 import _pploc_derivatives, _get_Gv_bases cell = pyscf.M( a = np.array([ [3.18693029, 0.0, 0.0], @@ -382,14 +383,18 @@ def test_pseudo_gradient_term_with_zero_nexp(self): ni = mf._numint dm0 = mf.make_rdm1() - rho_g = multigrid_v2.evaluate_density_on_g_mesh(ni, dm0, kpts) - rho_g = rho_g[0,0] - dx = 1e-5 + ni = multigrid_v3.MultiGridNumInt(cell) + ni.allow_mesh_reduction = False + Gv_bases = _get_Gv_bases(ni.mesh, cell.reciprocal_vectors()) + rho_g = multigrid_v3._eval_rhoG(ni, dm0, 1, kpts).ravel() + analytical_gradient = _pploc_derivatives(cell, ni.mesh, rho_g, Gv_bases)[0].get() + + dx = 1e-4 numerical_gradient = np.zeros([cell.natm, 3]) def get_pp_local_energy(cell): - vpplocG = multigrid_v1.eval_vpplocG(cell, cell.mesh) + vpplocG = multigrid_v1.eval_vpplocG(cell, ni.mesh) e_vpplocG = cp.einsum("g,g->", rho_g.conj(), vpplocG) / cell.vol assert abs(e_vpplocG.imag) < 1e-8 return float(e_vpplocG.real) @@ -412,9 +417,6 @@ def get_pp_local_energy(cell): numerical_gradient[i_atom, i_xyz] = (e_p - e_m) / (2 * dx) - analytical_gradient = multigrid_v1.eval_vpplocG_SI_gradient(cell, cell.mesh, rho_g) - analytical_gradient = analytical_gradient.get() - assert np.max(np.abs(numerical_gradient - analytical_gradient)) < 1e-8 diff --git a/gpu4pyscf/pbc/grad/tests/test_pbc_rks_stress.py b/gpu4pyscf/pbc/grad/tests/test_pbc_rks_stress.py index 653b254cb..82c466efa 100644 --- a/gpu4pyscf/pbc/grad/tests/test_pbc_rks_stress.py +++ b/gpu4pyscf/pbc/grad/tests/test_pbc_rks_stress.py @@ -15,6 +15,9 @@ import unittest import numpy as np +import cupy as cp +import pyscf +from pyscf.pbc.gto import pseudo from pyscf.pbc import dft, gto from pyscf.pbc.tools import pbc from pyscf.pbc.df import FFTDF @@ -200,6 +203,7 @@ def test_get_nuc(self): assert abs(dat[i,j] - de/2e-5) < 1e-8 def test_get_pp(self): + from gpu4pyscf.pbc.grad.rks_stress import _get_pp_nonloc_strain_derivatives a = np.eye(3) * 5 np.random.seed(5) a += np.random.rand(3, 3) - .5 @@ -212,10 +216,11 @@ def test_get_pp(self): xc = 'lda,' mf_grad = rks.Gradients(cell.RKS(xc=xc).to_gpu()) dat = rks_stress.get_vxc(mf_grad, cell, dm, with_nuc=True) + dat += _get_pp_nonloc_strain_derivatives(cell, cell.mesh, cp.array(dm)) ni = NumInt() kpt = np.zeros(3) for (i, j) in [(0, 0), (0, 1), (0, 2), (2, 1), (2, 2)]: - cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-5) + cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-4) cell1.precision = 1e-10 cell2.precision = 1e-10 vne1 = FFTDF(cell1).get_pp(kpt) @@ -224,7 +229,7 @@ def test_get_pp(self): exc2 = ni.nr_rks(cell2, UniformGrids(cell2), xc, dm)[1] de = np.einsum('ij,ji', dm, (vne1-vne2)) de += exc1 - exc2 - assert abs(dat[i,j] - de/2e-5) < 1e-8 + assert abs(dat[i,j] - de/2e-4) < 5e-7 def test_lda_vs_finite_difference(self): a = np.eye(3) * 3 @@ -325,6 +330,32 @@ def test_hse_vs_finite_difference(self): e2 = mf_scanner(cell2) assert abs(dat[i,j] - (e1-e2)/2e-3/vol) < 2e-7 + def test_get_vpplocG_strain_derivatives(self): + from gpu4pyscf.pbc.grad.rks_stress import _get_vpplocG_strain_derivatives + np.random.seed(8) + cell = pyscf.M( + atom='C 0 0 0;#C .2 .3 .7', + basis=[[0, [0.4, 1]]], + pseudo={'C': [[2, 2], 0.38, 4, [-8.8, 1.33, 0.85, 0.55]]}, + a=np.eye(3) * 2.5 + np.random.rand(3,3)*.5) + mesh = [11] * 3 + + disp = 1e-5 + ngrids = np.prod(mesh) + ref = cp.empty((3,3, ngrids), dtype=np.complex128) + SI = cp.array(cell.get_SI(mesh=mesh)) + for x in range(3): + for y in range(3): + cell1, cell2 = _finite_diff_cells(cell, x, y, disp) + vpplocG1 = pseudo.get_vlocG(cell1, cell1.get_Gv(mesh)) + vpplocG2 = pseudo.get_vlocG(cell2, cell2.get_Gv(mesh)) + vpplocG1 = -np.einsum('ij,ij->j', SI, vpplocG1) + vpplocG2 = -np.einsum('ij,ij->j', SI, vpplocG2) + ref[x,y] = cp.asarray((vpplocG1 - vpplocG2) / (2*disp)) + + dat = _get_vpplocG_strain_derivatives(cell, mesh)[1] + assert abs(dat - ref).max().get() < 1e-7 + if __name__ == "__main__": print("Full Tests for RKS Stress tensor") - unittest.main() + #unittest.main() diff --git a/gpu4pyscf/pbc/grad/uhf.py b/gpu4pyscf/pbc/grad/uhf.py index 85c20b485..2674d5028 100644 --- a/gpu4pyscf/pbc/grad/uhf.py +++ b/gpu4pyscf/pbc/grad/uhf.py @@ -23,8 +23,7 @@ import gpu4pyscf.grad.uhf as mol_uhf import gpu4pyscf.pbc.grad.rhf as rhf from gpu4pyscf.pbc.tools.k2gamma import kpts_to_kmesh -from gpu4pyscf.pbc.dft import multigrid_v2 -import gpu4pyscf.pbc.dft.multigrid as multigrid_v1 +from gpu4pyscf.pbc.dft import multigrid from gpu4pyscf.pbc.scf.rsjk import PBCJKMatrixOpt from gpu4pyscf.pbc.df.df import GDF from gpu4pyscf.pbc.gto import int1e @@ -49,7 +48,7 @@ def energy_ee(self, dm): # pseudo+GGA does not need to evaluate the gradients with PBCJKMatrixOpt ni = mf._numint - j_in_xc = isinstance(ni, multigrid_v2.MultiGridNumInt) + j_in_xc = isinstance(ni, multigrid.MultiGridNumIntBase) de = 0 xc = getattr(mf, 'xc', 'HF') if xc.upper() == 'HF': @@ -60,9 +59,9 @@ def energy_ee(self, dm): j_factor = 1 if j_in_xc: - de += multigrid_v2.get_veff_ip1( - ni, xc, dm, with_j=j_in_xc, - with_pseudo_vloc_orbital_derivative=True).get() + assert isinstance(ni, multigrid.MultiGridNumIntBase) + de += ni.energy_nuclear_gradient( + xc, dm, spin=1, with_j=j_in_xc, with_nuc=True) j_factor = 0 elif xc.upper() != 'HF': from gpu4pyscf.pbc.grad.kuks import get_vxc @@ -101,14 +100,8 @@ def grad_elec( de = self.energy_ee(dm0) ni = mf._numint - if isinstance(ni, multigrid_v2.MultiGridNumInt): - rhoG = multigrid_v2.evaluate_density_on_g_mesh(ni, dm0_sf) - rhoG = rhoG[0,0] - if cell._pseudo: - de += multigrid_v1.eval_vpplocG_SI_gradient(cell, ni.mesh, rhoG).get() - else: - de += multigrid_v1.eval_nucG_SI_gradient(cell, ni.mesh, rhoG).get() - + if isinstance(ni, multigrid.MultiGridNumIntBase): + # Vne or pploc contribution is evaluated in energy_ee dh1e_kin = int1e.int1e_ipkin(cell) de -= rhf.contract_h1e_dm(cell, dh1e_kin, dm0_sf, hermi=1) else: diff --git a/gpu4pyscf/pbc/grad/uhf_stress.py b/gpu4pyscf/pbc/grad/uhf_stress.py index 708cbcc3c..18c6f317d 100644 --- a/gpu4pyscf/pbc/grad/uhf_stress.py +++ b/gpu4pyscf/pbc/grad/uhf_stress.py @@ -22,8 +22,8 @@ from gpu4pyscf.lib import logger from gpu4pyscf.pbc.grad import uhf as uhf_grad from gpu4pyscf.pbc.gto import int1e -from gpu4pyscf.pbc.grad.rks_stress import _finite_diff_cells, ewald -from gpu4pyscf.pbc.grad.rhf_stress import get_nuc, get_veff +from gpu4pyscf.pbc.grad.rhf_stress import ( + get_nuc, get_veff, _get_pp_nonloc_strain_derivatives, ewald) def kernel(mf_grad): '''Compute the energy derivatives for strain tensor (e_ij) @@ -59,8 +59,10 @@ def kernel(mf_grad): dm0_sf = dm0[0] + dm0[1] dme0_sf = dme0[0] + dme0[1] sigma -= int1e.ovlp_strain_deriv(cell, dme0_sf) - sigma += int1e.kin_strain_deriv(cell, dm0) + sigma += int1e.kin_strain_deriv(cell, dm0_sf) sigma += get_nuc(mf_grad, cell, dm0_sf) + if cell._pseudo: + sigma += _get_pp_nonloc_strain_derivatives(cell, cell.mesh, dm0_sf) t0 = log.timer_debug1('hcore derivatives', *t0) sigma += get_veff(mf_grad, cell, dm0) diff --git a/gpu4pyscf/pbc/grad/uks_stress.py b/gpu4pyscf/pbc/grad/uks_stress.py index 62fb0d7a2..ee67ddd65 100644 --- a/gpu4pyscf/pbc/grad/uks_stress.py +++ b/gpu4pyscf/pbc/grad/uks_stress.py @@ -23,7 +23,7 @@ from gpu4pyscf.pbc.df import FFTDF from gpu4pyscf.pbc.df.aft import _get_ZSI from gpu4pyscf.pbc.dft.numint import NumInt, eval_ao_kpts, _GTOvalOpt -from gpu4pyscf.pbc.dft.multigrid_v2 import _uks_exc_strain_deriv, MultiGridNumInt +from gpu4pyscf.pbc.dft.multigrid_v3 import MultiGridNumInt from gpu4pyscf.pbc.grad import uks as uks_grad from gpu4pyscf.pbc.gto import int1e from gpu4pyscf.pbc.scf.rsjk import PBCJKMatrixOpt @@ -54,7 +54,8 @@ def get_veff(mf_grad, cell, dm, with_j=False, with_nuc=False): # TODO: with_nuc should be disabled for all-electron calculations if isinstance(ni, MultiGridNumInt): - sigma = _uks_exc_strain_deriv(ni, mf.xc, dm[:,None], None, with_j, with_nuc) + sigma = ni.energy_strain_gradient(mf.xc, dm[:,None], spin=1, + with_j=with_j, with_nuc=with_nuc) elif isinstance(ni, NumInt): sigma = get_vxc(mf_grad, cell, dm, with_j, with_nuc) else: @@ -230,7 +231,6 @@ def partial_dot(bra, ket): vpplocR = pbctools.ifft(vpplocG_0, mesh).real Ene = contract('xyg,g->xy', rho1, vpplocR).real.get() Ene += contract('g,xyg->xy', rhoG.conj(), vpplocG_1).real.get() * (1./ngrids) - Ene += _get_pp_nonloc_strain_derivatives(cell, mesh, dm.sum(axis=0)) else: # SI corresponds to Fourier components of the fractional atomic # positions within the cell. It does not respond to the strain @@ -273,23 +273,17 @@ def kernel(mf_grad): log.debug('Computing stress tensor') cell = mf.cell - dm0 = mf.make_rdm1().sum(axis=0) + dm0 = mf.make_rdm1() + dm0_sf = dm0[0] + dm0[1] dme0 = mf_grad.make_rdm1e().sum(axis=0) sigma = ewald(cell) sigma -= int1e.ovlp_strain_deriv(cell, dme0) - - disp = 1e-5 - for x in range(3): - for y in range(3): - cell1, cell2 = _finite_diff_cells(cell, x, y, disp) - t1 = int1e.int1e_kin(cell1) - t2 = int1e.int1e_kin(cell2) - t1 = cp.einsum('ij,ji->', t1, dm0) - t2 = cp.einsum('ij,ji->', t2, dm0) - sigma[x,y] += (t1 - t2) / (2*disp) + sigma += int1e.kin_strain_deriv(cell, dm0_sf) + if cell._pseudo: + # pploc contribution is evaluated in get_veff + sigma += _get_pp_nonloc_strain_derivatives(cell, cell.mesh, dm0_sf) t0 = log.timer_debug1('hcore derivatives', *t0) - dm0 = mf.make_rdm1() sigma += get_veff(mf_grad, cell, dm0, with_j=True, with_nuc=True) t0 = log.timer_debug1('Vxc and Coulomb derivatives', *t0) diff --git a/gpu4pyscf/pbc/gto/int1e.py b/gpu4pyscf/pbc/gto/int1e.py index 187529234..aa3f86553 100644 --- a/gpu4pyscf/pbc/gto/int1e.py +++ b/gpu4pyscf/pbc/gto/int1e.py @@ -76,7 +76,7 @@ def int1e_ipovlp(cell, kpts=None, bvk_kmesh=None, sort_output=True): return opt.intor('PBCint1e_ipovlp', 3, (1, 0), kpts, sort_output) def int1e_ipkin(cell, kpts=None, bvk_kmesh=None, sort_output=True): - opt = _check_opt(cell, 2, kpts, bvk_kmesh) + opt = _check_opt(cell, 2, kpts, bvk_kmesh, 1e-2) return opt.intor('PBCint1e_ipkin', 3, (3, 0), kpts, sort_output) def int1e_r2_origi(cell, kpts=None, bvk_kmesh=None, sort_output=True): @@ -144,7 +144,7 @@ def __init__(self, cell, hermi=0, bvk_kmesh=None): bvkcell = cell else: bvkcell = super_cell(cell, bvk_kmesh, wrap_around=True) - # PTR_BAS_COORD was not initialized in supe_rcell + # PTR_BAS_COORD was not initialized in the super_cell function bvkcell._bas[:,PTR_BAS_COORD] = bvkcell._atm[bvkcell._bas[:,ATOM_OF],PTR_COORD] Ls = asarray(bvkcell.get_lattice_Ls(rcut=cell.rcut)) Ls = Ls[cp.linalg.norm(Ls-.5, axis=1).argsort()] diff --git a/gpu4pyscf/pbc/gto/pseudo/pp_int.py b/gpu4pyscf/pbc/gto/pseudo/pp_int.py index ba0b7ca1c..3019b328b 100644 --- a/gpu4pyscf/pbc/gto/pseudo/pp_int.py +++ b/gpu4pyscf/pbc/gto/pseudo/pp_int.py @@ -111,7 +111,7 @@ def _contract_ppnl_gpu(cell, fakecell, hl_blocks, ppnl_half, comp=1, kpts=None): ppnl_k += cp.einsum('imp,ij,jmq->pq', ilp, hl_gpu, ilp) - ppnl.append(ppnl_k.get()) + ppnl.append(ppnl_k) if kpts is None or np.shape(kpts) == (3,): return ppnl[0] @@ -154,4 +154,4 @@ def get_pp_nl_gpu(cell, kpts=None): ppnl[k] += cp.einsum('ilp,ij,jlq->pq', ilp.conj(), hl_gpu, ilp) - return ppnl.get() + return ppnl diff --git a/gpu4pyscf/pbc/gto/pseudo/tests/test_pp_int.py b/gpu4pyscf/pbc/gto/pseudo/tests/test_pp_int.py index 439034507..f4ee42eb0 100644 --- a/gpu4pyscf/pbc/gto/pseudo/tests/test_pp_int.py +++ b/gpu4pyscf/pbc/gto/pseudo/tests/test_pp_int.py @@ -25,6 +25,7 @@ import unittest import numpy as np +import cupy as cp import pyscf from pyscf.pbc.gto.pseudo.pp_int import ( fake_cell_vnl, _int_vnl, _contract_ppnl, get_pp_nl) @@ -76,7 +77,7 @@ def _compare(self, cell, places=13): cpu = _contract_ppnl(cell, fakecell, hl_blocks, ppnl_half, kpts=kpts) gpu = _contract_ppnl_gpu(cell, fakecell, hl_blocks, ppnl_half, kpts=kpts) - err = np.max(np.abs(np.asarray(gpu) - np.asarray(cpu))) + err = np.max(np.abs(cp.asarray(gpu).get() - np.asarray(cpu))) self.assertAlmostEqual(err, 0, places, f"max|err|={err:.2e}") def test_carbon(self): @@ -96,7 +97,7 @@ def _compare(self, cell, places=12): from gpu4pyscf.pbc.gto.pseudo.pp_int import get_pp_nl_gpu cpu = get_pp_nl(cell) gpu = get_pp_nl_gpu(cell) - err = np.max(np.abs(np.asarray(gpu) - np.asarray(cpu))) + err = np.max(np.abs(cp.asnumpy(gpu) - np.asarray(cpu))) self.assertAlmostEqual(err, 0, places, f"max|err|={err:.2e}") def test_carbon(self): @@ -116,7 +117,7 @@ def _compare(self, cell, kpts, places=13): from gpu4pyscf.pbc.gto.pseudo.pp_int import get_pp_nl_gpu cpu = get_pp_nl(cell, kpts) gpu = get_pp_nl_gpu(cell, kpts) - err = np.max(np.abs(np.asarray(gpu) - np.asarray(cpu))) + err = np.max(np.abs(cp.asnumpy(gpu) - np.asarray(cpu))) self.assertAlmostEqual(err, 0, places, f"max|err|={err:.2e}") def test_silicon_single_kpt(self): diff --git a/gpu4pyscf/pbc/gto/tests/test_pbc_int1e.py b/gpu4pyscf/pbc/gto/tests/test_pbc_int1e.py index 2934760b2..999170570 100644 --- a/gpu4pyscf/pbc/gto/tests/test_pbc_int1e.py +++ b/gpu4pyscf/pbc/gto/tests/test_pbc_int1e.py @@ -36,13 +36,13 @@ def test_int1e_ovlp(): ref = pcell.pbc_intor('int1e_ovlp', hermi=1, kpts=kpts) dat = int1e.int1e_ovlp(cell).get() - assert abs(dat - ref[0]).max() < 1e-13 + assert abs(dat - ref[0]).max() < 1e-12 dat = int1e.int1e_ovlp(cell, kpts, kmesh).get() - assert abs(dat - ref).max() < 1e-13 + assert abs(dat - ref).max() < 1e-12 dat = int1e.int1e_ovlp(cell, kpts).get() - assert abs(dat - ref).max() < 1e-13 + assert abs(dat - ref).max() < 1e-12 cell = pyscf.M( atom='''C1 1.3 .2 .3 @@ -75,13 +75,13 @@ def test_int1e_kin(): ref = pcell.pbc_intor('int1e_kin', hermi=1, kpts=kpts) dat = int1e.int1e_kin(cell).get() - assert abs(dat - ref[0]).max() < 1e-10 + assert abs(dat - ref[0]).max() < 1e-9 dat = int1e.int1e_kin(cell, kpts, kmesh).get() - assert abs(dat - ref).max() < 1e-10 + assert abs(dat - ref).max() < 1e-9 dat = int1e.int1e_kin(cell, kpts).get() - assert abs(dat - ref).max() < 1e-10 + assert abs(dat - ref).max() < 1e-9 mol = cell.to_mol() dat = int1e.int1e_kin(mol).get() @@ -106,13 +106,13 @@ def test_int1e_ipovlp(): ref = np.asarray(pcell.pbc_intor('int1e_ipovlp', hermi=0, kpts=kpts)) dat = int1e.int1e_ipovlp(cell).get() - assert abs(dat - ref[0]).max() < 1e-8 + assert abs(dat - ref[0]).max() < 5e-8 dat = int1e.int1e_ipovlp(cell, kpts, kmesh).get() - assert abs(dat - ref).max() < 1e-8 + assert abs(dat - ref).max() < 5e-8 dat = int1e.int1e_ipovlp(cell, kpts).get() - assert abs(dat - ref).max() < 1e-8 + assert abs(dat - ref).max() < 5e-8 mol = cell.to_mol() dat = int1e.int1e_ipovlp(mol).get() @@ -137,13 +137,13 @@ def test_int1e_ipkin(): ref = np.asarray(pcell.pbc_intor('int1e_ipkin', hermi=0, kpts=kpts)) dat = int1e.int1e_ipkin(cell).get() - assert abs(dat - ref[0]).max() < 1e-8 + assert abs(dat - ref[0]).max() < 5e-8 dat = int1e.int1e_ipkin(cell, kpts, kmesh).get() - assert abs(dat - ref).max() < 1e-8 + assert abs(dat - ref).max() < 5e-8 dat = int1e.int1e_ipkin(cell, kpts).get() - assert abs(dat - ref).max() < 1e-8 + assert abs(dat - ref).max() < 5e-8 def test_int1e_ovlp1(): L = 4 diff --git a/gpu4pyscf/pbc/scf/hf.py b/gpu4pyscf/pbc/scf/hf.py index ef3f8ebff..3e3a73669 100644 --- a/gpu4pyscf/pbc/scf/hf.py +++ b/gpu4pyscf/pbc/scf/hf.py @@ -49,7 +49,7 @@ def get_bands(mf, kpts_band, cell=None, dm=None, kpt=None): single_kpt_band = (getattr(kpts_band, 'ndim', None) == 1) kpts_band = kpts_band.reshape(-1,3) - fock = cp.asarray(mf.get_veff(cell, dm, kpt=kpt, kpts_band=kpts_band)) + fock = mf.get_veff(cell, dm, kpt=kpt, kpts_band=kpts_band) fock += mf.get_hcore(cell, kpts_band) s1e = mf.get_ovlp(cell, kpts_band) mo_energy, mo_coeff = eigh_with_canonical_orth(fock, s1e) @@ -88,7 +88,7 @@ def get_rho(mf, dm=None, grids=None, kpt=None): '''Compute density in real space ''' from gpu4pyscf.pbc.dft import UniformGrids - from gpu4pyscf.pbc.dft import numint, multigrid, multigrid_v2 + from gpu4pyscf.pbc.dft import numint, multigrid if dm is None: dm = mf.make_rdm1() if getattr(dm, 'ndim', None) != 2: # UHF @@ -99,7 +99,7 @@ def get_rho(mf, dm=None, grids=None, kpt=None): ni = mf._numint if ni is None: ni = numint.NumInt() - if isinstance(ni, (multigrid.MultiGridNumInt, multigrid_v2.MultiGridNumInt)): + if isinstance(ni, multigrid.MultiGridNumIntBase): assert grids is None or isinstance(grids, UniformGrids) if grids is not None and any(grids.mesh != ni.mesh): ni = ni.copy().reset() @@ -215,8 +215,8 @@ def build(self, cell=None): # MultiGridNumInt integrator to evaluate Coulomb integrals, skipping the # self.with_df code path. if isinstance(self.with_df, df.FFTDF) and self._numint is None: - from gpu4pyscf.pbc.dft import multigrid_v2 - self._numint = multigrid_v2.MultiGridNumInt(self.cell) + from gpu4pyscf.pbc.dft import multigrid_v3 + self._numint = multigrid_v3.MultiGridNumInt(self.cell) if self.verbose >= logger.WARN: self.check_sanity() @@ -234,15 +234,15 @@ def get_ovlp(self, cell=None, kpt=None): return int1e.int1e_ovlp(cell, kpt) def get_hcore(self, cell=None, kpt=None): - from gpu4pyscf.pbc.dft import multigrid, multigrid_v2 + from gpu4pyscf.pbc.dft import multigrid, multigrid_v3 if cell is None: cell = self.cell if kpt is None: kpt = self.kpt - if isinstance(self._numint, (multigrid.MultiGridNumInt, multigrid_v2.MultiGridNumInt)): + if isinstance(self._numint, multigrid.MultiGridNumIntBase): ni = self._numint - elif np.prod(cell.mesh) < 500**3: + elif np.prod(cell.mesh) < 1000**3: # In the pseudo and all-electron mixed case, MultiGridNumInt is # still more efficient if Ecut is not too high. - ni = multigrid_v2.MultiGridNumInt(cell) + ni = multigrid_v3.MultiGridNumInt(cell) else: ni = self.with_df if cell.pseudo: @@ -362,8 +362,8 @@ def multigrid_numint(self, mesh=None): raise NotImplementedError def dump_chk(self, envs): - mol_hf.SCF.dump_chk(self, envs) if self.chkfile: + mol_hf.SCF.dump_chk(self, envs) with lib.H5FileWrap(self.chkfile, 'a') as fh5: fh5['scf/kpt'] = self.kpt return self diff --git a/gpu4pyscf/pbc/scf/khf.py b/gpu4pyscf/pbc/scf/khf.py index 0108bff8e..ba4574f2a 100644 --- a/gpu4pyscf/pbc/scf/khf.py +++ b/gpu4pyscf/pbc/scf/khf.py @@ -259,7 +259,7 @@ def get_rho(mf, dm=None, grids=None, kpts=None): '''Compute density in real space ''' from gpu4pyscf.pbc.dft import UniformGrids - from gpu4pyscf.pbc.dft import numint, multigrid, multigrid_v2 + from gpu4pyscf.pbc.dft import numint, multigrid if dm is None: dm = mf.make_rdm1() if getattr(dm[0], 'ndim', None) != 2: # KUHF @@ -270,7 +270,7 @@ def get_rho(mf, dm=None, grids=None, kpts=None): ni = mf._numint if ni is None: ni = numint.KNumInt() - if isinstance(ni, (multigrid.MultiGridNumInt, multigrid_v2.MultiGridNumInt)): + if isinstance(ni, multigrid.MultiGridNumIntBase): assert grids is None or isinstance(grids, UniformGrids) if grids is not None and any(grids.mesh != ni.mesh): ni = ni.copy().reset() @@ -372,8 +372,8 @@ def build(self, cell=None): # MultiGridNumInt integrator to evaluate Coulomb integrals, skipping the # self.with_df code path. if isinstance(self.with_df, df.FFTDF) and self._numint is None: - from gpu4pyscf.pbc.dft import multigrid_v2 - self._numint = multigrid_v2.MultiGridNumInt(self.cell) + from gpu4pyscf.pbc.dft import multigrid_v3 + self._numint = multigrid_v3.MultiGridNumInt(self.cell) if self.verbose >= logger.WARN: self.check_sanity() @@ -392,19 +392,19 @@ def get_ovlp(self, cell=None, kpts=None): return int1e.int1e_ovlp(cell, kpts, bvk_kmesh) def get_hcore(self, cell=None, kpts=None): - from gpu4pyscf.pbc.dft import multigrid, multigrid_v2 + from gpu4pyscf.pbc.dft import multigrid, multigrid_v3 if cell is None: cell = self.cell if kpts is None: kpts = self.kpts kpts_in_bvkcell = True else: kpts_in_bvkcell = len(kpts) == len(self.kpts) - if isinstance(self._numint, (multigrid.MultiGridNumInt, multigrid_v2.MultiGridNumInt)): + if isinstance(self._numint, multigrid.MultiGridNumIntBase): ni = self._numint - elif np.prod(cell.mesh) < 500**3: + elif np.prod(cell.mesh) < 1000**3: # In the pseudo and all-electron mixed case, MultiGridNumInt is # still more efficient if Ecut is not too high. - ni = multigrid_v2.MultiGridNumInt(cell) + ni = multigrid_v3.MultiGridNumInt(cell) else: ni = self.with_df if cell.pseudo: @@ -664,8 +664,8 @@ def get_bands(self, kpts_band, cell=None, dm_kpts=None, kpts=None): smearing = pbchf.SCF.smearing def dump_chk(self, envs): - mol_hf.SCF.dump_chk(self, envs) if self.chkfile: + mol_hf.SCF.dump_chk(self, envs) with lib.H5FileWrap(self.chkfile, 'a') as fh5: fh5['scf/kpts'] = cp.asnumpy(self.kpts) return self diff --git a/gpu4pyscf/pbc/scf/tests/test_pbc_scf_diffuse_orbital.py b/gpu4pyscf/pbc/scf/tests/test_pbc_scf_diffuse_orbital.py index 0d302bf47..be34cfff5 100644 --- a/gpu4pyscf/pbc/scf/tests/test_pbc_scf_diffuse_orbital.py +++ b/gpu4pyscf/pbc/scf/tests/test_pbc_scf_diffuse_orbital.py @@ -15,7 +15,7 @@ import unittest import numpy as np from pyscf.pbc import gto as pbcgto -from gpu4pyscf.pbc.dft.multigrid_v2 import MultiGridNumInt +from gpu4pyscf.pbc.dft.multigrid_v3 import MultiGridNumInt from gpu4pyscf.pbc.df import AFTDF import gpu4pyscf from gpu4pyscf.lib.multi_gpu import num_devices @@ -94,8 +94,8 @@ def test_rks(self): ref_gradient = np.array([[-3.75319097e-02, 9.32974780e-03, -4.73203810e-11], [ 3.75319095e-02, -9.32974768e-03, -4.51863943e-11]]) - assert abs(test_energy - ref_energy) < 2e-10 - assert np.max(np.abs(test_gradient - ref_gradient)) < 1e-7 + assert abs(test_energy - ref_energy) < 3e-9 + assert np.max(np.abs(test_gradient - ref_gradient)) < 1e-6 def test_krks(self): cell = self.cell @@ -121,14 +121,14 @@ def test_krks(self): ref_gradient = np.array([[ 2.77114603e-02, -2.47140931e-02, 3.94331325e-10], [-2.77114767e-02, 2.47502293e-02, 4.16056183e-10]]) - assert abs(test_energy - ref_energy) < 2e-10 - assert np.max(np.abs(test_gradient - ref_gradient)) < 1e-7 + assert abs(test_energy - ref_energy) < 1e-9 + assert np.max(np.abs(test_gradient - ref_gradient)) < 1e-6 e, c = mf.canonicalize(mf.mo_coeff, mf.mo_occ) assert abs(e[e<1e7] - mf.mo_energy[e<1e7]).max() < 5e-7 f = mf.get_fock() e1 = contract('kqi,kqi->ki', contract('kpi,kpq->kqi', c.conj(), f), c) - assert abs(e[e<1e7] - e1[e<1e7]).max() < 2e-10 + assert abs(e[e<1e7] - e1[e<1e7]).max() < 1e-9 @unittest.skipIf(num_devices > 1, '') def test_krks_aftdf(self): @@ -149,7 +149,7 @@ def test_krks_aftdf(self): ref_energy = -1.439218776262645 - assert abs(test_energy - ref_energy) < 2e-10 + assert abs(test_energy - ref_energy) < 1e-9 def test_uks(self): cell = self.cell @@ -174,8 +174,8 @@ def test_uks(self): ref_gradient = np.array([[-4.10163065e-02, 1.00148573e-02, -4.86782781e-11], [ 4.10163065e-02, -1.00148573e-02, -4.62222600e-11]]) - assert abs(test_energy - ref_energy) < 2e-10 - assert np.max(np.abs(test_gradient - ref_gradient)) < 1e-7 + assert abs(test_energy - ref_energy) < 3e-9 + assert np.max(np.abs(test_gradient - ref_gradient)) < 1e-6 def test_kuks(self): cell = self.cell @@ -220,8 +220,8 @@ def test_rks_smearing(self): ref_gradient = np.array([[-3.75888726e-02, 9.34455109e-03, -4.80129776e-11], [ 3.75888723e-02, -9.34455096e-03, -4.57646238e-11]]) - assert abs(test_energy - ref_energy) < 2e-10 - assert np.max(np.abs(test_gradient - ref_gradient)) < 1e-7 + assert abs(test_energy - ref_energy) < 3e-9 + assert np.max(np.abs(test_gradient - ref_gradient)) < 1e-6 @unittest.skipIf(num_devices > 1, '') def test_krks_smearing(self): @@ -240,8 +240,8 @@ def test_krks_smearing(self): ref_gradient = np.array([[-2.80927897e-04, -1.55260904e-02, -7.70325423e-11], [ 2.80905389e-04, 1.55173700e-02, -8.01210477e-11]]) - assert abs(test_energy - ref_energy) < 2e-10 - assert np.max(np.abs(test_gradient - ref_gradient)) < 1e-7 + assert abs(test_energy - ref_energy) < 1e-9 + assert np.max(np.abs(test_gradient - ref_gradient)) < 1e-6 @unittest.skipIf(num_devices > 1, '') def test_uks_smearing(self): @@ -259,8 +259,8 @@ def test_uks_smearing(self): ref_gradient = np.array([[-4.10910618e-02, 1.00340747e-02, -4.86288417e-11], [ 4.10910618e-02, -1.00340747e-02, -4.60603924e-11]]) - assert abs(test_energy - ref_energy) < 2e-10 - assert np.max(np.abs(test_gradient - ref_gradient)) < 1e-7 + assert abs(test_energy - ref_energy) < 3e-9 + assert np.max(np.abs(test_gradient - ref_gradient)) < 1e-6 @unittest.skipIf(num_devices > 1, '') def test_kuks_smearing(self): diff --git a/gpu4pyscf/pbc/scf/tests/test_pbc_scf_hf.py b/gpu4pyscf/pbc/scf/tests/test_pbc_scf_hf.py index af3759ef8..cda9399f9 100644 --- a/gpu4pyscf/pbc/scf/tests/test_pbc_scf_hf.py +++ b/gpu4pyscf/pbc/scf/tests/test_pbc_scf_hf.py @@ -63,8 +63,8 @@ def test_rhf_exx_ewald(self): e1, c1 = mf.get_bands(kpts_band) e0, c0 = kmf.get_bands(kpts_band) self.assertAlmostEqual(abs(e0-e1).get().max(), 0, 7) - self.assertAlmostEqual(lib.fp(e1[0].get()), -6.2986775452228283, 6) - self.assertAlmostEqual(lib.fp(e1[1].get()), -7.6616273746782362, 6) + self.assertAlmostEqual(lib.fp(e1[0].get()), -6.2986775452228283, delta=1e-5) + self.assertAlmostEqual(lib.fp(e1[1].get()), -7.6616273746782362, delta=1e-5) def test_rhf_exx_ewald_with_kpt(self): np.random.seed(1) @@ -73,13 +73,13 @@ def test_rhf_exx_ewald_with_kpt(self): kmf = scf.KRHF(cell, k, exxdiv='ewald') kmf.time_reversal_symmetry = False e0 = kmf.kernel() - self.assertAlmostEqual(e0, -4.2048655827967139, 7) + self.assertAlmostEqual(e0, -4.2048655827967139, delta=1e-6) # test bands np.random.seed(1) kpt_band = np.random.random(3) e0, c0 = kmf.get_bands(kpt_band) - self.assertAlmostEqual(lib.fp(e0.get()), -6.8312867098806249, 6) + self.assertAlmostEqual(lib.fp(e0.get()), -6.8312867098806249, delta=1e-5) def test_rhf_exx_None(self): cell = self.cell @@ -98,7 +98,7 @@ def test_rhf_exx_None(self): kmf.time_reversal_symmetry = False kmf.init_guess = 'hcore' e0 = kmf.kernel() - self.assertAlmostEqual(e0, -2.7862168430230341, 7) + self.assertAlmostEqual(e0, -2.7862168430230341, delta=1e-6) def test_jk(self): cell = self.cell @@ -128,7 +128,7 @@ def test_krhf_bands(self): kpts_bands = np.random.random((1,3)) e = kmf.get_bands(kpts_bands)[0] e_ref = kmf_cpu.get_bands(kpts_bands)[0] - self.assertAlmostEqual(abs(e.get()-e_ref).max(), 0, 6) + self.assertAlmostEqual(abs(e.get()-e_ref).max(), 0, delta=1e-5) def test_density_fit(self): from gpu4pyscf.pbc.df.df import GDF diff --git a/gpu4pyscf/pbc/scf/tests/test_pbc_scf_uhf.py b/gpu4pyscf/pbc/scf/tests/test_pbc_scf_uhf.py index d1b820729..8e8cd440e 100644 --- a/gpu4pyscf/pbc/scf/tests/test_pbc_scf_uhf.py +++ b/gpu4pyscf/pbc/scf/tests/test_pbc_scf_uhf.py @@ -71,7 +71,7 @@ def test_uhf_bands(self): self.assertAlmostEqual(lib.fp(pop), -0.04691820429296646, 5) np.random.seed(1) - kpts_bands = np.random.random((1,3)) + kpts_bands = np.random.random((4,3)) e = mf.get_bands(kpts_bands)[0] e_ref = mf_cpu.get_bands(kpts_bands)[0] self.assertAlmostEqual(abs(e.get()-e_ref).max(), 0, 6) diff --git a/gpu4pyscf/pbc/scf/tests/test_pbc_soscf.py b/gpu4pyscf/pbc/scf/tests/test_pbc_soscf.py index 0eab51dd6..7e76255ac 100644 --- a/gpu4pyscf/pbc/scf/tests/test_pbc_soscf.py +++ b/gpu4pyscf/pbc/scf/tests/test_pbc_soscf.py @@ -65,7 +65,7 @@ def test_nr_rhf(self): mf.conv_tol_grad = 1e-4 mf.kernel() self.assertAlmostEqual(mf.e_tot, ref.e_tot, 8) - self.assertAlmostEqual(mf.e_tot, -9.870755717258616, 8) + self.assertAlmostEqual(mf.e_tot, -9.870755717258616, 7) def test_nr_uhf(self): mf = scf.UHF(cell).density_fit() @@ -74,7 +74,7 @@ def test_nr_uhf(self): mf.conv_tol_grad = 1e-4 mf.kernel() self.assertAlmostEqual(mf.e_tot, ref.e_tot, 8) - self.assertAlmostEqual(mf.e_tot, -9.87737235242746, 8) + self.assertAlmostEqual(mf.e_tot, -9.87737235242746, 7) def test_nr_rks_lda(self): mf = dft.RKS(cell, xc='lda,') @@ -84,7 +84,7 @@ def test_nr_rks_lda(self): mf.conv_tol_grad = 1e-4 mf.kernel() self.assertAlmostEqual(mf.e_tot, ref.e_tot, 8) - self.assertAlmostEqual(mf.e_tot, -9.454224868598352, 8) + self.assertAlmostEqual(mf.e_tot, -9.454224868598352, 7) def test_nr_uks_lda(self): mf = dft.RKS(cell, xc='lda,') @@ -113,7 +113,7 @@ def test_nr_uks_gga(self): mf.conv_tol_grad = 1e-4 mf.kernel() self.assertAlmostEqual(mf.e_tot, ref.e_tot, 8) - self.assertAlmostEqual(mf.e_tot, -10.021211009946274, 8) + self.assertAlmostEqual(mf.e_tot, -10.021211009946274, 7) def test_nr_uks_rsh(self): mf = dft.UKS(cell, xc='camb3lyp') @@ -124,7 +124,7 @@ def test_nr_uks_rsh(self): mf.conv_tol_grad = 1e-4 mf.kernel() self.assertAlmostEqual(mf.e_tot, ref.e_tot, 8) - self.assertAlmostEqual(mf.e_tot, -10.04741991140345, 8) + self.assertAlmostEqual(mf.e_tot, -10.04741991140345, 7) def test_nr_krhf(self): mf = scf.KRHF(cell, cell.make_kpts([2,1,1])) @@ -204,7 +204,7 @@ def test_rks_gen_g_hop(self): hcore = mf.get_hcore() g, hop, hdiag = mf.gen_g_hop(mo, mo_occ, hcore) dat = hop(dm1) - self.assertAlmostEqual(lib.fp(dat.get()), -2.580202335681725-0.28028319954452735j, 9) + self.assertAlmostEqual(lib.fp(dat.get()), -2.580202335681725-0.28028319954452735j, 8) mf_ref = cell.KRKS(kpts=cell.make_kpts([2,1,1])) mf_ref.xc = 'b3lyp5' @@ -220,7 +220,7 @@ def test_rks_gen_g_hop(self): dm1 = cp.append(dm1, 0.5) g, hop, hdiag = mf.gen_g_hop(mo, mo_occ, hcore) dat = hop(dm1) - self.assertAlmostEqual(lib.fp(dat.get()), 7.474044381526779+0.0640933714545938j, 9) + self.assertAlmostEqual(lib.fp(dat.get()), 7.474044381526779+0.0640933714545938j, delta=3e-7) g_ref, hop_ref, hdiag_ref = mf_ref.gen_g_hop(mo.get(), mo_occ.get(), hcore.get()) self.assertAlmostEqual(abs(dat.get() - hop_ref(dm1.get())).max(), 0, 7) @@ -243,7 +243,7 @@ def test_uks_gen_g_hop(self): hcore = cp.array([mf.get_hcore()]*2) g, hop, hdiag = mf.gen_g_hop(mo, mo_occ, hcore) dat = hop(dm1) - self.assertAlmostEqual(lib.fp(dat.get()), -10.115070233160678-0.818925838878515j, 9) + self.assertAlmostEqual(lib.fp(dat.get()), -10.115070233160678-0.818925838878515j, delta=1e-7) mf_ref = cell.KUKS(kpts=cell.make_kpts([2,1,1])) mf_ref.xc = 'b3lyp5' @@ -259,7 +259,7 @@ def test_uks_gen_g_hop(self): dm1 = cp.append(dm1, 0.5) g, hop, hdiag = mf.gen_g_hop(mo, mo_occ, hcore) dat = hop(dm1) - self.assertAlmostEqual(lib.fp(dat.get()), 0.21794953957922497-0.08227815483834572j, 9) + self.assertAlmostEqual(lib.fp(dat.get()), 0.21794953957922497-0.08227815483834572j, delta=1e-7) g_ref, hop_ref, hdiag_ref = mf_ref.gen_g_hop(mo.get(), mo_occ.get(), hcore.get()) self.assertAlmostEqual(abs(dat.get() - hop_ref(dm1.get())).max(), 0, delta=3e-8) diff --git a/gpu4pyscf/scf/hf.py b/gpu4pyscf/scf/hf.py index 246f559c5..51be35c76 100644 --- a/gpu4pyscf/scf/hf.py +++ b/gpu4pyscf/scf/hf.py @@ -228,12 +228,15 @@ def _kernel(mf, conv_tol=1e-10, conv_tol_grad=None, assert isinstance(dm0, cupy.ndarray) + dm, dm0 = dm0, None + # Call get_veff before get_hcore. In PBC, the initialization for + # two-electron integrals can be reused by get_hcore, avoiding redundant + # initialization. + vhf = mf.get_veff(mol, dm) + h1e = cupy.asarray(mf.get_hcore()) s1e = cupy.asarray(mf.get_ovlp()) - t1 = log.timer_debug1('hcore', *t1) - dm, dm0 = dm0, None - vhf = mf.get_veff(mol, dm) e_tot = mf.energy_tot(dm, h1e, vhf) log.info('init E= %.15g', e_tot) x_orth = mf.check_linear_dependency(s1e, log) @@ -301,7 +304,13 @@ def _kernel(mf, conv_tol=1e-10, conv_tol_grad=None, cycle+1, e_tot, e_tot-last_hf_e, norm_gorb, norm_ddm) if dump_chk: - mf.dump_chk(locals()) + mf.dump_chk({ + 'mol': mol, + 'mo_energy': mo_energy, + 'mo_occ': mo_occ, + 'mo_coeff': mo_coeff, + 'e_tot': e_tot, + }) if callable(callback): callback(locals()) @@ -315,8 +324,10 @@ def _kernel(mf, conv_tol=1e-10, conv_tol_grad=None, mf.cycles = cycle + 1 if scf_conv and mf.level_shift is not None: + mo_coeff = mo_occ = mo_energy = mf_diis = None # An extra diagonalization, to remove level shift mo_energy, mo_coeff = mf.eig(fock, s1e, x=x_orth) + fock = None mo_occ = mf.get_occ(mo_energy, mo_coeff) dm, dm_last = mf.make_rdm1(mo_coeff, mo_occ), dm vhf = mf.get_veff(mol, dm, dm_last, vhf) @@ -922,8 +933,8 @@ def reset(self, mol=None): return self def dump_chk(self, envs): - assert isinstance(envs, dict) if self.chkfile: + assert isinstance(envs, dict) chkfile.dump_scf( self.mol, self.chkfile, envs['e_tot'], cupy.asnumpy(envs['mo_energy']), cupy.asnumpy(envs['mo_coeff']),