From 28e2f5c3d196e80ae789f762049c17e6ba756e27 Mon Sep 17 00:00:00 2001 From: Qiming Sun Date: Tue, 15 Sep 2026 10:26:44 -0700 Subject: [PATCH 1/7] ppnl derivatives kernels --- gpu4pyscf/lib/pbc/CMakeLists.txt | 3 +- gpu4pyscf/lib/pbc/overlap.cu | 942 +--------------- gpu4pyscf/lib/pbc/ppnl.cu | 1033 ++++++++++++++++++ gpu4pyscf/pbc/grad/krhf.py | 3 +- gpu4pyscf/pbc/grad/kuhf.py | 3 +- gpu4pyscf/pbc/grad/pp.py | 240 ++-- gpu4pyscf/pbc/grad/rhf.py | 3 +- gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py | 100 ++ gpu4pyscf/pbc/gto/int1e.py | 11 +- gpu4pyscf/pbc/gto/pseudo/pp_int.py | 36 +- 10 files changed, 1313 insertions(+), 1061 deletions(-) create mode 100644 gpu4pyscf/lib/pbc/ppnl.cu diff --git a/gpu4pyscf/lib/pbc/CMakeLists.txt b/gpu4pyscf/lib/pbc/CMakeLists.txt index a5072ac71..e855ef3f2 100644 --- a/gpu4pyscf/lib/pbc/CMakeLists.txt +++ b/gpu4pyscf/lib/pbc/CMakeLists.txt @@ -6,7 +6,8 @@ add_library(pbc SHARED fill_int3c2e.cu contract_int3c2e.cu ejk_int3c2e_ip1.cu fill_int2c2e.cu fill_int2c2e_ip1.cu e_int2c2e_ip1.cu contract_int3c2e_pvp.cu - overlap.cu supmol_sr_estimator.cu + overlap.cu ppnl.cu + supmol_sr_estimator.cu rys_roots_dat.cu rys_contract_k.cu rys_contract_jk_ip1.cu rys_contract_j.cu diff --git a/gpu4pyscf/lib/pbc/overlap.cu b/gpu4pyscf/lib/pbc/overlap.cu index 7685c249d..bf997d2b4 100644 --- a/gpu4pyscf/lib/pbc/overlap.cu +++ b/gpu4pyscf/lib/pbc/overlap.cu @@ -23,524 +23,15 @@ #include "gvhf-rys/rys_contract_k.cuh" #include "pbc.cuh" #include "int3c2e.cuh" +#include "overlap.cuh" -#define PI_POW_1_5 5.568327996831707845 #define GOUT_WIDTH 36 #define GOUT_WIDTH_IP1 18 -#define REMOTE_THRESHOLD 50 -__inline__ __device__ -void vrr_hrr(double *gx, double *rjri, double ai, double aj, double cicj, - int li, int lj, int gout_id, int gout_stride, int nsp_per_block) -{ - int stride_j = li + 1; - int g_size = (li + 1) * (lj + 1); - int gx_len = g_size * nsp_per_block; - double aij = ai + aj; - double aj_aij = aj / aij; - if (gout_id == 0) { - double theta = ai * aj_aij; - double theta_rr = theta * rjri[3*nsp_per_block]; - gx[gx_len*2] = cicj / (aij*sqrt(aij)) * exp(-theta_rr); - } - int lij = li + lj; - if (lij > 0) { - __syncthreads(); - double s0x, s1x, s2x; - double b = .5 / aij; - for (int n = gout_id; n < 3; n += gout_stride) { - double *_gx = gx + n * gx_len; - double xjxi = rjri[n*nsp_per_block]; - double xpa = xjxi * aj_aij; - s0x = _gx[0]; - s1x = xpa * s0x; - _gx[nsp_per_block] = s1x; - for (int i = 1; i < lij; ++i) { - s2x = xpa * s1x + i * b * s0x; - _gx[(i+1)*nsp_per_block] = s2x; - s0x = s1x; - s1x = s2x; - } - for (int j = 0; j < lj; ++j) { - int ij = (lij-j) + j*stride_j; - s1x = _gx[ij*nsp_per_block]; - for (--ij; ij >= j*stride_j; --ij) { - s0x = _gx[ij*nsp_per_block]; - _gx[(ij+stride_j)*nsp_per_block] = s1x - xjxi * s0x; - s1x = s0x; - } - } - } - } - __syncthreads(); -} - -__global__ static -void int1e_ovlp_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, - int *shl_pair_offsets, int *gout_stride_lookup, - int naoi, int naoj, size_t ij_offset) -{ - int sp_block_id = blockIdx.x; - int thread_id = threadIdx.x; - int nbas = envs.cell0_nbas * envs.bvk_ncells; - int *bas = envs.bas; - double *env = envs.env; - double *img_coords = envs.img_coords; - __shared__ int shl_pair0, shl_pair1; - __shared__ int li, lj, iprim, jprim; - __shared__ int gout_stride, nsp_per_block; - if (thread_id == 0) { - shl_pair0 = shl_pair_offsets[sp_block_id]; - shl_pair1 = shl_pair_offsets[sp_block_id+1]; - int bas_ij0 = bas_ij_idx[shl_pair0]; - int ish0 = bas_ij0 / nbas; - int jsh0 = bas_ij0 % nbas; - li = bas[ish0*BAS_SLOTS+ANG_OF]; - lj = bas[jsh0*BAS_SLOTS+ANG_OF]; - iprim = bas[ish0*BAS_SLOTS+NPRIM_OF]; - jprim = bas[jsh0*BAS_SLOTS+NPRIM_OF]; - gout_stride = gout_stride_lookup[li*L_AUX1+lj]; - nsp_per_block = THREADS / gout_stride; - } - __syncthreads(); - int sp_id = thread_id % nsp_per_block; - int gout_id = thread_id / nsp_per_block; - int g_size = (li + 1) * (lj + 1); - int gx_len = g_size * nsp_per_block; - extern __shared__ double g[]; - double *gx = g + sp_id; - double *gy = g + gx_len + sp_id; - double *gz = g + gx_len * 2 + sp_id; - double *rjri = g + gx_len * 3 + sp_id; - int idx_i = lex_xyz_offset(li); - int idx_j = lex_xyz_offset(lj); - if (gout_id == 0) { - gx[0] = PI_POW_1_5; - gy[0] = 1.; - } - - for (int pair_ij = shl_pair0+sp_id; pair_ij < shl_pair1+sp_id; pair_ij += nsp_per_block) { - double gout[GOUT_WIDTH]; -#pragma unroll - for (int n = 0; n < GOUT_WIDTH; ++n) { - gout[n] = 0.; - } - int bas_ij; - if (pair_ij >= shl_pair1) { - bas_ij = bas_ij_idx[shl_pair0]; - } else { - bas_ij = bas_ij_idx[pair_ij]; - } - int ish = bas_ij / nbas; - int jsh = bas_ij % nbas; - int ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; - int rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; - int expi = bas[ish*BAS_SLOTS+PTR_EXP]; - int expj = bas[jsh*BAS_SLOTS+PTR_EXP]; - int ci = bas[ish*BAS_SLOTS+PTR_COEFF]; - int cj = bas[jsh*BAS_SLOTS+PTR_COEFF]; - for (int img = 0; img < envs.nimgs; img++) { - if (gout_id == 0) { - double xjL = img_coords[img*3+0]; - double yjL = img_coords[img*3+1]; - double zjL = img_coords[img*3+2]; - double xjxi = env[rj+0] + xjL - env[ri+0]; - double yjyi = env[rj+1] + yjL - env[ri+1]; - double zjzi = env[rj+2] + zjL - env[ri+2]; - double rr_ij = xjxi*xjxi + yjyi*yjyi + zjzi*zjzi; - rjri[0*nsp_per_block] = xjxi; - rjri[1*nsp_per_block] = yjyi; - rjri[2*nsp_per_block] = zjzi; - rjri[3*nsp_per_block] = rr_ij; - } - int ijprim = iprim * jprim; - for (int ijp = 0; ijp < ijprim; ++ijp) { - __syncthreads(); - int ip = ijp % iprim; - int jp = ijp / iprim; - double ai = env[expi+ip]; - double aj = env[expj+jp]; - double cicj = env[ci+ip] * env[cj+jp]; - vrr_hrr(gx, rjri, ai, aj, cicj, li, lj, gout_id, gout_stride, - nsp_per_block); - if (pair_ij >= shl_pair1) { - continue; - } - int nsp = nsp_per_block; - int stride_j = li + 1; - int nfi = c_nf[li]; - int nfj = c_nf[lj]; - int nfij = nfi * nfj; - float div_nfi = c_div_nf[li]; -#pragma unroll - for (int n = 0; n < GOUT_WIDTH; ++n) { - uint32_t ij = gout_id + n * gout_stride; - if (ij >= nfij) break; - uint32_t j = ij * div_nfi; - uint32_t i = ij - j * nfi; - 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 + jx*stride_j) * nsp; - int addry = (iy + jy*stride_j) * nsp; - int addrz = (iz + jz*stride_j) * nsp; - gout[n] += gx[addrx] * gy[addry] * gz[addrz]; - } - } - } - - if (pair_ij < shl_pair1) { - int nfi = c_nf[li]; - int nfj = c_nf[lj]; - int nfij = nfi * nfj; - int *ao_loc = envs.ao_loc; - int nbas = envs.cell0_nbas; - int cell_id = jsh / nbas; - int jshp = jsh % nbas; - int i0 = ao_loc[ish]; - int j0 = ao_loc[jshp]; - double *out_subblock = out + (cell_id*naoi+i0) * naoj + j0 - ij_offset; -#pragma unroll - for (int n = 0; n < GOUT_WIDTH; ++n) { - int ij = n*gout_stride+gout_id; - if (ij >= nfij) break; - int j = ij / nfi; - int i = ij % nfi; - out_subblock[i*naoj+j] = gout[n]; - } - } - } -} - -static __global__ -void int1e_kin_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, - int *shl_pair_offsets, int *gout_stride_lookup, - int naoi, int naoj, size_t ij_offset) -{ - int sp_block_id = blockIdx.x; - int thread_id = threadIdx.x; - int nbas = envs.cell0_nbas * envs.bvk_ncells; - int *bas = envs.bas; - double *env = envs.env; - double *img_coords = envs.img_coords; - __shared__ int shl_pair0, shl_pair1; - __shared__ int li, lj, iprim, jprim; - __shared__ int gout_stride, nsp_per_block; - if (thread_id == 0) { - shl_pair0 = shl_pair_offsets[sp_block_id]; - shl_pair1 = shl_pair_offsets[sp_block_id+1]; - int bas_ij0 = bas_ij_idx[shl_pair0]; - int ish0 = bas_ij0 / nbas; - int jsh0 = bas_ij0 % nbas; - li = bas[ish0*BAS_SLOTS+ANG_OF]; - lj = bas[jsh0*BAS_SLOTS+ANG_OF]; - iprim = bas[ish0*BAS_SLOTS+NPRIM_OF]; - jprim = bas[jsh0*BAS_SLOTS+NPRIM_OF]; - gout_stride = gout_stride_lookup[li*L_AUX1+lj]; - nsp_per_block = THREADS / gout_stride; - } - __syncthreads(); - int sp_id = thread_id % nsp_per_block; - int gout_id = thread_id / nsp_per_block; - int g_size = (li + 3) * (lj + 1); - int gx_len = g_size * nsp_per_block; - extern __shared__ double g[]; - double *gx = g + sp_id; - double *gy = g + gx_len + sp_id; - double *gz = g + gx_len * 2 + sp_id; - double *rjri = g + gx_len * 3 + sp_id; - int idx_i = lex_xyz_offset(li); - int idx_j = lex_xyz_offset(lj); - if (gout_id == 0) { - gx[0] = PI_POW_1_5; - gy[0] = -.5; - } - - for (int pair_ij = shl_pair0+sp_id; pair_ij < shl_pair1+sp_id; pair_ij += nsp_per_block) { - double gout[GOUT_WIDTH]; -#pragma unroll - for (int n = 0; n < GOUT_WIDTH; ++n) { - gout[n] = 0.; - } - int bas_ij; - if (pair_ij >= shl_pair1) { - bas_ij = bas_ij_idx[shl_pair0]; - } else { - bas_ij = bas_ij_idx[pair_ij]; - } - int ish = bas_ij / nbas; - int jsh = bas_ij % nbas; - int ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; - int rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; - int expi = bas[ish*BAS_SLOTS+PTR_EXP]; - int expj = bas[jsh*BAS_SLOTS+PTR_EXP]; - int ci = bas[ish*BAS_SLOTS+PTR_COEFF]; - int cj = bas[jsh*BAS_SLOTS+PTR_COEFF]; - for (int img = 0; img < envs.nimgs; img++) { - if (gout_id == 0) { - double xjL = img_coords[img*3+0]; - double yjL = img_coords[img*3+1]; - double zjL = img_coords[img*3+2]; - double xjxi = env[rj+0] + xjL - env[ri+0]; - double yjyi = env[rj+1] + yjL - env[ri+1]; - double zjzi = env[rj+2] + zjL - env[ri+2]; - double rr_ij = xjxi*xjxi + yjyi*yjyi + zjzi*zjzi; - rjri[0*nsp_per_block] = xjxi; - rjri[1*nsp_per_block] = yjyi; - rjri[2*nsp_per_block] = zjzi; - rjri[3*nsp_per_block] = rr_ij; - } - int ijprim = iprim * jprim; - for (int ijp = 0; ijp < ijprim; ++ijp) { - __syncthreads(); - int ip = ijp % iprim; - int jp = ijp / iprim; - double ai = env[expi+ip]; - double aj = env[expj+jp]; - double cicj = env[ci+ip] * env[cj+jp]; - vrr_hrr(gx, rjri, ai, aj, cicj, li+2, lj, gout_id, gout_stride, - nsp_per_block); - if (pair_ij >= shl_pair1) { - continue; - } - int nsp = nsp_per_block; - int stride_j = li + 3; - int i_1 = nsp_per_block; - double ai2 = ai * -2; - float div_nfi = c_div_nf[li]; - int nfi = c_nf[li]; - int nfj = c_nf[lj]; - int nfij = nfi * nfj; -#pragma unroll - for (int n = 0; n < GOUT_WIDTH; ++n) { - uint32_t ij = gout_id + n * gout_stride; - if (ij >= nfij) break; - uint32_t j = ij * div_nfi; - uint32_t i = ij - j * nfi; - 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 + jx*stride_j) * nsp; - int addry = (iy + jy*stride_j) * nsp; - int addrz = (iz + jz*stride_j) * nsp; - double fx0 = gx[addrx]; - double fy0 = gy[addry]; - double fz0 = gz[addrz]; - double fx2 = ai2 * ((ix*2+1)*fx0 + ai2*gx[addrx+i_1*2]); - double fy2 = ai2 * ((iy*2+1)*fy0 + ai2*gy[addry+i_1*2]); - double fz2 = ai2 * ((iz*2+1)*fz0 + ai2*gz[addrz+i_1*2]); - if (ix > 1) fx2 += ix*(ix-1) * gx[addrx-i_1*2]; - if (iy > 1) fy2 += iy*(iy-1) * gy[addry-i_1*2]; - if (iz > 1) fz2 += iz*(iz-1) * gz[addrz-i_1*2]; - gout[n] += fx2 * fy0 * fz0; - gout[n] += fx0 * fy2 * fz0; - gout[n] += fx0 * fy0 * fz2; - } - } - } - - if (pair_ij < shl_pair1) { - int nfi = c_nf[li]; - int nfj = c_nf[lj]; - int nfij = nfi * nfj; - int *ao_loc = envs.ao_loc; - int nbas = envs.cell0_nbas; - int cell_id = jsh / nbas; - int jshp = jsh % nbas; - int i0 = ao_loc[ish]; - int j0 = ao_loc[jshp]; - double *out_subblock = out + (cell_id*naoi+i0) * naoj + j0 - ij_offset; -#pragma unroll - for (int n = 0; n < GOUT_WIDTH; ++n) { - int ij = n*gout_stride+gout_id; - if (ij >= nfij) break; - int j = ij / nfi; - int i = ij % nfi; - out_subblock[i*naoj+j] = gout[n]; - } - } - } -} - -// r^2 moment about origin: with r measured from (0,0,0). -// -// Trick: x = (x - Bx) + Bx, so x^2 = (x-Bx)^2 + 2*Bx*(x-Bx) + Bx^2. This -// lets us reuse the existing OS recursion for S(ix, jx) = -// and assemble the moment via three j-shifted samples. The recursion is -// extended by lj+2 in the j-axis (lij bumped by 2, HRR loop bumped by 2). -// Final integrand for r^2 is x^2 + y^2 + z^2. -__global__ static -void int1e_r2_origi_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, - int *shl_pair_offsets, int *gout_stride_lookup, - int naoi, int naoj, size_t ij_offset) -{ - int sp_block_id = blockIdx.x; - int thread_id = threadIdx.x; - int nbas = envs.cell0_nbas * envs.bvk_ncells; - int *bas = envs.bas; - double *env = envs.env; - double *img_coords = envs.img_coords; - __shared__ int shl_pair0, shl_pair1; - __shared__ int li, lj, iprim, jprim; - __shared__ int gout_stride, nsp_per_block; - if (thread_id == 0) { - shl_pair0 = shl_pair_offsets[sp_block_id]; - shl_pair1 = shl_pair_offsets[sp_block_id+1]; - int bas_ij0 = bas_ij_idx[shl_pair0]; - int ish0 = bas_ij0 / nbas; - int jsh0 = bas_ij0 % nbas; - li = bas[ish0*BAS_SLOTS+ANG_OF]; - lj = bas[jsh0*BAS_SLOTS+ANG_OF]; - iprim = bas[ish0*BAS_SLOTS+NPRIM_OF]; - jprim = bas[jsh0*BAS_SLOTS+NPRIM_OF]; - gout_stride = gout_stride_lookup[li*L_AUX1+lj]; - nsp_per_block = THREADS / gout_stride; - } - __syncthreads(); - int sp_id = thread_id % nsp_per_block; - int gout_id = thread_id / nsp_per_block; - int g_size = (li + 1) * (lj + 3); - int gx_len = g_size * nsp_per_block; - extern __shared__ double g[]; - double *gx = g + sp_id; - double *gy = g + gx_len + sp_id; - double *gz = g + gx_len * 2 + sp_id; - double *rjri = g + gx_len * 3 + sp_id; - int idx_i = lex_xyz_offset(li); - int idx_j = lex_xyz_offset(lj); - if (gout_id == 0) { - gx[0] = PI_POW_1_5; - gy[0] = 1.; - } - - for (int pair_ij = shl_pair0+sp_id; pair_ij < shl_pair1+sp_id; pair_ij += nsp_per_block) { - double gout[GOUT_WIDTH]; -#pragma unroll - for (int n = 0; n < GOUT_WIDTH; ++n) { - gout[n] = 0.; - } - int bas_ij; - if (pair_ij >= shl_pair1) { - bas_ij = bas_ij_idx[shl_pair0]; - } else { - bas_ij = bas_ij_idx[pair_ij]; - } - int ish = bas_ij / nbas; - int jsh = bas_ij % nbas; - int ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; - int rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; - int expi = bas[ish*BAS_SLOTS+PTR_EXP]; - int expj = bas[jsh*BAS_SLOTS+PTR_EXP]; - int ci = bas[ish*BAS_SLOTS+PTR_COEFF]; - int cj = bas[jsh*BAS_SLOTS+PTR_COEFF]; - for (int img = 0; img < envs.nimgs; img++) { - if (gout_id == 0) { - double xjL = img_coords[img*3+0]; - double yjL = img_coords[img*3+1]; - double zjL = img_coords[img*3+2]; - double xjxi = env[rj+0] + xjL - env[ri+0]; - double yjyi = env[rj+1] + yjL - env[ri+1]; - double zjzi = env[rj+2] + zjL - env[ri+2]; - double rr_ij = xjxi*xjxi + yjyi*yjyi + zjzi*zjzi; - rjri[0*nsp_per_block] = xjxi; - rjri[1*nsp_per_block] = yjyi; - rjri[2*nsp_per_block] = zjzi; - rjri[3*nsp_per_block] = rr_ij; - } - // Separation vector (ket - bra) for moment expansion about bra center. - // libcint's "origi" convention: r^2 measured from bra center. - double Bx = env[rj+0] + img_coords[img*3+0] - env[ri+0]; - double By = env[rj+1] + img_coords[img*3+1] - env[ri+1]; - double Bz = env[rj+2] + img_coords[img*3+2] - env[ri+2]; - int ijprim = iprim * jprim; - for (int ijp = 0; ijp < ijprim; ++ijp) { - __syncthreads(); - int ip = ijp % iprim; - int jp = ijp / iprim; - double ai = env[expi+ip]; - double aj = env[expj+jp]; - double cicj = env[ci+ip] * env[cj+jp]; - vrr_hrr(gx, rjri, ai, aj, cicj, li, lj+2, gout_id, gout_stride, - nsp_per_block); - if (pair_ij >= shl_pair1) { - continue; - } - int nsp = nsp_per_block; - int stride_j = li + 1; - int nfi = c_nf[li]; - int nfj = c_nf[lj]; - int nfij = nfi * nfj; - float div_nfi = c_div_nf[li]; -#pragma unroll - for (int n = 0; n < GOUT_WIDTH; ++n) { - uint32_t ij = gout_id + n * gout_stride; - if (ij >= nfij) break; - uint32_t j = ij * div_nfi; - uint32_t i = ij - j * nfi; - 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 addrx0 = (ix + (jx+0)*stride_j) * nsp; - int addrx1 = (ix + (jx+1)*stride_j) * nsp; - int addrx2 = (ix + (jx+2)*stride_j) * nsp; - int addry0 = (iy + (jy+0)*stride_j) * nsp; - int addry1 = (iy + (jy+1)*stride_j) * nsp; - int addry2 = (iy + (jy+2)*stride_j) * nsp; - int addrz0 = (iz + (jz+0)*stride_j) * nsp; - int addrz1 = (iz + (jz+1)*stride_j) * nsp; - int addrz2 = (iz + (jz+2)*stride_j) * nsp; - double sx = gx[addrx0]; - double sy = gy[addry0]; - double sz = gz[addrz0]; - {double mx = gx[addrx2] + 2.*Bx*gx[addrx1] + Bx*Bx*sx; - double my = gy[addry2] + 2.*By*gy[addry1] + By*By*sy; - double mz = gz[addrz2] + 2.*Bz*gz[addrz1] + Bz*Bz*sz; - gout[n] += mx*sy*sz + sx*my*sz + sx*sy*mz;} - } - } - } - - if (pair_ij < shl_pair1) { - int nfi = c_nf[li]; - int nfj = c_nf[lj]; - int nfij = nfi * nfj; - int *ao_loc = envs.ao_loc; - int nbas = envs.cell0_nbas; - int cell_id = jsh / nbas; - int jshp = jsh % nbas; - int i0 = ao_loc[ish]; - int j0 = ao_loc[jshp]; - double *out_subblock = out + (cell_id*naoi+i0) * naoj + j0 - ij_offset; -#pragma unroll - for (int n = 0; n < GOUT_WIDTH; ++n) { - int ij = n*gout_stride+gout_id; - if (ij >= nfij) break; - int j = ij / nfi; - int i = ij % nfi; - out_subblock[i*naoj+j] = gout[n]; - } - } - } -} - -// r^4 moment about origin: with r measured from (0,0,0). -// -// r^4 = x^4+y^4+z^4 + 2(x^2*y^2 + y^2*z^2 + x^2*z^2). -// Each 1D moment x^n uses binomial expansion x^n = sum_k C(n,k) Bx^(n-k) (x-Bx)^k. -// Recursion extended by lj+4 in j-axis. __global__ static -void int1e_r4_origi_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, - int *shl_pair_offsets, int *gout_stride_lookup, - int naoi, int naoj, size_t ij_offset) +void int1e_ovlp_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, + int *shl_pair_offsets, int *gout_stride_lookup, + int naoi, int naoj, size_t ij_offset) { int sp_block_id = blockIdx.x; int thread_id = threadIdx.x; @@ -567,7 +58,7 @@ void int1e_r4_origi_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, __syncthreads(); int sp_id = thread_id % nsp_per_block; int gout_id = thread_id / nsp_per_block; - int g_size = (li + 1) * (lj + 5); + int g_size = (li + 1) * (lj + 1); int gx_len = g_size * nsp_per_block; extern __shared__ double g[]; double *gx = g + sp_id; @@ -615,9 +106,6 @@ void int1e_r4_origi_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, rjri[2*nsp_per_block] = zjzi; rjri[3*nsp_per_block] = rr_ij; } - double Bx = env[rj+0] + img_coords[img*3+0] - env[ri+0]; - double By = env[rj+1] + img_coords[img*3+1] - env[ri+1]; - double Bz = env[rj+2] + img_coords[img*3+2] - env[ri+2]; int ijprim = iprim * jprim; for (int ijp = 0; ijp < ijprim; ++ijp) { __syncthreads(); @@ -626,7 +114,7 @@ void int1e_r4_origi_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, double ai = env[expi+ip]; double aj = env[expj+jp]; double cicj = env[ci+ip] * env[cj+jp]; - vrr_hrr(gx, rjri, ai, aj, cicj, li, lj+4, gout_id, gout_stride, + vrr_hrr(gx, rjri, ai, aj, cicj, li, lj, gout_id, gout_stride, nsp_per_block); if (pair_ij >= shl_pair1) { continue; @@ -637,7 +125,6 @@ void int1e_r4_origi_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, int nfj = c_nf[lj]; int nfij = nfi * nfj; float div_nfi = c_div_nf[li]; - double Bx2 = Bx*Bx, By2 = By*By, Bz2 = Bz*Bz; #pragma unroll for (int n = 0; n < GOUT_WIDTH; ++n) { uint32_t ij = gout_id + n * gout_stride; @@ -650,33 +137,10 @@ void int1e_r4_origi_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, 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]; - // Load S(ix, jx+k) for k=0..4 (and same for y, z) - double Sx0 = gx[(ix + (jx+0)*stride_j) * nsp]; - double Sx1 = gx[(ix + (jx+1)*stride_j) * nsp]; - double Sx2 = gx[(ix + (jx+2)*stride_j) * nsp]; - double Sx3 = gx[(ix + (jx+3)*stride_j) * nsp]; - double Sx4 = gx[(ix + (jx+4)*stride_j) * nsp]; - double Sy0 = gy[(iy + (jy+0)*stride_j) * nsp]; - double Sy1 = gy[(iy + (jy+1)*stride_j) * nsp]; - double Sy2 = gy[(iy + (jy+2)*stride_j) * nsp]; - double Sy3 = gy[(iy + (jy+3)*stride_j) * nsp]; - double Sy4 = gy[(iy + (jy+4)*stride_j) * nsp]; - double Sz0 = gz[(iz + (jz+0)*stride_j) * nsp]; - double Sz1 = gz[(iz + (jz+1)*stride_j) * nsp]; - double Sz2 = gz[(iz + (jz+2)*stride_j) * nsp]; - double Sz3 = gz[(iz + (jz+3)*stride_j) * nsp]; - double Sz4 = gz[(iz + (jz+4)*stride_j) * nsp]; - // x^2 moment: S(ix,jx+2) + 2*Bx*S(ix,jx+1) + Bx^2*S(ix,jx) - double mx2 = Sx2 + 2.*Bx*Sx1 + Bx2*Sx0; - double my2 = Sy2 + 2.*By*Sy1 + By2*Sy0; - double mz2 = Sz2 + 2.*Bz*Sz1 + Bz2*Sz0; - // x^4 moment: S4 + 4*Bx*S3 + 6*Bx^2*S2 + 4*Bx^3*S1 + Bx^4*S0 - double mx4 = Sx4 + 4.*Bx*Sx3 + 6.*Bx2*Sx2 + 4.*Bx2*Bx*Sx1 + Bx2*Bx2*Sx0; - double my4 = Sy4 + 4.*By*Sy3 + 6.*By2*Sy2 + 4.*By2*By*Sy1 + By2*By2*Sy0; - double mz4 = Sz4 + 4.*Bz*Sz3 + 6.*Bz2*Sz2 + 4.*Bz2*Bz*Sz1 + Bz2*Bz2*Sz0; - // r^4 = x^4+y^4+z^4 + 2(x^2*y^2 + y^2*z^2 + x^2*z^2) - gout[n] += mx4*Sy0*Sz0 + Sx0*my4*Sz0 + Sx0*Sy0*mz4 - + 2.*(mx2*my2*Sz0 + Sx0*my2*mz2 + mx2*Sy0*mz2); + int addrx = (ix + jx*stride_j) * nsp; + int addry = (iy + jy*stride_j) * nsp; + int addrz = (iz + jz*stride_j) * nsp; + gout[n] += gx[addrx] * gy[addry] * gz[addrz]; } } } @@ -704,15 +168,10 @@ void int1e_r4_origi_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, } } -// ip2 derivative of r^2 moment about origin: d/dB . -// d/dBx acts only on the j-basis (the weight r^2 is fixed at origin). -// Uses j-side derivative: D_x[f(jx)] = -2*aj*f(jx+1) + jx*f(jx-1). -// Needs S(ix, jx-1..jx+3) => lij = li+lj+3, g_size = (li+1)*(lj+4). -// Output: 3 components (goutx, gouty, goutz) matching ip convention. -__global__ static -void int1e_r2_origi_ip2_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, - int *shl_pair_offsets, int *gout_stride_lookup, - int naoi, int naoj, size_t ij_offset) +static __global__ +void int1e_kin_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, + int *shl_pair_offsets, int *gout_stride_lookup, + int naoi, int naoj, size_t ij_offset) { int sp_block_id = blockIdx.x; int thread_id = threadIdx.x; @@ -739,7 +198,7 @@ void int1e_r2_origi_ip2_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, __syncthreads(); int sp_id = thread_id % nsp_per_block; int gout_id = thread_id / nsp_per_block; - int g_size = (li + 1) * (lj + 4); + int g_size = (li + 3) * (lj + 1); int gx_len = g_size * nsp_per_block; extern __shared__ double g[]; double *gx = g + sp_id; @@ -750,18 +209,14 @@ void int1e_r2_origi_ip2_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, int idx_j = lex_xyz_offset(lj); if (gout_id == 0) { gx[0] = PI_POW_1_5; - gy[0] = 1.; + gy[0] = -.5; } for (int pair_ij = shl_pair0+sp_id; pair_ij < shl_pair1+sp_id; pair_ij += nsp_per_block) { - double goutx[GOUT_WIDTH_IP1]; - double gouty[GOUT_WIDTH_IP1]; - double goutz[GOUT_WIDTH_IP1]; + double gout[GOUT_WIDTH]; #pragma unroll - for (int n = 0; n < GOUT_WIDTH_IP1; ++n) { - goutx[n] = 0.; - gouty[n] = 0.; - goutz[n] = 0.; + for (int n = 0; n < GOUT_WIDTH; ++n) { + gout[n] = 0.; } int bas_ij; if (pair_ij >= shl_pair1) { @@ -791,10 +246,6 @@ void int1e_r2_origi_ip2_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, rjri[2*nsp_per_block] = zjzi; rjri[3*nsp_per_block] = rr_ij; } - double Bx = env[rj+0] + img_coords[img*3+0] - env[ri+0]; - double By = env[rj+1] + img_coords[img*3+1] - env[ri+1]; - double Bz = env[rj+2] + img_coords[img*3+2] - env[ri+2]; - double Bx2 = Bx*Bx, By2 = By*By, Bz2 = Bz*Bz; int ijprim = iprim * jprim; for (int ijp = 0; ijp < ijprim; ++ijp) { __syncthreads(); @@ -803,20 +254,21 @@ void int1e_r2_origi_ip2_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, double ai = env[expi+ip]; double aj = env[expj+jp]; double cicj = env[ci+ip] * env[cj+jp]; - vrr_hrr(gx, rjri, ai, aj, cicj, li, lj+3, gout_id, gout_stride, + vrr_hrr(gx, rjri, ai, aj, cicj, li+2, lj, gout_id, gout_stride, nsp_per_block); if (pair_ij >= shl_pair1) { continue; } int nsp = nsp_per_block; - int stride_j = li + 1; - double aj2 = aj * -2; + int stride_j = li + 3; + int i_1 = nsp_per_block; + double ai2 = ai * -2; float div_nfi = c_div_nf[li]; int nfi = c_nf[li]; int nfj = c_nf[lj]; int nfij = nfi * nfj; #pragma unroll - for (int n = 0; n < GOUT_WIDTH_IP1; ++n) { + for (int n = 0; n < GOUT_WIDTH; ++n) { uint32_t ij = gout_id + n * gout_stride; if (ij >= nfij) break; uint32_t j = ij * div_nfi; @@ -827,283 +279,43 @@ void int1e_r2_origi_ip2_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, 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]; - // Load S(ix, jx+k) for k = -1..3 (and y, z analogues) - // k=-1 only contributes when jx > 0 (multiplied by jx) - #define SADDR(g,a,joff) ((g)[((a) + (joff)*stride_j) * nsp]) - double Sxm1 = (jx > 0) ? SADDR(gx,ix,jx-1) : 0.; - double Sx0 = SADDR(gx,ix,jx); - double Sx1 = SADDR(gx,ix,jx+1); - double Sx2 = SADDR(gx,ix,jx+2); - double Sx3 = SADDR(gx,ix,jx+3); - double Sym1 = (jy > 0) ? SADDR(gy,iy,jy-1) : 0.; - double Sy0 = SADDR(gy,iy,jy); - double Sy1 = SADDR(gy,iy,jy+1); - double Sy2 = SADDR(gy,iy,jy+2); - double Sy3 = SADDR(gy,iy,jy+3); - double Szm1 = (jz > 0) ? SADDR(gz,iz,jz-1) : 0.; - double Sz0 = SADDR(gz,iz,jz); - double Sz1 = SADDR(gz,iz,jz+1); - double Sz2 = SADDR(gz,iz,jz+2); - double Sz3 = SADDR(gz,iz,jz+3); - #undef SADDR - // Undifferentiated overlaps and x^2 moments: - double sx = Sx0, sy = Sy0, sz = Sz0; - double mx = Sx2 + 2.*Bx*Sx1 + Bx2*Sx0; - double my = Sy2 + 2.*By*Sy1 + By2*Sy0; - double mz = Sz2 + 2.*Bz*Sz1 + Bz2*Sz0; - // j-side derivative D[f(j)] = -2*aj*f(j+1) + j_alpha*f(j-1) - // Applied to overlap: - double Dsx = aj2*Sx1 + jx*Sxm1; - double Dsy = aj2*Sy1 + jy*Sym1; - double Dsz = aj2*Sz1 + jz*Szm1; - // Applied to moment (D acts on j-index inside binomial): - // D[mx] = aj2*(S(jx+3)+2Bx*S(jx+2)+Bx^2*S(jx+1)) - // + jx*(S(jx+1)+2Bx*S(jx)+Bx^2*S(jx-1)) - double Dmx = aj2*(Sx3 + 2.*Bx*Sx2 + Bx2*Sx1) - + jx*(Sx1 + 2.*Bx*Sx0 + Bx2*Sxm1); - double Dmy = aj2*(Sy3 + 2.*By*Sy2 + By2*Sy1) - + jy*(Sy1 + 2.*By*Sy0 + By2*Sym1); - double Dmz = aj2*(Sz3 + 2.*Bz*Sz2 + Bz2*Sz1) - + jz*(Sz1 + 2.*Bz*Sz0 + Bz2*Szm1); - // d/dBx: only x-axis j-basis depends on Bx - goutx[n] += Dmx*sy*sz + Dsx*(my*sz + sy*mz); - // d/dBy: only y-axis j-basis depends on By - gouty[n] += Dmy*sx*sz + Dsy*(mx*sz + sx*mz); - // d/dBz: only z-axis j-basis depends on Bz - goutz[n] += Dmz*sx*sy + Dsz*(mx*sy + sx*my); + int addrx = (ix + jx*stride_j) * nsp; + int addry = (iy + jy*stride_j) * nsp; + int addrz = (iz + jz*stride_j) * nsp; + double fx0 = gx[addrx]; + double fy0 = gy[addry]; + double fz0 = gz[addrz]; + double fx2 = ai2 * ((ix*2+1)*fx0 + ai2*gx[addrx+i_1*2]); + double fy2 = ai2 * ((iy*2+1)*fy0 + ai2*gy[addry+i_1*2]); + double fz2 = ai2 * ((iz*2+1)*fz0 + ai2*gz[addrz+i_1*2]); + if (ix > 1) fx2 += ix*(ix-1) * gx[addrx-i_1*2]; + if (iy > 1) fy2 += iy*(iy-1) * gy[addry-i_1*2]; + if (iz > 1) fz2 += iz*(iz-1) * gz[addrz-i_1*2]; + gout[n] += fx2 * fy0 * fz0; + gout[n] += fx0 * fy2 * fz0; + gout[n] += fx0 * fy0 * fz2; } } } if (pair_ij < shl_pair1) { - int *ao_loc = envs.ao_loc; - int nbas = envs.cell0_nbas; - size_t nao2 = naoi * naoj; - int cell_id = jsh / nbas; - int jshp = jsh % nbas; - int i0 = ao_loc[ish]; - int j0 = ao_loc[jshp]; - double *outx = out + cell_id*nao2*3 + i0 * naoj + j0 - ij_offset; - double *outy = outx + nao2; - double *outz = outx + nao2 * 2; int nfi = c_nf[li]; int nfj = c_nf[lj]; int nfij = nfi * nfj; -#pragma unroll - for (int n = 0; n < GOUT_WIDTH_IP1; ++n) { - int ij = n*gout_stride+gout_id; - if (ij >= nfij) break; - int j = ij / nfi; - int i = ij % nfi; - outx[i*naoj+j] = goutx[n]; - outy[i*naoj+j] = gouty[n]; - outz[i*naoj+j] = goutz[n]; - } - } - } -} - -// ip2 derivative of r^4 moment about origin: d/dB . -// Same j-derivative approach. r^4 = x^4+y^4+z^4+2(x^2*y^2+y^2*z^2+x^2*z^2). -// Needs S(ix, jx-1..jx+5) => lij = li+lj+5, g_size = (li+1)*(lj+6). -__global__ static -void int1e_r4_origi_ip2_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, - int *shl_pair_offsets, int *gout_stride_lookup, - int naoi, int naoj, size_t ij_offset) -{ - int sp_block_id = blockIdx.x; - int thread_id = threadIdx.x; - int nbas = envs.cell0_nbas * envs.bvk_ncells; - int *bas = envs.bas; - double *env = envs.env; - double *img_coords = envs.img_coords; - __shared__ int shl_pair0, shl_pair1; - __shared__ int li, lj, iprim, jprim; - __shared__ int gout_stride, nsp_per_block; - if (thread_id == 0) { - shl_pair0 = shl_pair_offsets[sp_block_id]; - shl_pair1 = shl_pair_offsets[sp_block_id+1]; - int bas_ij0 = bas_ij_idx[shl_pair0]; - int ish0 = bas_ij0 / nbas; - int jsh0 = bas_ij0 % nbas; - li = bas[ish0*BAS_SLOTS+ANG_OF]; - lj = bas[jsh0*BAS_SLOTS+ANG_OF]; - iprim = bas[ish0*BAS_SLOTS+NPRIM_OF]; - jprim = bas[jsh0*BAS_SLOTS+NPRIM_OF]; - gout_stride = gout_stride_lookup[li*L_AUX1+lj]; - nsp_per_block = THREADS / gout_stride; - } - __syncthreads(); - int sp_id = thread_id % nsp_per_block; - int gout_id = thread_id / nsp_per_block; - int g_size = (li + 1) * (lj + 6); - int gx_len = g_size * nsp_per_block; - extern __shared__ double g[]; - double *gx = g + sp_id; - double *gy = g + gx_len + sp_id; - double *gz = g + gx_len * 2 + sp_id; - double *rjri = g + gx_len * 3 + sp_id; - int idx_i = lex_xyz_offset(li); - int idx_j = lex_xyz_offset(lj); - if (gout_id == 0) { - gx[0] = PI_POW_1_5; - gy[0] = 1.; - } - - for (int pair_ij = shl_pair0+sp_id; pair_ij < shl_pair1+sp_id; pair_ij += nsp_per_block) { - double goutx[GOUT_WIDTH_IP1]; - double gouty[GOUT_WIDTH_IP1]; - double goutz[GOUT_WIDTH_IP1]; -#pragma unroll - for (int n = 0; n < GOUT_WIDTH_IP1; ++n) { - goutx[n] = 0.; - gouty[n] = 0.; - goutz[n] = 0.; - } - int bas_ij; - if (pair_ij >= shl_pair1) { - bas_ij = bas_ij_idx[shl_pair0]; - } else { - bas_ij = bas_ij_idx[pair_ij]; - } - int ish = bas_ij / nbas; - int jsh = bas_ij % nbas; - int ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; - int rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; - int expi = bas[ish*BAS_SLOTS+PTR_EXP]; - int expj = bas[jsh*BAS_SLOTS+PTR_EXP]; - int ci = bas[ish*BAS_SLOTS+PTR_COEFF]; - int cj = bas[jsh*BAS_SLOTS+PTR_COEFF]; - for (int img = 0; img < envs.nimgs; img++) { - if (gout_id == 0) { - double xjL = img_coords[img*3+0]; - double yjL = img_coords[img*3+1]; - double zjL = img_coords[img*3+2]; - double xjxi = env[rj+0] + xjL - env[ri+0]; - double yjyi = env[rj+1] + yjL - env[ri+1]; - double zjzi = env[rj+2] + zjL - env[ri+2]; - double rr_ij = xjxi*xjxi + yjyi*yjyi + zjzi*zjzi; - rjri[0*nsp_per_block] = xjxi; - rjri[1*nsp_per_block] = yjyi; - rjri[2*nsp_per_block] = zjzi; - rjri[3*nsp_per_block] = rr_ij; - } - double Bx = env[rj+0] + img_coords[img*3+0] - env[ri+0]; - double By = env[rj+1] + img_coords[img*3+1] - env[ri+1]; - double Bz = env[rj+2] + img_coords[img*3+2] - env[ri+2]; - double Bx2 = Bx*Bx, By2 = By*By, Bz2 = Bz*Bz; - int ijprim = iprim * jprim; - for (int ijp = 0; ijp < ijprim; ++ijp) { - __syncthreads(); - int ip = ijp % iprim; - int jp = ijp / iprim; - double ai = env[expi+ip]; - double aj = env[expj+jp]; - double cicj = env[ci+ip] * env[cj+jp]; - vrr_hrr(gx, rjri, ai, aj, cicj, li, lj+5, gout_id, gout_stride, - nsp_per_block); - if (pair_ij >= shl_pair1) { - continue; - } - int nsp = nsp_per_block; - int stride_j = li + 1; - double aj2 = aj * -2; - float div_nfi = c_div_nf[li]; - int nfi = c_nf[li]; - int nfj = c_nf[lj]; - int nfij = nfi * nfj; -#pragma unroll - for (int n = 0; n < GOUT_WIDTH_IP1; ++n) { - uint32_t ij = gout_id + n * gout_stride; - if (ij >= nfij) break; - uint32_t j = ij * div_nfi; - uint32_t i = ij - j * nfi; - 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]; - #define SADDR(g,a,joff) ((g)[((a) + (joff)*stride_j) * nsp]) - double Sxm1 = (jx > 0) ? SADDR(gx,ix,jx-1) : 0.; - double Sx0 = SADDR(gx,ix,jx); double Sx1 = SADDR(gx,ix,jx+1); - double Sx2 = SADDR(gx,ix,jx+2); double Sx3 = SADDR(gx,ix,jx+3); - double Sx4 = SADDR(gx,ix,jx+4); double Sx5 = SADDR(gx,ix,jx+5); - double Sym1 = (jy > 0) ? SADDR(gy,iy,jy-1) : 0.; - double Sy0 = SADDR(gy,iy,jy); double Sy1 = SADDR(gy,iy,jy+1); - double Sy2 = SADDR(gy,iy,jy+2); double Sy3 = SADDR(gy,iy,jy+3); - double Sy4 = SADDR(gy,iy,jy+4); double Sy5 = SADDR(gy,iy,jy+5); - double Szm1 = (jz > 0) ? SADDR(gz,iz,jz-1) : 0.; - double Sz0 = SADDR(gz,iz,jz); double Sz1 = SADDR(gz,iz,jz+1); - double Sz2 = SADDR(gz,iz,jz+2); double Sz3 = SADDR(gz,iz,jz+3); - double Sz4 = SADDR(gz,iz,jz+4); double Sz5 = SADDR(gz,iz,jz+5); - #undef SADDR - // Undifferentiated overlaps, x^2 moments, x^4 moments: - double sx = Sx0, sy = Sy0, sz = Sz0; - double mx2 = Sx2 + 2.*Bx*Sx1 + Bx2*Sx0; - double my2 = Sy2 + 2.*By*Sy1 + By2*Sy0; - double mz2 = Sz2 + 2.*Bz*Sz1 + Bz2*Sz0; - double mx4 = Sx4 + 4.*Bx*Sx3 + 6.*Bx2*Sx2 + 4.*Bx2*Bx*Sx1 + Bx2*Bx2*Sx0; - double my4 = Sy4 + 4.*By*Sy3 + 6.*By2*Sy2 + 4.*By2*By*Sy1 + By2*By2*Sy0; - double mz4 = Sz4 + 4.*Bz*Sz3 + 6.*Bz2*Sz2 + 4.*Bz2*Bz*Sz1 + Bz2*Bz2*Sz0; - // j-derivative of overlaps: - double Dsx = aj2*Sx1 + jx*Sxm1; - double Dsy = aj2*Sy1 + jy*Sym1; - double Dsz = aj2*Sz1 + jz*Szm1; - // j-derivative of x^2 moments: - double Dmx2 = aj2*(Sx3 + 2.*Bx*Sx2 + Bx2*Sx1) - + jx*(Sx1 + 2.*Bx*Sx0 + Bx2*Sxm1); - double Dmy2 = aj2*(Sy3 + 2.*By*Sy2 + By2*Sy1) - + jy*(Sy1 + 2.*By*Sy0 + By2*Sym1); - double Dmz2 = aj2*(Sz3 + 2.*Bz*Sz2 + Bz2*Sz1) - + jz*(Sz1 + 2.*Bz*Sz0 + Bz2*Szm1); - // j-derivative of x^4 moments: - double Dmx4 = aj2*(Sx5+4.*Bx*Sx4+6.*Bx2*Sx3+4.*Bx2*Bx*Sx2+Bx2*Bx2*Sx1) - + jx*(Sx3+4.*Bx*Sx2+6.*Bx2*Sx1+4.*Bx2*Bx*Sx0+Bx2*Bx2*Sxm1); - double Dmy4 = aj2*(Sy5+4.*By*Sy4+6.*By2*Sy3+4.*By2*By*Sy2+By2*By2*Sy1) - + jy*(Sy3+4.*By*Sy2+6.*By2*Sy1+4.*By2*By*Sy0+By2*By2*Sym1); - double Dmz4 = aj2*(Sz5+4.*Bz*Sz4+6.*Bz2*Sz3+4.*Bz2*Bz*Sz2+Bz2*Bz2*Sz1) - + jz*(Sz3+4.*Bz*Sz2+6.*Bz2*Sz1+4.*Bz2*Bz*Sz0+Bz2*Bz2*Szm1); - // r^4 = x^4+y^4+z^4 + 2(x^2*y^2+y^2*z^2+x^2*z^2) - // d/dBx: acts on x-axis j-basis - // d/dBx[x^4*sy*sz] = Dmx4*sy*sz - // d/dBx[sx*y^4*sz] = Dsx*my4*sz (etc.) - // d/dBx[2*x^2*y^2*sz] = 2*Dmx2*my2*sz - // d/dBx[2*sx*y^2*z^2] = 2*Dsx*my2*mz2 (etc.) - // d/dBx[2*x^2*sz*z^2] = 2*Dmx2*sy*mz2 - goutx[n] += Dmx4*sy*sz + Dsx*my4*sz + Dsx*sy*mz4 - + 2.*(Dmx2*my2*sz + Dsx*my2*mz2 + Dmx2*sy*mz2); - gouty[n] += mx4*Dsy*sz + sx*Dmy4*sz + sx*Dsy*mz4 - + 2.*(mx2*Dmy2*sz + sx*Dmy2*mz2 + mx2*Dsy*mz2); - goutz[n] += mx4*sy*Dsz + sx*my4*Dsz + sx*sy*Dmz4 - + 2.*(mx2*my2*Dsz + sx*my2*Dmz2 + mx2*sy*Dmz2); - } - } - } - - if (pair_ij < shl_pair1) { int *ao_loc = envs.ao_loc; int nbas = envs.cell0_nbas; - size_t nao2 = naoi * naoj; int cell_id = jsh / nbas; int jshp = jsh % nbas; int i0 = ao_loc[ish]; int j0 = ao_loc[jshp]; - double *outx = out + cell_id*nao2*3 + i0 * naoj + j0 - ij_offset; - double *outy = outx + nao2; - double *outz = outx + nao2 * 2; - int nfi = c_nf[li]; - int nfj = c_nf[lj]; - int nfij = nfi * nfj; + double *out_subblock = out + (cell_id*naoi+i0) * naoj + j0 - ij_offset; #pragma unroll - for (int n = 0; n < GOUT_WIDTH_IP1; ++n) { + for (int n = 0; n < GOUT_WIDTH; ++n) { int ij = n*gout_stride+gout_id; if (ij >= nfij) break; int j = ij / nfi; int i = ij % nfi; - outx[i*naoj+j] = goutx[n]; - outy[i*naoj+j] = gouty[n]; - outz[i*naoj+j] = goutz[n]; + out_subblock[i*naoj+j] = gout[n]; } } } @@ -1163,6 +375,7 @@ void int1e_ipovlp_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, gouty[n] = 0.; goutz[n] = 0.; } + __syncthreads(); int bas_ij; if (pair_ij >= shl_pair1) { bas_ij = bas_ij_idx[shl_pair0]; @@ -1178,6 +391,7 @@ void int1e_ipovlp_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, int ci = bas[ish*BAS_SLOTS+PTR_COEFF]; int cj = bas[jsh*BAS_SLOTS+PTR_COEFF]; for (int img = 0; img < envs.nimgs; img++) { + __syncthreads(); if (gout_id == 0) { double xjL = img_coords[img*3+0]; double yjL = img_coords[img*3+1]; @@ -1986,74 +1200,6 @@ int PBCint1e_kin(double *out, PBCIntEnvVars *envs, int shm_size, return 0; } -int PBCint1e_r2_origi(double *out, PBCIntEnvVars *envs, int shm_size, - int nbatches_shl_pair, int *bas_ij_idx, - int *shl_pair_offsets, int *gout_stride_lookup, - int naoi, int naoj, size_t ij_offset) -{ - cudaFuncSetAttribute(int1e_r2_origi_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); - int1e_r2_origi_kernel<<>>( - out, *envs, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, - naoi, naoj, ij_offset); - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error in int1e_r2_origi kernel: %s\n", cudaGetErrorString(err)); - return 1; - } - return 0; -} - -int PBCint1e_r4_origi(double *out, PBCIntEnvVars *envs, int shm_size, - int nbatches_shl_pair, int *bas_ij_idx, - int *shl_pair_offsets, int *gout_stride_lookup, - int naoi, int naoj, size_t ij_offset) -{ - cudaFuncSetAttribute(int1e_r4_origi_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); - int1e_r4_origi_kernel<<>>( - out, *envs, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, - naoi, naoj, ij_offset); - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error in int1e_r4_origi kernel: %s\n", cudaGetErrorString(err)); - return 1; - } - return 0; -} - -int PBCint1e_r2_origi_ip2(double *out, PBCIntEnvVars *envs, int shm_size, - int nbatches_shl_pair, int *bas_ij_idx, - int *shl_pair_offsets, int *gout_stride_lookup, - int naoi, int naoj, size_t ij_offset) -{ - cudaFuncSetAttribute(int1e_r2_origi_ip2_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); - int1e_r2_origi_ip2_kernel<<>>( - out, *envs, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, - naoi, naoj, ij_offset); - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error in int1e_r2_origi_ip2 kernel: %s\n", cudaGetErrorString(err)); - return 1; - } - return 0; -} - -int PBCint1e_r4_origi_ip2(double *out, PBCIntEnvVars *envs, int shm_size, - int nbatches_shl_pair, int *bas_ij_idx, - int *shl_pair_offsets, int *gout_stride_lookup, - int naoi, int naoj, size_t ij_offset) -{ - cudaFuncSetAttribute(int1e_r4_origi_ip2_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); - int1e_r4_origi_ip2_kernel<<>>( - out, *envs, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, - naoi, naoj, ij_offset); - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - fprintf(stderr, "CUDA Error in int1e_r4_origi_ip2 kernel: %s\n", cudaGetErrorString(err)); - return 1; - } - return 0; -} - int PBCint1e_ipovlp(double *out, PBCIntEnvVars *envs, int shm_size, int nbatches_shl_pair, int *bas_ij_idx, int *shl_pair_offsets, int *gout_stride_lookup, diff --git a/gpu4pyscf/lib/pbc/ppnl.cu b/gpu4pyscf/lib/pbc/ppnl.cu new file mode 100644 index 000000000..2c2caf135 --- /dev/null +++ b/gpu4pyscf/lib/pbc/ppnl.cu @@ -0,0 +1,1033 @@ +/* + * 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. + */ + +// Included by overlap.cu to share vrr_hrr and its constants. +// All three radial channels contract each image directly into nuclear and +// strain derivatives, without allocating derivative integral tensors. + +#include +#include +#include +#include +#include +#include "gvhf-rys/vhf.cuh" +#include "gvhf-rys/rys_contract_k.cuh" +#include "pbc.cuh" +#include "int3c2e.cuh" +#include "overlap.cuh" + +#define GOUT_WIDTH 36 +#define GOUT_WIDTH_IP1 18 + +// : raise projector-side powers by two, with A the i center. +__global__ static +void int1e_r2_origi_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, + int *shl_pair_offsets, int *gout_stride_lookup, + int naoi, int naoj, size_t ij_offset) +{ + int sp_block_id = blockIdx.x; + int thread_id = threadIdx.x; + int nbas = envs.cell0_nbas * envs.bvk_ncells; + int *bas = envs.bas; + double *env = envs.env; + double *img_coords = envs.img_coords; + __shared__ int shl_pair0, shl_pair1; + __shared__ int li, lj, iprim, jprim; + __shared__ int gout_stride, nsp_per_block; + if (thread_id == 0) { + shl_pair0 = shl_pair_offsets[sp_block_id]; + shl_pair1 = shl_pair_offsets[sp_block_id+1]; + int bas_ij0 = bas_ij_idx[shl_pair0]; + int ish0 = bas_ij0 / nbas; + int jsh0 = bas_ij0 % nbas; + li = bas[ish0*BAS_SLOTS+ANG_OF]; + lj = bas[jsh0*BAS_SLOTS+ANG_OF]; + iprim = bas[ish0*BAS_SLOTS+NPRIM_OF]; + jprim = bas[jsh0*BAS_SLOTS+NPRIM_OF]; + gout_stride = gout_stride_lookup[li*L_AUX1+lj]; + nsp_per_block = THREADS / gout_stride; + } + __syncthreads(); + int sp_id = thread_id % nsp_per_block; + int gout_id = thread_id / nsp_per_block; + int stride_j = li + 3; + int g_size = stride_j * (lj + 1); + int gx_len = g_size * nsp_per_block; + extern __shared__ double g[]; + double *gx = g + sp_id; + double *gy = g + gx_len + sp_id; + double *gz = g + gx_len * 2 + sp_id; + double *rjri = g + gx_len * 3 + sp_id; + int idx_i = lex_xyz_offset(li); + int idx_j = lex_xyz_offset(lj); + if (gout_id == 0) { + gx[0] = PI_POW_1_5; + gy[0] = 1.; + } + + for (int pair_ij = shl_pair0+sp_id; pair_ij < shl_pair1+sp_id; pair_ij += nsp_per_block) { + double gout[GOUT_WIDTH]; +#pragma unroll + for (int n = 0; n < GOUT_WIDTH; ++n) { + gout[n] = 0.; + } + int bas_ij; + if (pair_ij >= shl_pair1) { + bas_ij = bas_ij_idx[shl_pair0]; + } else { + bas_ij = bas_ij_idx[pair_ij]; + } + int ish = bas_ij / nbas; + int jsh = bas_ij % nbas; + int ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; + int rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; + int expi = bas[ish*BAS_SLOTS+PTR_EXP]; + int expj = bas[jsh*BAS_SLOTS+PTR_EXP]; + int ci = bas[ish*BAS_SLOTS+PTR_COEFF]; + int cj = bas[jsh*BAS_SLOTS+PTR_COEFF]; + for (int img = 0; img < envs.nimgs; img++) { + __syncthreads(); + if (gout_id == 0) { + double xjL = img_coords[img*3+0]; + double yjL = img_coords[img*3+1]; + double zjL = img_coords[img*3+2]; + double xjxi = env[rj+0] + xjL - env[ri+0]; + double yjyi = env[rj+1] + yjL - env[ri+1]; + double zjzi = env[rj+2] + zjL - env[ri+2]; + double rr_ij = xjxi*xjxi + yjyi*yjyi + zjzi*zjzi; + rjri[0*nsp_per_block] = xjxi; + rjri[1*nsp_per_block] = yjyi; + rjri[2*nsp_per_block] = zjzi; + rjri[3*nsp_per_block] = rr_ij; + } + int ijprim = iprim * jprim; + for (int ijp = 0; ijp < ijprim; ++ijp) { + __syncthreads(); + int ip = ijp % iprim; + int jp = ijp / iprim; + double ai = env[expi+ip]; + double aj = env[expj+jp]; + double cicj = env[ci+ip] * env[cj+jp]; + vrr_hrr(gx, rjri, ai, aj, cicj, li+2, lj, gout_id, gout_stride, + nsp_per_block); + if (pair_ij >= shl_pair1) { + continue; + } + int nsp = nsp_per_block; + int i_1 = nsp_per_block; + int stride_j = li + 3; + int nfi = c_nf[li]; + int nfj = c_nf[lj]; + int nfij = nfi * nfj; + float div_nfi = c_div_nf[li]; +#pragma unroll + for (int n = 0; n < GOUT_WIDTH; ++n) { + uint32_t ij = gout_id + n * gout_stride; + if (ij >= nfij) break; + uint32_t j = ij * div_nfi; + uint32_t i = ij - j * nfi; + 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 + jx*stride_j) * nsp; + int addry = (iy + jy*stride_j) * nsp; + int addrz = (iz + jz*stride_j) * nsp; + double sx = gx[addrx]; + double sy = gy[addry]; + double sz = gz[addrz]; + double mx = gx[addrx+i_1*2]; + double my = gy[addry+i_1*2]; + double mz = gz[addrz+i_1*2]; + gout[n] += mx*sy*sz + sx*my*sz + sx*sy*mz; + } + } + } + + if (pair_ij < shl_pair1) { + int nfi = c_nf[li]; + int nfj = c_nf[lj]; + int nfij = nfi * nfj; + int *ao_loc = envs.ao_loc; + int nbas = envs.cell0_nbas; + int cell_id = jsh / nbas; + int jshp = jsh % nbas; + int i0 = ao_loc[ish]; + int j0 = ao_loc[jshp]; + double *out_subblock = out + (cell_id*naoi+i0) * naoj + j0 - ij_offset; +#pragma unroll + for (int n = 0; n < GOUT_WIDTH; ++n) { + int ij = n*gout_stride+gout_id; + if (ij >= nfij) break; + int j = ij / nfi; + int i = ij % nfi; + out_subblock[i*naoj+j] = gout[n]; + } + } + } +} + +// : raise projector-side powers by four, with A the i center. +// r^4 = x^4+y^4+z^4 + 2(x^2*y^2+y^2*z^2+x^2*z^2). +__global__ static +void int1e_r4_origi_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, + int *shl_pair_offsets, int *gout_stride_lookup, + int naoi, int naoj, size_t ij_offset) +{ + int sp_block_id = blockIdx.x; + int thread_id = threadIdx.x; + int nbas = envs.cell0_nbas * envs.bvk_ncells; + int *bas = envs.bas; + double *env = envs.env; + double *img_coords = envs.img_coords; + __shared__ int shl_pair0, shl_pair1; + __shared__ int li, lj, iprim, jprim; + __shared__ int gout_stride, nsp_per_block; + if (thread_id == 0) { + shl_pair0 = shl_pair_offsets[sp_block_id]; + shl_pair1 = shl_pair_offsets[sp_block_id+1]; + int bas_ij0 = bas_ij_idx[shl_pair0]; + int ish0 = bas_ij0 / nbas; + int jsh0 = bas_ij0 % nbas; + li = bas[ish0*BAS_SLOTS+ANG_OF]; + lj = bas[jsh0*BAS_SLOTS+ANG_OF]; + iprim = bas[ish0*BAS_SLOTS+NPRIM_OF]; + jprim = bas[jsh0*BAS_SLOTS+NPRIM_OF]; + gout_stride = gout_stride_lookup[li*L_AUX1+lj]; + nsp_per_block = THREADS / gout_stride; + } + __syncthreads(); + int sp_id = thread_id % nsp_per_block; + int gout_id = thread_id / nsp_per_block; + int stride_j = li + 5; + int g_size = stride_j * (lj + 1); + int gx_len = g_size * nsp_per_block; + extern __shared__ double g[]; + double *gx = g + sp_id; + double *gy = g + gx_len + sp_id; + double *gz = g + gx_len * 2 + sp_id; + double *rjri = g + gx_len * 3 + sp_id; + int idx_i = lex_xyz_offset(li); + int idx_j = lex_xyz_offset(lj); + if (gout_id == 0) { + gx[0] = PI_POW_1_5; + gy[0] = 1.; + } + + for (int pair_ij = shl_pair0+sp_id; pair_ij < shl_pair1+sp_id; pair_ij += nsp_per_block) { + double gout[GOUT_WIDTH]; +#pragma unroll + for (int n = 0; n < GOUT_WIDTH; ++n) { + gout[n] = 0.; + } + int bas_ij; + if (pair_ij >= shl_pair1) { + bas_ij = bas_ij_idx[shl_pair0]; + } else { + bas_ij = bas_ij_idx[pair_ij]; + } + int ish = bas_ij / nbas; + int jsh = bas_ij % nbas; + int ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; + int rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; + int expi = bas[ish*BAS_SLOTS+PTR_EXP]; + int expj = bas[jsh*BAS_SLOTS+PTR_EXP]; + int ci = bas[ish*BAS_SLOTS+PTR_COEFF]; + int cj = bas[jsh*BAS_SLOTS+PTR_COEFF]; + for (int img = 0; img < envs.nimgs; img++) { + __syncthreads(); + if (gout_id == 0) { + double xjL = img_coords[img*3+0]; + double yjL = img_coords[img*3+1]; + double zjL = img_coords[img*3+2]; + double xjxi = env[rj+0] + xjL - env[ri+0]; + double yjyi = env[rj+1] + yjL - env[ri+1]; + double zjzi = env[rj+2] + zjL - env[ri+2]; + double rr_ij = xjxi*xjxi + yjyi*yjyi + zjzi*zjzi; + rjri[0*nsp_per_block] = xjxi; + rjri[1*nsp_per_block] = yjyi; + rjri[2*nsp_per_block] = zjzi; + rjri[3*nsp_per_block] = rr_ij; + } + int ijprim = iprim * jprim; + for (int ijp = 0; ijp < ijprim; ++ijp) { + __syncthreads(); + int ip = ijp % iprim; + int jp = ijp / iprim; + double ai = env[expi+ip]; + double aj = env[expj+jp]; + double cicj = env[ci+ip] * env[cj+jp]; + vrr_hrr(gx, rjri, ai, aj, cicj, li+4, lj, gout_id, gout_stride, + nsp_per_block); + if (pair_ij >= shl_pair1) { + continue; + } + int nsp = nsp_per_block; + int i_1 = nsp_per_block; + int stride_j = li + 5; + int nfi = c_nf[li]; + int nfj = c_nf[lj]; + int nfij = nfi * nfj; + float div_nfi = c_div_nf[li]; +#pragma unroll + for (int n = 0; n < GOUT_WIDTH; ++n) { + uint32_t ij = gout_id + n * gout_stride; + if (ij >= nfij) break; + uint32_t j = ij * div_nfi; + uint32_t i = ij - j * nfi; + 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 + jx*stride_j) * nsp; + int addry = (iy + jy*stride_j) * nsp; + int addrz = (iz + jz*stride_j) * nsp; + double Sx0 = gx[addrx+i_1*0]; + double Sx2 = gx[addrx+i_1*2]; + double Sx4 = gx[addrx+i_1*4]; + double Sy0 = gy[addry+i_1*0]; + double Sy2 = gy[addry+i_1*2]; + double Sy4 = gy[addry+i_1*4]; + double Sz0 = gz[addrz+i_1*0]; + double Sz2 = gz[addrz+i_1*2]; + double Sz4 = gz[addrz+i_1*4]; + // r^4 = x^4+y^4+z^4 + 2(x^2*y^2 + y^2*z^2 + x^2*z^2) + gout[n] += Sx4*Sy0*Sz0 + Sx0*Sy4*Sz0 + Sx0*Sy0*Sz4 + + 2.*(Sx2*Sy2*Sz0 + Sx0*Sy2*Sz2 + Sx2*Sy0*Sz2); + } + } + } + + if (pair_ij < shl_pair1) { + int nfi = c_nf[li]; + int nfj = c_nf[lj]; + int nfij = nfi * nfj; + int *ao_loc = envs.ao_loc; + int nbas = envs.cell0_nbas; + int cell_id = jsh / nbas; + int jshp = jsh % nbas; + int i0 = ao_loc[ish]; + int j0 = ao_loc[jshp]; + double *out_subblock = out + (cell_id*naoi+i0) * naoj + j0 - ij_offset; +#pragma unroll + for (int n = 0; n < GOUT_WIDTH; ++n) { + int ij = n*gout_stride+gout_id; + if (ij >= nfij) break; + int j = ij / nfi; + int i = ij % nfi; + out_subblock[i*naoj+j] = gout[n]; + } + } + } +} + +// , with A the projector center. +// Raise i by 2 for the moment and j by one for its spatial derivative. +__global__ static +void int1e_r2_origi_ip2_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, + int *shl_pair_offsets, int *gout_stride_lookup, + int naoi, int naoj, size_t ij_offset) +{ + int sp_block_id = blockIdx.x; + int thread_id = threadIdx.x; + int nbas = envs.cell0_nbas * envs.bvk_ncells; + int *bas = envs.bas; + double *env = envs.env; + double *img_coords = envs.img_coords; + __shared__ int shl_pair0, shl_pair1; + __shared__ int li, lj, iprim, jprim; + __shared__ int gout_stride, nsp_per_block; + if (thread_id == 0) { + shl_pair0 = shl_pair_offsets[sp_block_id]; + shl_pair1 = shl_pair_offsets[sp_block_id+1]; + int bas_ij0 = bas_ij_idx[shl_pair0]; + int ish0 = bas_ij0 / nbas; + int jsh0 = bas_ij0 % nbas; + li = bas[ish0*BAS_SLOTS+ANG_OF]; + lj = bas[jsh0*BAS_SLOTS+ANG_OF]; + iprim = bas[ish0*BAS_SLOTS+NPRIM_OF]; + jprim = bas[jsh0*BAS_SLOTS+NPRIM_OF]; + gout_stride = gout_stride_lookup[li*L_AUX1+lj]; + nsp_per_block = THREADS / gout_stride; + } + __syncthreads(); + int sp_id = thread_id % nsp_per_block; + int gout_id = thread_id / nsp_per_block; + int stride_j = li + 3; + int g_size = stride_j * (lj + 2); + int gx_len = g_size * nsp_per_block; + extern __shared__ double g[]; + double *gx = g + sp_id; + double *gy = g + gx_len + sp_id; + double *gz = g + gx_len * 2 + sp_id; + double *rjri = g + gx_len * 3 + sp_id; + int idx_i = lex_xyz_offset(li); + int idx_j = lex_xyz_offset(lj); + if (gout_id == 0) { + gx[0] = PI_POW_1_5; + gy[0] = 1.; + } + + for (int pair_ij = shl_pair0+sp_id; pair_ij < shl_pair1+sp_id; pair_ij += nsp_per_block) { + double goutx[GOUT_WIDTH_IP1]; + double gouty[GOUT_WIDTH_IP1]; + double goutz[GOUT_WIDTH_IP1]; +#pragma unroll + for (int n = 0; n < GOUT_WIDTH_IP1; ++n) { + goutx[n] = 0.; + gouty[n] = 0.; + goutz[n] = 0.; + } + __syncthreads(); + int bas_ij; + if (pair_ij >= shl_pair1) { + bas_ij = bas_ij_idx[shl_pair0]; + } else { + bas_ij = bas_ij_idx[pair_ij]; + } + int ish = bas_ij / nbas; + int jsh = bas_ij % nbas; + int ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; + int rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; + int expi = bas[ish*BAS_SLOTS+PTR_EXP]; + int expj = bas[jsh*BAS_SLOTS+PTR_EXP]; + int ci = bas[ish*BAS_SLOTS+PTR_COEFF]; + int cj = bas[jsh*BAS_SLOTS+PTR_COEFF]; + for (int img = 0; img < envs.nimgs; img++) { + __syncthreads(); + if (gout_id == 0) { + double xjL = img_coords[img*3+0]; + double yjL = img_coords[img*3+1]; + double zjL = img_coords[img*3+2]; + double xjxi = env[rj+0] + xjL - env[ri+0]; + double yjyi = env[rj+1] + yjL - env[ri+1]; + double zjzi = env[rj+2] + zjL - env[ri+2]; + double rr_ij = xjxi*xjxi + yjyi*yjyi + zjzi*zjzi; + rjri[0*nsp_per_block] = xjxi; + rjri[1*nsp_per_block] = yjyi; + rjri[2*nsp_per_block] = zjzi; + rjri[3*nsp_per_block] = rr_ij; + } + int ijprim = iprim * jprim; + for (int ijp = 0; ijp < ijprim; ++ijp) { + __syncthreads(); + int ip = ijp % iprim; + int jp = ijp / iprim; + double ai = env[expi+ip]; + double aj = env[expj+jp]; + double cicj = env[ci+ip] * env[cj+jp]; + vrr_hrr(gx, rjri, ai, aj, cicj, li+2, lj+1, gout_id, gout_stride, + nsp_per_block); + if (pair_ij >= shl_pair1) { + continue; + } + int nsp = nsp_per_block; + int stride_j = li + 3; + int j_1 = stride_j * nsp; + double aj2 = aj * -2; + float div_nfi = c_div_nf[li]; + int nfi = c_nf[li]; + int nfj = c_nf[lj]; + int nfij = nfi * nfj; +#pragma unroll + for (int n = 0; n < GOUT_WIDTH_IP1; ++n) { + uint32_t ij = gout_id + n * gout_stride; + if (ij >= nfij) break; + uint32_t j = ij * div_nfi; + uint32_t i = ij - j * nfi; + 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 + jx*stride_j) * nsp; + int addry = (iy + jy*stride_j) * nsp; + int addrz = (iz + jz*stride_j) * nsp; + double sx = gx[addrx]; + double sy = gy[addry]; + double sz = gz[addrz]; + double Dsx = aj2*gx[addrx+j_1]; + double Dsy = aj2*gy[addry+j_1]; + double Dsz = aj2*gz[addrz+j_1]; + if (jx > 0) Dsx += jx*gx[addrx-j_1]; + if (jy > 0) Dsy += jy*gy[addry-j_1]; + if (jz > 0) Dsz += jz*gz[addrz-j_1]; + double mx = gx[addrx+2*nsp]; + double my = gy[addry+2*nsp]; + double mz = gz[addrz+2*nsp]; + double Dmx = aj2*gx[addrx+2*nsp+j_1]; + double Dmy = aj2*gy[addry+2*nsp+j_1]; + double Dmz = aj2*gz[addrz+2*nsp+j_1]; + if (jx > 0) Dmx += jx*gx[addrx+2*nsp-j_1]; + if (jy > 0) Dmy += jy*gy[addry+2*nsp-j_1]; + if (jz > 0) Dmz += jz*gz[addrz+2*nsp-j_1]; + goutx[n] += Dmx*sy*sz + Dsx*(my*sz + sy*mz); + gouty[n] += Dmy*sx*sz + Dsy*(mx*sz + sx*mz); + goutz[n] += Dmz*sx*sy + Dsz*(mx*sy + sx*my); + } + } + } + + if (pair_ij < shl_pair1) { + int *ao_loc = envs.ao_loc; + int nbas = envs.cell0_nbas; + size_t nao2 = naoi * naoj; + int cell_id = jsh / nbas; + int jshp = jsh % nbas; + int i0 = ao_loc[ish]; + int j0 = ao_loc[jshp]; + double *outx = out + cell_id*nao2*3 + i0 * naoj + j0 - ij_offset; + double *outy = outx + nao2; + double *outz = outx + nao2 * 2; + int nfi = c_nf[li]; + int nfj = c_nf[lj]; + int nfij = nfi * nfj; +#pragma unroll + for (int n = 0; n < GOUT_WIDTH_IP1; ++n) { + int ij = n*gout_stride+gout_id; + if (ij >= nfij) break; + int j = ij / nfi; + int i = ij % nfi; + outx[i*naoj+j] = goutx[n]; + outy[i*naoj+j] = gouty[n]; + outz[i*naoj+j] = goutz[n]; + } + } + } +} + +// , with A the projector center. +// Raise i by 4 for the moment and j by one for its spatial derivative. +__global__ static +void int1e_r4_origi_ip2_kernel(double *out, PBCIntEnvVars envs, int *bas_ij_idx, + int *shl_pair_offsets, int *gout_stride_lookup, + int naoi, int naoj, size_t ij_offset) +{ + int sp_block_id = blockIdx.x; + int thread_id = threadIdx.x; + int nbas = envs.cell0_nbas * envs.bvk_ncells; + int *bas = envs.bas; + double *env = envs.env; + double *img_coords = envs.img_coords; + __shared__ int shl_pair0, shl_pair1; + __shared__ int li, lj, iprim, jprim; + __shared__ int gout_stride, nsp_per_block; + if (thread_id == 0) { + shl_pair0 = shl_pair_offsets[sp_block_id]; + shl_pair1 = shl_pair_offsets[sp_block_id+1]; + int bas_ij0 = bas_ij_idx[shl_pair0]; + int ish0 = bas_ij0 / nbas; + int jsh0 = bas_ij0 % nbas; + li = bas[ish0*BAS_SLOTS+ANG_OF]; + lj = bas[jsh0*BAS_SLOTS+ANG_OF]; + iprim = bas[ish0*BAS_SLOTS+NPRIM_OF]; + jprim = bas[jsh0*BAS_SLOTS+NPRIM_OF]; + gout_stride = gout_stride_lookup[li*L_AUX1+lj]; + nsp_per_block = THREADS / gout_stride; + } + __syncthreads(); + int sp_id = thread_id % nsp_per_block; + int gout_id = thread_id / nsp_per_block; + int stride_j = li + 5; + int g_size = stride_j * (lj + 2); + int gx_len = g_size * nsp_per_block; + extern __shared__ double g[]; + double *gx = g + sp_id; + double *gy = g + gx_len + sp_id; + double *gz = g + gx_len * 2 + sp_id; + double *rjri = g + gx_len * 3 + sp_id; + int idx_i = lex_xyz_offset(li); + int idx_j = lex_xyz_offset(lj); + if (gout_id == 0) { + gx[0] = PI_POW_1_5; + gy[0] = 1.; + } + + for (int pair_ij = shl_pair0+sp_id; pair_ij < shl_pair1+sp_id; pair_ij += nsp_per_block) { + double goutx[GOUT_WIDTH_IP1]; + double gouty[GOUT_WIDTH_IP1]; + double goutz[GOUT_WIDTH_IP1]; +#pragma unroll + for (int n = 0; n < GOUT_WIDTH_IP1; ++n) { + goutx[n] = 0.; + gouty[n] = 0.; + goutz[n] = 0.; + } + __syncthreads(); + int bas_ij; + if (pair_ij >= shl_pair1) { + bas_ij = bas_ij_idx[shl_pair0]; + } else { + bas_ij = bas_ij_idx[pair_ij]; + } + int ish = bas_ij / nbas; + int jsh = bas_ij % nbas; + int ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; + int rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; + int expi = bas[ish*BAS_SLOTS+PTR_EXP]; + int expj = bas[jsh*BAS_SLOTS+PTR_EXP]; + int ci = bas[ish*BAS_SLOTS+PTR_COEFF]; + int cj = bas[jsh*BAS_SLOTS+PTR_COEFF]; + for (int img = 0; img < envs.nimgs; img++) { + __syncthreads(); + if (gout_id == 0) { + double xjL = img_coords[img*3+0]; + double yjL = img_coords[img*3+1]; + double zjL = img_coords[img*3+2]; + double xjxi = env[rj+0] + xjL - env[ri+0]; + double yjyi = env[rj+1] + yjL - env[ri+1]; + double zjzi = env[rj+2] + zjL - env[ri+2]; + double rr_ij = xjxi*xjxi + yjyi*yjyi + zjzi*zjzi; + rjri[0*nsp_per_block] = xjxi; + rjri[1*nsp_per_block] = yjyi; + rjri[2*nsp_per_block] = zjzi; + rjri[3*nsp_per_block] = rr_ij; + } + int ijprim = iprim * jprim; + for (int ijp = 0; ijp < ijprim; ++ijp) { + __syncthreads(); + int ip = ijp % iprim; + int jp = ijp / iprim; + double ai = env[expi+ip]; + double aj = env[expj+jp]; + double cicj = env[ci+ip] * env[cj+jp]; + vrr_hrr(gx, rjri, ai, aj, cicj, li+4, lj+1, gout_id, gout_stride, + nsp_per_block); + if (pair_ij >= shl_pair1) { + continue; + } + int nsp = nsp_per_block; + int stride_j = li + 5; + int j_1 = stride_j * nsp; + double aj2 = aj * -2; + float div_nfi = c_div_nf[li]; + int nfi = c_nf[li]; + int nfj = c_nf[lj]; + int nfij = nfi * nfj; +#pragma unroll + for (int n = 0; n < GOUT_WIDTH_IP1; ++n) { + uint32_t ij = gout_id + n * gout_stride; + if (ij >= nfij) break; + uint32_t j = ij * div_nfi; + uint32_t i = ij - j * nfi; + 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 + jx*stride_j) * nsp; + int addry = (iy + jy*stride_j) * nsp; + int addrz = (iz + jz*stride_j) * nsp; + double sx = gx[addrx]; + double sy = gy[addry]; + double sz = gz[addrz]; + double Dsx = aj2*gx[addrx+j_1]; + double Dsy = aj2*gy[addry+j_1]; + double Dsz = aj2*gz[addrz+j_1]; + if (jx > 0) Dsx += jx*gx[addrx-j_1]; + if (jy > 0) Dsy += jy*gy[addry-j_1]; + if (jz > 0) Dsz += jz*gz[addrz-j_1]; + double mx2 = gx[addrx+2*nsp]; + double my2 = gy[addry+2*nsp]; + double mz2 = gz[addrz+2*nsp]; + double Dmx2 = aj2*gx[addrx+2*nsp+j_1]; + double Dmy2 = aj2*gy[addry+2*nsp+j_1]; + double Dmz2 = aj2*gz[addrz+2*nsp+j_1]; + if (jx > 0) Dmx2 += jx*gx[addrx+2*nsp-j_1]; + if (jy > 0) Dmy2 += jy*gy[addry+2*nsp-j_1]; + if (jz > 0) Dmz2 += jz*gz[addrz+2*nsp-j_1]; + double mx4 = gx[addrx+4*nsp]; + double my4 = gy[addry+4*nsp]; + double mz4 = gz[addrz+4*nsp]; + double Dmx4 = aj2*gx[addrx+4*nsp+j_1]; + double Dmy4 = aj2*gy[addry+4*nsp+j_1]; + double Dmz4 = aj2*gz[addrz+4*nsp+j_1]; + if (jx > 0) Dmx4 += jx*gx[addrx+4*nsp-j_1]; + if (jy > 0) Dmy4 += jy*gy[addry+4*nsp-j_1]; + if (jz > 0) Dmz4 += jz*gz[addrz+4*nsp-j_1]; + // r^4 = x^4+y^4+z^4 + 2(x^2*y^2+y^2*z^2+x^2*z^2) + goutx[n] += Dmx4*sy*sz + Dsx*my4*sz + Dsx*sy*mz4 + + 2.*(Dmx2*my2*sz + Dsx*my2*mz2 + Dmx2*sy*mz2); + gouty[n] += mx4*Dsy*sz + sx*Dmy4*sz + sx*Dsy*mz4 + + 2.*(mx2*Dmy2*sz + sx*Dmy2*mz2 + mx2*Dsy*mz2); + goutz[n] += mx4*sy*Dsz + sx*my4*Dsz + sx*sy*Dmz4 + + 2.*(mx2*my2*Dsz + sx*my2*Dmz2 + mx2*sy*Dmz2); + } + } + } + + if (pair_ij < shl_pair1) { + int *ao_loc = envs.ao_loc; + int nbas = envs.cell0_nbas; + size_t nao2 = naoi * naoj; + int cell_id = jsh / nbas; + int jshp = jsh % nbas; + int i0 = ao_loc[ish]; + int j0 = ao_loc[jshp]; + double *outx = out + cell_id*nao2*3 + i0 * naoj + j0 - ij_offset; + double *outy = outx + nao2; + double *outz = outx + nao2 * 2; + int nfi = c_nf[li]; + int nfj = c_nf[lj]; + int nfij = nfi * nfj; +#pragma unroll + for (int n = 0; n < GOUT_WIDTH_IP1; ++n) { + int ij = n*gout_stride+gout_id; + if (ij >= nfij) break; + int j = ij / nfi; + int i = ij % nfi; + outx[i*naoj+j] = goutx[n]; + outy[i*naoj+j] = gouty[n]; + outz[i*naoj+j] = goutz[n]; + } + } + } +} + +template +static __global__ +void ppnl_derivatives_kernel(double *grad, double *sigma, double *dm, PBCIntEnvVars envs, + int *shl_pair_offsets, int *bas_ij_idx, + int *gout_stride_lookup, int naoi, int naoj) +{ + int sp_block_id = blockIdx.x; + int thread_id = threadIdx.x; + int *bas = envs.bas; + int cell0_nbas = envs.cell0_nbas; + int nbas = envs.cell0_nbas * envs.bvk_ncells; + int *ao_loc = envs.ao_loc; + double *env = envs.env; + double *img_coords = envs.img_coords; + __shared__ int shl_pair0, shl_pair1; + __shared__ int li, lj, iprim, jprim; + __shared__ int gout_stride, nsp_per_block; + if (thread_id == 0) { + shl_pair0 = shl_pair_offsets[sp_block_id]; + shl_pair1 = shl_pair_offsets[sp_block_id+1]; + int bas_ij0 = bas_ij_idx[shl_pair0]; + int ish0 = bas_ij0 / nbas; + int jsh0 = bas_ij0 % nbas; + li = bas[ish0*BAS_SLOTS+ANG_OF]; + lj = bas[jsh0*BAS_SLOTS+ANG_OF]; + iprim = bas[ish0*BAS_SLOTS+NPRIM_OF]; + jprim = bas[jsh0*BAS_SLOTS+NPRIM_OF]; + gout_stride = gout_stride_lookup[li*L_AUX1+lj]; + nsp_per_block = THREADS / gout_stride; + } + __syncthreads(); + int sp_id = thread_id % nsp_per_block; + int gout_id = thread_id / nsp_per_block; + + int g_size = (li + RADIAL + 1) * (lj + 2); + int gx_len = g_size * nsp_per_block; + extern __shared__ double g[]; + double *gx = g + sp_id; + double *gy = g + gx_len + sp_id; + double *gz = g + gx_len * 2 + sp_id; + double *rjri = g + gx_len * 3 + sp_id; + if (gout_id == 0) { + gx[0] = PI_POW_1_5; + gy[0] = 1.; + } + int idx_i = lex_xyz_offset(li); + int idx_j = lex_xyz_offset(lj); + + double sigma_xx = 0; + double sigma_xy = 0; + double sigma_xz = 0; + double sigma_yx = 0; + double sigma_yy = 0; + double sigma_yz = 0; + double sigma_zx = 0; + double sigma_zy = 0; + double sigma_zz = 0; + for (int pair_ij = shl_pair0+sp_id; pair_ij < shl_pair1+sp_id; pair_ij += nsp_per_block) { + __syncthreads(); + int bas_ij = bas_ij_idx[shl_pair0]; + if (pair_ij < shl_pair1) { + bas_ij = bas_ij_idx[pair_ij]; + } + int ish = bas_ij / nbas; + int jsh = bas_ij - nbas * ish; + int cell_j = jsh / cell0_nbas; + int jsh_cell0 = jsh - cell0_nbas * cell_j; + int i0 = ao_loc[ish]; + int j0 = ao_loc[jsh_cell0] - naoi; + // Rectangular weights are stored as (image, projector AO, cell AO). + const double *dm_ij = dm + (size_t)cell_j*naoi*naoj + (size_t)i0*naoj + j0; + int ri = bas[ish*BAS_SLOTS+PTR_BAS_COORD]; + int rj = bas[jsh*BAS_SLOTS+PTR_BAS_COORD]; + int expi = bas[ish*BAS_SLOTS+PTR_EXP]; + int expj = bas[jsh*BAS_SLOTS+PTR_EXP]; + int ci = bas[ish*BAS_SLOTS+PTR_COEFF]; + int cj = bas[jsh*BAS_SLOTS+PTR_COEFF]; + double grad_ix = 0; + double grad_iy = 0; + double grad_iz = 0; + for (int img = 0; img < envs.nimgs; img++) { + double xi = env[ri+0]; + double yi = env[ri+1]; + double zi = env[ri+2]; + double xj = env[rj+0] + img_coords[img*3+0]; + double yj = env[rj+1] + img_coords[img*3+1]; + double zj = env[rj+2] + img_coords[img*3+2]; + __syncthreads(); + if (gout_id == 0) { + double xjxi = xj - xi; + double yjyi = yj - yi; + double zjzi = zj - zi; + double rr_ij = xjxi*xjxi + yjyi*yjyi + zjzi*zjzi; + rjri[0*nsp_per_block] = xjxi; + rjri[1*nsp_per_block] = yjyi; + rjri[2*nsp_per_block] = zjzi; + rjri[3*nsp_per_block] = rr_ij; + } + double v_ix = 0; + double v_iy = 0; + double v_iz = 0; + int ijprim = iprim * jprim; + for (int ijp = 0; ijp < ijprim; ++ijp) { + __syncthreads(); + int ip = ijp % iprim; + int jp = ijp / iprim; + double ai = env[expi+ip]; + double aj = env[expj+jp]; + double cicj = env[ci+ip] * env[cj+jp]; + vrr_hrr(gx, rjri, ai, aj, cicj, li+RADIAL, lj+1, gout_id, gout_stride, + nsp_per_block); + if (pair_ij >= shl_pair1) { + continue; + } + int stride_j = li + RADIAL + 1; + float div_nfi = c_div_nf[li]; + int nfi = c_nf[li]; + int nfj = c_nf[lj]; + int nfij = nfi * nfj; +#pragma unroll + for (int n = 0; n < GOUT_WIDTH; ++n) { + uint32_t ij = gout_id + n * gout_stride; + if (ij >= nfij) break; + uint32_t j = ij * div_nfi; + uint32_t i = ij - j * nfi; + 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]; + double v[3] = {0., 0., 0.}; + // Expand |r-A|^RADIAL on the projector (i) center. + // Translational invariance gives d/dA = . + // This includes motion of the radial moment's origin. +#pragma unroll + for (int rx = 0; rx <= RADIAL/2; ++rx) { +#pragma unroll + for (int ry = 0; ry <= RADIAL/2-rx; ++ry) { + int rz = RADIAL/2-rx-ry; + double weight = 1.; + if (RADIAL == 4 && (rx == 1 || ry == 1 || rz == 1)) { + weight = 2.; + } + int addrx = (ix + 2*rx + jx*stride_j) * nsp_per_block; + int addry = (iy + 2*ry + jy*stride_j) * nsp_per_block; + int addrz = (iz + 2*rz + jz*stride_j) * nsp_per_block; + int dj = stride_j * nsp_per_block; + double sx = gx[addrx]; + double sy = gy[addry]; + double sz = gz[addrz]; + double dx = -2*aj*gx[addrx+dj]; + double dy = -2*aj*gy[addry+dj]; + double dz = -2*aj*gz[addrz+dj]; + if (jx > 0) dx += jx*gx[addrx-dj]; + if (jy > 0) dy += jy*gy[addry-dj]; + if (jz > 0) dz += jz*gz[addrz-dj]; + v[0] += weight * dx*sy*sz; + v[1] += weight * sx*dy*sz; + v[2] += weight * sx*sy*dz; + } + } + double dm_val = dm_ij[(size_t)i*naoj+j]; + v_ix += v[0] * dm_val; + v_iy += v[1] * dm_val; + v_iz += v[2] * dm_val; + } + } + double xjxi = rjri[0*nsp_per_block]; + double yjyi = rjri[1*nsp_per_block]; + double zjzi = rjri[2*nsp_per_block]; + sigma_xx -= v_ix * xjxi; + sigma_xy -= v_ix * yjyi; + sigma_xz -= v_ix * zjzi; + sigma_yx -= v_iy * xjxi; + sigma_yy -= v_iy * yjyi; + sigma_yz -= v_iy * zjzi; + sigma_zx -= v_iz * xjxi; + sigma_zy -= v_iz * yjyi; + sigma_zz -= v_iz * zjzi; + grad_ix += v_ix; + grad_iy += v_iy; + grad_iz += v_iz; + } + int ish_cell0 = ish; + int ia = bas[ish_cell0*BAS_SLOTS+ATOM_OF]; + int ja = bas[jsh_cell0*BAS_SLOTS+ATOM_OF]; + atomicAdd(grad+ia*3+0, grad_ix); + atomicAdd(grad+ia*3+1, grad_iy); + atomicAdd(grad+ia*3+2, grad_iz); + atomicAdd(grad+ja*3+0, -grad_ix); + atomicAdd(grad+ja*3+1, -grad_iy); + atomicAdd(grad+ja*3+2, -grad_iz); + } + atomicAdd(sigma+0, sigma_xx); + atomicAdd(sigma+1, sigma_xy); + atomicAdd(sigma+2, sigma_xz); + atomicAdd(sigma+3, sigma_yx); + atomicAdd(sigma+4, sigma_yy); + atomicAdd(sigma+5, sigma_yz); + atomicAdd(sigma+6, sigma_zx); + atomicAdd(sigma+7, sigma_zy); + atomicAdd(sigma+8, sigma_zz); +} + +template +static int ppnl_derivatives(double *grad, double *sigma, double *dm, + PBCIntEnvVars *envs, int shm_size, int nbatches_shl_pair, + int *shl_pair_offsets, int *bas_ij_idx, + int *gout_stride_lookup, int naoi, int naoj) +{ + if (nbatches_shl_pair == 0) return 0; + cudaError_t err = cudaFuncSetAttribute( + ppnl_derivatives_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); + if (err == cudaSuccess) { + ppnl_derivatives_kernel<<>>( + grad, sigma, dm, *envs, shl_pair_offsets, bas_ij_idx, + gout_stride_lookup, naoi, naoj); + err = cudaGetLastError(); + } + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in ppnl derivatives (r^%d): %s\n", + RADIAL, cudaGetErrorString(err)); + return 1; + } + return 0; +} + +extern "C" { +int PBCint1e_r2_origi(double *out, PBCIntEnvVars *envs, int shm_size, + int nbatches_shl_pair, int *bas_ij_idx, + int *shl_pair_offsets, int *gout_stride_lookup, + int naoi, int naoj, size_t ij_offset) +{ + cudaFuncSetAttribute(int1e_r2_origi_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); + int1e_r2_origi_kernel<<>>( + out, *envs, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, + naoi, naoj, ij_offset); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in int1e_r2_origi kernel: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} + +int PBCint1e_r4_origi(double *out, PBCIntEnvVars *envs, int shm_size, + int nbatches_shl_pair, int *bas_ij_idx, + int *shl_pair_offsets, int *gout_stride_lookup, + int naoi, int naoj, size_t ij_offset) +{ + cudaFuncSetAttribute(int1e_r4_origi_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); + int1e_r4_origi_kernel<<>>( + out, *envs, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, + naoi, naoj, ij_offset); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in int1e_r4_origi kernel: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} + +int PBCint1e_r2_origi_ip2(double *out, PBCIntEnvVars *envs, int shm_size, + int nbatches_shl_pair, int *bas_ij_idx, + int *shl_pair_offsets, int *gout_stride_lookup, + int naoi, int naoj, size_t ij_offset) +{ + cudaFuncSetAttribute(int1e_r2_origi_ip2_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); + int1e_r2_origi_ip2_kernel<<>>( + out, *envs, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, + naoi, naoj, ij_offset); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in int1e_r2_origi_ip2 kernel: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} + +int PBCint1e_r4_origi_ip2(double *out, PBCIntEnvVars *envs, int shm_size, + int nbatches_shl_pair, int *bas_ij_idx, + int *shl_pair_offsets, int *gout_stride_lookup, + int naoi, int naoj, size_t ij_offset) +{ + cudaFuncSetAttribute(int1e_r4_origi_ip2_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shm_size); + int1e_r4_origi_ip2_kernel<<>>( + out, *envs, bas_ij_idx, shl_pair_offsets, gout_stride_lookup, + naoi, naoj, ij_offset); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "CUDA Error in int1e_r4_origi_ip2 kernel: %s\n", cudaGetErrorString(err)); + return 1; + } + return 0; +} + +int PBCovlp_cross_derivatives(double *grad, double *sigma, double *dm, + PBCIntEnvVars *envs, int shm_size, int nbatches_shl_pair, + int *shl_pair_offsets, int *bas_ij_idx, + int *gout_stride_lookup, int naoi, int naoj) +{ + return ppnl_derivatives<0>(grad, sigma, dm, envs, shm_size, + nbatches_shl_pair, shl_pair_offsets, bas_ij_idx, + gout_stride_lookup, naoi, naoj); +} + +int PBCint1e_r2_origi_derivatives(double *grad, double *sigma, double *dm, + PBCIntEnvVars *envs, int shm_size, int nbatches_shl_pair, + int *shl_pair_offsets, int *bas_ij_idx, + int *gout_stride_lookup, int naoi, int naoj) +{ + return ppnl_derivatives<2>(grad, sigma, dm, envs, shm_size, + nbatches_shl_pair, shl_pair_offsets, bas_ij_idx, + gout_stride_lookup, naoi, naoj); +} + +int PBCint1e_r4_origi_derivatives(double *grad, double *sigma, double *dm, + PBCIntEnvVars *envs, int shm_size, int nbatches_shl_pair, + int *shl_pair_offsets, int *bas_ij_idx, + int *gout_stride_lookup, int naoi, int naoj) +{ + return ppnl_derivatives<4>(grad, sigma, dm, envs, shm_size, + nbatches_shl_pair, shl_pair_offsets, bas_ij_idx, + gout_stride_lookup, naoi, naoj); +} + +} diff --git a/gpu4pyscf/pbc/grad/krhf.py b/gpu4pyscf/pbc/grad/krhf.py index ad3c61a71..5953b9519 100644 --- a/gpu4pyscf/pbc/grad/krhf.py +++ b/gpu4pyscf/pbc/grad/krhf.py @@ -34,8 +34,7 @@ from gpu4pyscf.pbc.gto import int1e from gpu4pyscf.pbc.scf.rsjk import PBCJKMatrixOpt from gpu4pyscf.pbc import tools as pbctools -from gpu4pyscf.pbc.grad.pp import ( - vppnl_nuc_grad, _get_pp_nonloc_strain_derivatives) +from gpu4pyscf.pbc.grad.pp import vppnl_nuc_grad from gpu4pyscf.pbc.grad.rhf import contract_h1e_dm, _get_ejk_derivatives from gpu4pyscf.pbc.grad import rhf as pbchf_grad diff --git a/gpu4pyscf/pbc/grad/kuhf.py b/gpu4pyscf/pbc/grad/kuhf.py index 046a553d4..a94ef30cb 100644 --- a/gpu4pyscf/pbc/grad/kuhf.py +++ b/gpu4pyscf/pbc/grad/kuhf.py @@ -23,8 +23,7 @@ from gpu4pyscf.lib import logger from gpu4pyscf.pbc.grad import krhf as krhf_grad from gpu4pyscf.lib.cupy_helper import contract -from gpu4pyscf.pbc.grad.pp import ( - vppnl_nuc_grad, _get_pp_nonloc_strain_derivatives) +from gpu4pyscf.pbc.grad.pp import vppnl_nuc_grad from gpu4pyscf.pbc.dft import multigrid from gpu4pyscf.pbc.gto import int1e diff --git a/gpu4pyscf/pbc/grad/pp.py b/gpu4pyscf/pbc/grad/pp.py index f600365b7..81792f5ca 100644 --- a/gpu4pyscf/pbc/grad/pp.py +++ b/gpu4pyscf/pbc/grad/pp.py @@ -13,158 +13,134 @@ # See the License for the specific language governing permissions and # limitations under the License. +import ctypes import numpy as np import cupy as cp -from pyscf import lib, gto -from pyscf.pbc.gto import pseudo -from pyscf.gto.mole import ATOM_OF from pyscf.pbc.lib.kpts_helper import gamma_point -from gpu4pyscf.gto.mole import groupby -from gpu4pyscf.lib import logger +from pyscf.pbc.tools.k2gamma import translation_vectors_for_kmesh from gpu4pyscf.lib.cupy_helper import contract, asarray -from gpu4pyscf.pbc.df import ft_ao -from gpu4pyscf.pbc.df.aft import get_SI +from gpu4pyscf.gto.mole import SortedGTO from gpu4pyscf.pbc.gto.pseudo.pp_int import _int_vnl_gpu, _sorted_fake_cell_vnl +from gpu4pyscf.pbc.gto import int1e +from gpu4pyscf.pbc.gto.int1e import libpbc +from gpu4pyscf.pbc.tools import k2gamma def vppnl_nuc_grad(cell, dm, kpts=None): '''Nuclear gradients of the non-local part of the GTH pseudo potential, contracted with the density matrix. + ''' + return vppnl_derivatives(cell, dm, kpts)[:-3] + +def vppnl_derivatives(cell, dm, kpts=None): + '''Nonlocal GTH atomic and strain derivatives, averaged over k-points. - Uses GPU CUDA kernels for the r^2/r^4 moment integrals at gamma point, - with CPU fallback via pyscf _int_vnl for multi-k-point calculations. + Returns a real (natm+3, 3) array: nuclear gradients followed by dE/d(strain_xy) ''' if kpts is None: - kpts_lst = np.zeros((1, 3)) + kpts = np.zeros((1, 3)) else: - kpts_lst = np.reshape(kpts, (-1, 3)) - nkpts = len(kpts_lst) - - # pattern stores the unique [hl_dim, l] combinations - fakecell, hl_blocks, pattern, splits = _sorted_fake_cell_vnl(cell) - - intors_d = ('int1e_ipovlp', 'int1e_r2_origi_ip2', 'int1e_r4_origi_ip2') - ppnl_half = _int_vnl_gpu(cell, fakecell, hl_blocks, kpts_lst) - ppnl_half_ip2 = _int_vnl_gpu(cell, fakecell, hl_blocks, kpts_lst, intors_d, comp=3) - if len(ppnl_half_ip2[0]) > 0: - ppnl_half_ip2[0] *= -1 + kpts = np.reshape(kpts, (-1, 3)) + is_gamma_point = gamma_point(kpts) + nkpts = len(kpts) nao = cell.nao dm = cp.asarray(dm).reshape(-1, nao, nao) - if gamma_point(kpts_lst): + if len(dm) != nkpts: + raise ValueError('Expected one density matrix per k-point') + if is_gamma_point: dm = dm.real - dm_dmH = dm + dm.transpose(0, 2, 1).conj() - - grad = np.zeros([cell.natm, 3], dtype=cp.complex128) - dppnl = cp.zeros((nao, 3), dtype=cp.complex128) - - hl_offset = [0] * 3 - for ii, (i0, i1) in enumerate(zip(splits[:-1], splits[1:])): - hl_dim, l = pattern[ii] - nd = 2 * l + 1 - hl_block = cp.asarray(np.stack(hl_blocks[i0:i1])) - n_hl = len(hl_block) + dm_dmH = dm + dm.conj().transpose(0, 2, 1) - ilp = cp.empty((nkpts, n_hl, hl_dim, nd, nao), dtype=cp.complex128) - dilp = cp.empty((nkpts, 3, n_hl, hl_dim, nd, nao), dtype=cp.complex128) - for i in range(hl_dim): - p0 = hl_offset[i] - p1 = p0 + n_hl * nd - ilp[:,:,i] = ppnl_half[i][:,p0:p1].reshape(nkpts, n_hl, nd, nao) - dilp[:,:,:,i] = ppnl_half_ip2[i][:,:,p0:p1].reshape(nkpts, 3, n_hl, nd, nao).conj() - hl_offset[i] = p1 - - tmp = contract('nij,knjlq->knilq', hl_block, ilp) - ilp = contract('knilq,kqp->knilp', tmp, dm_dmH, out=ilp) + grad_sigma = np.zeros((cell.natm+3, 3)) - value = contract('kdnilp,knilp->nd', dilp, ilp) - np.add.at(grad, fakecell._bas[i0:i1, ATOM_OF], value.get()) + fakecell, hl_blocks, pattern, splits = _sorted_fake_cell_vnl(cell) + if not hl_blocks: + return grad_sigma - dppnl += contract('kdnilp,knilp->pd', dilp, ilp) + sorted_cell = SortedGTO.from_cell(cell, decontract=True) + ppnl_half = _int_vnl_gpu(sorted_cell, fakecell, hl_blocks, kpts) - ao_loc = cell.ao_loc - atm_labels = np.repeat(cell._bas[:,ATOM_OF], ao_loc[1:]-ao_loc[:-1]) - grad -= groupby(atm_labels, dppnl.get(), 'sum') + derivative_kernels = ( + ('PBCovlp_cross_derivatives', (0, 1)), + ('PBCint1e_r2_origi_derivatives', (2, 1)), + ('PBCint1e_r4_origi_derivatives', (4, 1)), + ) - grad_max_imag = np.max(np.abs(grad.imag)) - if grad_max_imag >= 1e-8: - logger.warn(cell, f"Large imaginary part ({grad_max_imag:e}) from pseudopotential non-local term gradient.") - return grad.real + bvk_kmesh = k2gamma.kpts_to_kmesh(cell, kpts) + bvkmesh_Ls = translation_vectors_for_kmesh(cell, bvk_kmesh, True) + expLk = cp.exp(-1j*asarray(bvkmesh_Ls).dot(asarray(kpts).T)) + dtype = np.float64 if is_gamma_point else np.complex128 -def _get_pp_nonloc_strain_derivatives(cell, mesh, dm_kpts, kpts=None): - from gpu4pyscf.pbc.grad.rhf import _finite_diff_cells - if kpts is None: - assert dm_kpts.ndim == 2 - dm_kpts = dm_kpts[None,:,:] - kpts = np.zeros((1, 3)) - fakemol = gto.Mole() - fakemol._atm = np.zeros((1,gto.ATM_SLOTS), dtype=np.int32) - fakemol._bas = np.zeros((1,gto.BAS_SLOTS), dtype=np.int32) - ptr = gto.PTR_ENV_START - fakemol._env = np.zeros(ptr+10) - fakemol._bas[0,gto.NPRIM_OF ] = 1 - fakemol._bas[0,gto.NCTR_OF ] = 1 - fakemol._bas[0,gto.PTR_EXP ] = ptr+3 - fakemol._bas[0,gto.PTR_COEFF] = ptr+4 - - ngrids = np.prod(mesh) - buf = np.empty((48,ngrids), dtype=np.complex128) - scaled_kpts = kpts.dot(cell.lattice_vectors().T) - nkpts = len(kpts) - - def eval_pp_nonloc(cell): - vol = cell.vol - b = cell.reciprocal_vectors(norm_to=1) - Gv = cell.get_Gv(mesh) - SI = get_SI(cell, mesh=mesh) - # buf for SPG_lmi upto l=0..3 and nl=3 - vppnl = 0 - for k, dm in enumerate(dm_kpts): - kpt = scaled_kpts[k].dot(b) - Gk = Gv + kpt - G_rad = lib.norm(Gk, axis=1) - aokG = ft_ao.ft_ao(cell, Gv, kpt=kpt) * (1/vol)**.5 - for ia in range(cell.natm): - symb = cell.atom_symbol(ia) - if symb not in cell._pseudo: - continue - pp = cell._pseudo[symb] - p1 = 0 - for l, proj in enumerate(pp[5:]): - rl, nl, hl = proj - if nl > 0: - fakemol._bas[0,gto.ANG_OF] = l - fakemol._env[ptr+3] = .5*rl**2 - fakemol._env[ptr+4] = rl**(l+1.5)*np.pi**1.25 - pYlm_part = fakemol.eval_gto('GTOval', Gk) - - p0, p1 = p1, p1+nl*(l*2+1) - # pYlm is real, SI[ia] is complex - pYlm = np.ndarray((nl,l*2+1,ngrids), dtype=np.complex128, buffer=buf[p0:p1]) - for k in range(nl): - qkl = pseudo.pp._qli(G_rad*rl, l, k) - pYlm[k] = pYlm_part.T * qkl - if p1 > 0: - SPG_lmi = asarray(buf[:p1]) - SPG_lmi *= SI[ia].conj() - SPG_lm_aoGs = SPG_lmi.dot(aokG) - rho = SPG_lm_aoGs.dot(dm).dot(SPG_lm_aoGs.conj().T).real.get() - p1 = 0 - for l, proj in enumerate(pp[5:]): - rl, nl, hl = proj - if nl > 0: - nf = l * 2 + 1 - p0, p1 = p1, p1+nl*nf - hl = np.asarray(hl) - rho_sub = rho[p0:p1,p0:p1].reshape(nl, nf, nl, nf) - vppnl += np.einsum('ij,jmim->', hl, rho_sub) - return vppnl / (nkpts*vol) - - disp = max(1e-5, (cell.precision*.1)**.5) - out = np.empty((3, 3)) - for i in range(3): - for j in range(3): - cell1, cell2 = _finite_diff_cells(cell, i, j, disp) - e1 = eval_pp_nonloc(cell1) - e2 = eval_pp_nonloc(cell2) - out[i,j] = (e1 - e2) / (2*disp) - return out + hl_offset = [0] * 3 + for (hl_dim, l), i0, i1 in zip(pattern, splits[:-1], splits[1:]): + if hl_dim == 0: + continue + nd = 2 * l + 1 + n_hl = i1 - i0 + hl_block = asarray(np.stack(hl_blocks[i0:i1])) + ilp = cp.empty((hl_dim, nkpts, n_hl, nd, nao), dtype=dtype) + for rn in range(hl_dim): + p0 = hl_offset[rn] + p1 = p0 + n_hl * nd + ilp[rn] = ppnl_half[rn][:,p0:p1].reshape(nkpts, n_hl, nd, nao) + hl_offset[rn] = p1 + tmp = contract('nij,jknlq->iknlq', hl_block, ilp) + weights = contract('iknlq,kqp->iknlp', tmp, dm_dmH, out=ilp) + + if not is_gamma_point: + weights = contract('Lk,iknlp->iLnlp', expLk, weights).real + + pcell = fakecell.copy(deep=False) + pcell._bas = fakecell._bas[i0:i1] + opt = int1e.CrossInt1e(pcell, sorted_cell, bvk_kmesh) + for rn in range(hl_dim): + kern, deriv = derivative_kernels[rn] + as_dm = weights[rn].reshape(nkpts, n_hl*nd, nao) + grad_sigma += _derivatives_intor(opt, as_dm, kern, deriv) + return grad_sigma + +def _derivatives_intor(cross_int1e, dm, kern, deriv): + '''Contract Re[dm * conjugate(dI)] for a rectangular cross integral. + ''' + cell = cross_int1e.cell + assert dm.ndim == 3 + assert dm.dtype == np.float64 + nkpts = len(dm) + + tmp = cross_int1e.cell2.apply_C_dot(dm, axis=2) + dm = cross_int1e.cell1.apply_C_dot(tmp, axis=1) + dm = cp.asarray(dm, order='C') + + gout_stride_lookup, shm_size = int1e._gout_stride_lookup_table(cell, deriv) + nbatches_shl_pair = len(cross_int1e.shl_pair_offsets) - 1 + + if nbatches_shl_pair == 0: + return np.zeros([cell.natm+3, 3]) + + grad = cp.zeros((cell.natm, 3)) + sigma = cp.zeros((3, 3)) + drv = getattr(libpbc, kern) + err = drv( + ctypes.cast(grad.data.ptr, ctypes.c_void_p), + ctypes.cast(sigma.data.ptr, ctypes.c_void_p), + ctypes.cast(dm.data.ptr, ctypes.c_void_p), + ctypes.byref(cross_int1e.int1e_envs), + ctypes.c_int(shm_size), + ctypes.c_int(nbatches_shl_pair), + ctypes.cast(cross_int1e.shl_pair_offsets.data.ptr, ctypes.c_void_p), + ctypes.cast(cross_int1e.bas_ij_idx.data.ptr, ctypes.c_void_p), + ctypes.cast(gout_stride_lookup.data.ptr, ctypes.c_void_p), + ctypes.c_int(cross_int1e.cell1.nao), + ctypes.c_int(cross_int1e.cell2.nao)) + if err != 0: + raise RuntimeError(f'{kern} failed') + + # CrossInt1e concatenates the projector and AO atom lists. + natm = cross_int1e.cell1.natm + if natm != cross_int1e.cell2.natm: + raise ValueError( + 'fakecell for ppnl must have the same number of atoms as the AO cell') + grad = (grad[:natm] + grad[natm:]).get() + grad_sigma = np.vstack([grad, sigma.get()]) + grad_sigma /= nkpts + return grad_sigma diff --git a/gpu4pyscf/pbc/grad/rhf.py b/gpu4pyscf/pbc/grad/rhf.py index a2424ab5f..d6c760365 100644 --- a/gpu4pyscf/pbc/grad/rhf.py +++ b/gpu4pyscf/pbc/grad/rhf.py @@ -30,8 +30,7 @@ from gpu4pyscf.pbc.df import aft_jk, AFTDF, GDF from gpu4pyscf.pbc.gto import int1e from gpu4pyscf.pbc.dft import KohnShamDFT, BeckeGrids -from gpu4pyscf.pbc.grad.pp import ( - vppnl_nuc_grad, _get_pp_nonloc_strain_derivatives) +from gpu4pyscf.pbc.grad.pp import vppnl_nuc_grad from gpu4pyscf.gto.mole import groupby __all__ = ['Gradients'] diff --git a/gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py b/gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py index 10d63e544..7f8b07750 100644 --- a/gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py +++ b/gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py @@ -130,6 +130,69 @@ def _cpu_vppnl_nuc_grad(cell, dm, kpts=None): return grad.real +def vppnl_nuc_grad(cell, dm, kpts=None): + '''Nuclear gradients of the non-local part of the GTH pseudo potential, + contracted with the density matrix. + ''' + from gpu4pyscf.lib import logger + from gpu4pyscf.gto.mole import groupby + if kpts is None: + kpts_lst = np.zeros((1, 3)) + else: + kpts_lst = np.reshape(kpts, (-1, 3)) + nkpts = len(kpts_lst) + + # pattern stores the unique [hl_dim, l] combinations + fakecell, hl_blocks, pattern, splits = _sorted_fake_cell_vnl(cell) + + intors_d = ('int1e_ipovlp', 'int1e_r2_origi_ip2', 'int1e_r4_origi_ip2') + ppnl_half = _int_vnl_gpu(cell, fakecell, hl_blocks, kpts_lst) + ppnl_half_ip2 = _int_vnl_gpu(cell, fakecell, hl_blocks, kpts_lst, intors_d, comp=3) + if len(ppnl_half_ip2[0]) > 0: + ppnl_half_ip2[0] *= -1 + + nao = cell.nao + dm = cp.asarray(dm).reshape(-1, nao, nao) + if gamma_point(kpts_lst): + dm = dm.real + dm_dmH = dm + dm.transpose(0, 2, 1).conj() + + grad = np.zeros([cell.natm, 3], dtype=cp.complex128) + dppnl = cp.zeros((nao, 3), dtype=cp.complex128) + + hl_offset = [0] * 3 + for ii, (i0, i1) in enumerate(zip(splits[:-1], splits[1:])): + hl_dim, l = pattern[ii] + nd = 2 * l + 1 + hl_block = cp.asarray(np.stack(hl_blocks[i0:i1])) + n_hl = len(hl_block) + + ilp = cp.empty((hl_dim, nkpts, n_hl, nd, nao), dtype=cp.complex128) + dilp = cp.empty((hl_dim, nkpts, 3, n_hl, nd, nao), dtype=cp.complex128) + for i in range(hl_dim): + p0 = hl_offset[i] + p1 = p0 + n_hl * nd + ilp[i] = ppnl_half[i][:,p0:p1].reshape(nkpts, n_hl, nd, nao) + dilp[i] = ppnl_half_ip2[i][:,:,p0:p1].reshape(nkpts, 3, n_hl, nd, nao).conj() + hl_offset[i] = p1 + + tmp = contract('nij,jknlq->iknlq', hl_block, ilp) + ilp = contract('iknlq,kqp->iknlp', tmp, dm_dmH, out=ilp) + + value = contract('ikdnlp,iknlp->nd', dilp, ilp) + np.add.at(grad, fakecell._bas[i0:i1, ATOM_OF], value.get()) + + dppnl += contract('ikdnlp,iknlp->pd', dilp, ilp) + + ao_loc = cell.ao_loc + atm_labels = np.repeat(cell._bas[:,ATOM_OF], ao_loc[1:]-ao_loc[:-1]) + grad -= groupby(atm_labels, dppnl.get(), 'sum') + + grad_max_imag = np.max(np.abs(grad.imag)) + if grad_max_imag >= 1e-8: + logger.warn(cell, f"Large imaginary part ({grad_max_imag:e}) from pseudopotential non-local term gradient.") + return grad.real + class TestCrossBasisIntegrals(unittest.TestCase): """Test GPU _int_vnl_gpu against CPU _int_vnl for each element.""" @@ -310,6 +373,43 @@ def _fd_check(self, cell, atom_id=1, cart_id=0, places=5): def test_carbon_fd(self): self._fd_check(cell_c, atom_id=1, cart_id=0, places=5) + def _strain_fd_check(self, cell, scaled_kpts): + from pyscf.pbc.gto.pseudo.pp_int import get_pp_nl + from gpu4pyscf.pbc.grad.pp import _get_pp_nonloc_strain_derivatives + from gpu4pyscf.pbc.grad.rhf import _finite_diff_cells + + kpts = cell.get_abs_kpts(np.asarray(scaled_kpts)) + rng = np.random.default_rng(19) + dm = rng.normal(size=(len(kpts), cell.nao, cell.nao)) + if not gamma_point(kpts): + dm = dm + 1j * rng.normal(size=dm.shape) + dm = (dm + dm.transpose(0, 2, 1).conj()) / cell.nao + analytical = _get_pp_nonloc_strain_derivatives(cell, cell.mesh, dm, kpts) + reference = np.empty((3, 3)) + # Keep fractional k-points and AO density fixed. The CPU real-space + # PP energy independently checks all nine strain components, including + # the image displacement and k-point averaging in the GPU derivative. + step = 1e-5 + for x in range(3): + for y in range(3): + energies = [] + for strained in _finite_diff_cells(cell, x, y, disp=step): + vpp = get_pp_nl(strained, strained.get_abs_kpts(scaled_kpts)) + energies.append(np.einsum('kij,kji->', vpp, dm).real / len(kpts)) + reference[x, y] = (energies[0] - energies[1]) / (2 * step) + np.testing.assert_allclose(analytical, reference, atol=2e-6, rtol=1e-5) + + def test_ppnl_strain_gamma(self): + self._strain_fd_check(cell_c, [[0., 0., 0.]]) + + def test_ppnl_strain_kpts(self): + self._strain_fd_check(cell_si, [[.13, .07, -.11], [-.21, .16, .09]]) + + @pytest.mark.slow + def test_ppnl_strain_iron(self): + # Fe exercises higher angular momentum and r^2/r^4 projectors. + self._strain_fd_check(cell_fe, [[.13, .07, -.11]]) + def test_silicon_fd(self): self._fd_check(cell_si, atom_id=0, cart_id=2, places=5) diff --git a/gpu4pyscf/pbc/gto/int1e.py b/gpu4pyscf/pbc/gto/int1e.py index 21b877f51..a1d66e971 100644 --- a/gpu4pyscf/pbc/gto/int1e.py +++ b/gpu4pyscf/pbc/gto/int1e.py @@ -85,19 +85,19 @@ def int1e_ipkin(cell, kpts=None, bvk_kmesh=None, sort_output=True): def int1e_r2_origi(cell, kpts=None, bvk_kmesh=None, sort_output=True): opt = _check_opt(cell, 0, kpts, bvk_kmesh) - return opt.intor('PBCint1e_r2_origi', 1, (0, 2), kpts, sort_output) + return opt.intor('PBCint1e_r2_origi', 1, (2, 0), kpts, sort_output) def int1e_r4_origi(cell, kpts=None, bvk_kmesh=None, sort_output=True): opt = _check_opt(cell, 0, kpts, bvk_kmesh) - return opt.intor('PBCint1e_r4_origi', 1, (0, 4), kpts, sort_output) + return opt.intor('PBCint1e_r4_origi', 1, (4, 0), kpts, sort_output) def int1e_r2_origi_ip2(cell, kpts=None, bvk_kmesh=None, sort_output=True): opt = _check_opt(cell, 0, kpts, bvk_kmesh) - return opt.intor('PBCint1e_r2_origi_ip2', 3, (0, 3), kpts, sort_output) + return opt.intor('PBCint1e_r2_origi_ip2', 3, (2, 1), kpts, sort_output) def int1e_r4_origi_ip2(cell, kpts=None, bvk_kmesh=None, sort_output=True): opt = _check_opt(cell, 0, kpts, bvk_kmesh) - return opt.intor('PBCint1e_r4_origi_ip2', 3, (0, 5), kpts, sort_output) + return opt.intor('PBCint1e_r4_origi_ip2', 3, (4, 1), kpts, sort_output) def ovlp_derivatives(cell, dm, kpts=None, kmesh=None): assert isinstance(cell, Cell) @@ -429,8 +429,7 @@ def __init__(self, cell1, cell2, bvk_kmesh=None): def intor(self, kern, comp, deriv_ij, kpts=None, sort_output=True, out=None, buf=None, shls_slice=None): shls_slice = (0, self.cell1.nbas, self.cell1.nbas, self.cell.nbas) - out = super().intor(kern, comp, deriv_ij, kpts, False, out, buf, - shls_slice) + out = super().intor(kern, comp, deriv_ij, kpts, False, out, buf, shls_slice) if sort_output: leading_shape = out.shape[:-2] n1, n2 = out.shape[-2:] diff --git a/gpu4pyscf/pbc/gto/pseudo/pp_int.py b/gpu4pyscf/pbc/gto/pseudo/pp_int.py index 366e116e0..cd2a0c6ce 100644 --- a/gpu4pyscf/pbc/gto/pseudo/pp_int.py +++ b/gpu4pyscf/pbc/gto/pseudo/pp_int.py @@ -17,27 +17,26 @@ import cupy as cp from pyscf import gto, lib from pyscf.pbc.gto.cell import _estimate_rcut -from pyscf.pbc.gto.pseudo.pp_int import fake_cell_vnl, _int_vnl +from pyscf.pbc.gto.pseudo.pp_int import fake_cell_vnl from pyscf.pbc.lib.kpts_helper import gamma_point from gpu4pyscf.lib.cupy_helper import contract from gpu4pyscf.gto.mole import most_diffuse_pgto, SortedGTO from gpu4pyscf.pbc.gto import int1e from gpu4pyscf.pbc.tools import k2gamma +kern_map = { + 'int1e_ovlp': ('PBCint1e_ovlp', 1, (0, 0)), + 'int1e_r2_origi': ('PBCint1e_r2_origi', 1, (2, 0)), + 'int1e_r4_origi': ('PBCint1e_r4_origi', 1, (4, 0)), + 'int1e_ipovlp': ('PBCint1e_ipovlp', 3, (1, 0)), + 'int1e_r2_origi_ip2': ('PBCint1e_r2_origi_ip2', 3, (2, 1)), + 'int1e_r4_origi_ip2': ('PBCint1e_r4_origi_ip2', 3, (4, 1)), +} def _int_vnl_gpu(cell, fakecell, hl_blocks, kpts, intors=None, comp=1): if intors is None: intors = ['int1e_ovlp', 'int1e_r2_origi', 'int1e_r4_origi'] - kern_map = { - 'int1e_ovlp': ('PBCint1e_ovlp', 1, (0, 0)), - 'int1e_r2_origi': ('PBCint1e_r2_origi', 1, (0, 2)), - 'int1e_r4_origi': ('PBCint1e_r4_origi', 1, (0, 4)), - 'int1e_ipovlp': ('PBCint1e_ipovlp', 3, (1, 0)), - 'int1e_r2_origi_ip2': ('PBCint1e_r2_origi_ip2', 3, (0, 3)), - 'int1e_r4_origi_ip2': ('PBCint1e_r4_origi_ip2', 3, (0, 5)), - } - hl_dims = np.asarray([len(hl) for hl in hl_blocks]) cell = SortedGTO.from_cell(cell, decontract=True) @@ -59,6 +58,9 @@ def int_ket(_bas_fake, intor_name): def _sorted_fake_cell_vnl(cell): fakecell, hl_blocks = fake_cell_vnl(cell) + # GTH projectors are spherical even when the AO basis is Cartesian. + fakecell.cart = False + hl_dims = np.asarray([len(hl) for hl in hl_blocks]) ls = fakecell._bas[:,gto.ANG_OF] # groupby [hl_dim, l] @@ -83,8 +85,9 @@ def get_pp_nl_gpu(cell, kpts=None): ppnl_half = _int_vnl_gpu(cell, fakecell, hl_blocks, kpts_lst) + dtype = np.float64 if is_gamma_point else np.complex128 nao = cell.nao - ppnl = cp.zeros((nkpts, nao, nao), dtype=cp.complex128) + ppnl = cp.zeros((nkpts, nao, nao), dtype=dtype) hl_offset = [0] * 3 for ii, (i0, i1) in enumerate(zip(splits[:-1], splits[1:])): @@ -93,17 +96,14 @@ def get_pp_nl_gpu(cell, kpts=None): hl_block = cp.asarray(np.stack(hl_blocks[i0:i1])) n_hl = len(hl_block) - ilp = cp.empty((nkpts, n_hl, hl_dim, nd, nao), dtype=cp.complex128) + ilp = cp.empty((hl_dim, nkpts, n_hl, nd, nao), dtype=dtype) for i in range(hl_dim): p0 = hl_offset[i] p1 = p0 + n_hl * nd - ilp[:,:,i] = ppnl_half[i][:,p0:p1].reshape(nkpts, n_hl, nd, nao) + ilp[i] = ppnl_half[i][:,p0:p1].reshape(nkpts, n_hl, nd, nao) hl_offset[i] = p1 - tmp = contract('nij,knjlq->knilq', hl_block, ilp) + tmp = contract('nij,jknlq->iknlq', hl_block, ilp) ilp_conj = cp.conjugate(ilp, out=ilp) - contract('knilp,knilq->kpq', ilp_conj, tmp, beta=1, out=ppnl) - - if kpts is None or gamma_point(kpts): - ppnl = ppnl.real + contract('iknlp,iknlq->kpq', ilp_conj, tmp, beta=1, out=ppnl) return ppnl From c45d60c1a45e6c84e766272b0d885e482b60ecb9 Mon Sep 17 00:00:00 2001 From: Qiming Sun Date: Tue, 15 Sep 2026 12:16:07 -0700 Subject: [PATCH 2/7] Add tests --- gpu4pyscf/pbc/grad/pp.py | 4 +- gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py | 186 +++++++++++++++---- gpu4pyscf/pbc/gto/pseudo/pp_int.py | 1 + 3 files changed, 152 insertions(+), 39 deletions(-) diff --git a/gpu4pyscf/pbc/grad/pp.py b/gpu4pyscf/pbc/grad/pp.py index 81792f5ca..a4bf720c9 100644 --- a/gpu4pyscf/pbc/grad/pp.py +++ b/gpu4pyscf/pbc/grad/pp.py @@ -29,9 +29,9 @@ def vppnl_nuc_grad(cell, dm, kpts=None): '''Nuclear gradients of the non-local part of the GTH pseudo potential, contracted with the density matrix. ''' - return vppnl_derivatives(cell, dm, kpts)[:-3] + return ppnl_derivatives(cell, dm, kpts)[:-3] -def vppnl_derivatives(cell, dm, kpts=None): +def ppnl_derivatives(cell, dm, kpts=None): '''Nonlocal GTH atomic and strain derivatives, averaged over k-points. Returns a real (natm+3, 3) array: nuclear gradients followed by dE/d(strain_xy) diff --git a/gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py b/gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py index 7f8b07750..62538f11a 100644 --- a/gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py +++ b/gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py @@ -35,6 +35,9 @@ from pyscf.pbc.lib.kpts_helper import gamma_point import gpu4pyscf.pbc.dft.multigrid as multigrid_v1 from gpu4pyscf.pbc.dft import multigrid_v3 +from gpu4pyscf.pbc.gto.pseudo import pp_int +from gpu4pyscf.pbc.lib.kpts_helper import fft_matrix +from gpu4pyscf.pbc.grad.pp import ppnl_derivatives import pytest disp = 1e-4 @@ -134,8 +137,12 @@ def vppnl_nuc_grad(cell, dm, kpts=None): '''Nuclear gradients of the non-local part of the GTH pseudo potential, contracted with the density matrix. ''' + from pyscf.gto import ATOM_OF + from pyscf.pbc.lib.kpts_helper import gamma_point + from gpu4pyscf.lib.cupy_helper import contract from gpu4pyscf.lib import logger from gpu4pyscf.gto.mole import groupby + from gpu4pyscf.pbc.gto.pseudo.pp_int import _int_vnl_gpu, _sorted_fake_cell_vnl if kpts is None: kpts_lst = np.zeros((1, 3)) else: @@ -193,6 +200,85 @@ def vppnl_nuc_grad(cell, dm, kpts=None): logger.warn(cell, f"Large imaginary part ({grad_max_imag:e}) from pseudopotential non-local term gradient.") return grad.real +def _get_pp_nonloc_strain_derivatives(cell, mesh, dm_kpts, kpts=None): + from gpu4pyscf.pbc.grad.rhf import _finite_diff_cells + if kpts is None: + assert dm_kpts.ndim == 2 + dm_kpts = dm_kpts[None,:,:] + kpts = np.zeros((1, 3)) + fakemol = gto.Mole() + fakemol._atm = np.zeros((1,gto.ATM_SLOTS), dtype=np.int32) + fakemol._bas = np.zeros((1,gto.BAS_SLOTS), dtype=np.int32) + ptr = gto.PTR_ENV_START + fakemol._env = np.zeros(ptr+10) + fakemol._bas[0,gto.NPRIM_OF ] = 1 + fakemol._bas[0,gto.NCTR_OF ] = 1 + fakemol._bas[0,gto.PTR_EXP ] = ptr+3 + fakemol._bas[0,gto.PTR_COEFF] = ptr+4 + + ngrids = np.prod(mesh) + buf = np.empty((48,ngrids), dtype=np.complex128) + scaled_kpts = kpts.dot(cell.lattice_vectors().T) + nkpts = len(kpts) + + def eval_pp_nonloc(cell): + vol = cell.vol + b = cell.reciprocal_vectors(norm_to=1) + Gv = cell.get_Gv(mesh) + SI = get_SI(cell, mesh=mesh) + # buf for SPG_lmi upto l=0..3 and nl=3 + vppnl = 0 + for k, dm in enumerate(dm_kpts): + kpt = scaled_kpts[k].dot(b) + Gk = Gv + kpt + G_rad = lib.norm(Gk, axis=1) + aokG = ft_ao.ft_ao(cell, Gv, kpt=kpt) * (1/vol)**.5 + for ia in range(cell.natm): + symb = cell.atom_symbol(ia) + if symb not in cell._pseudo: + continue + pp = cell._pseudo[symb] + p1 = 0 + for l, proj in enumerate(pp[5:]): + rl, nl, hl = proj + if nl > 0: + fakemol._bas[0,gto.ANG_OF] = l + fakemol._env[ptr+3] = .5*rl**2 + fakemol._env[ptr+4] = rl**(l+1.5)*np.pi**1.25 + pYlm_part = fakemol.eval_gto('GTOval', Gk) + + p0, p1 = p1, p1+nl*(l*2+1) + # pYlm is real, SI[ia] is complex + pYlm = np.ndarray((nl,l*2+1,ngrids), dtype=np.complex128, buffer=buf[p0:p1]) + for k in range(nl): + qkl = pseudo.pp._qli(G_rad*rl, l, k) + pYlm[k] = pYlm_part.T * qkl + if p1 > 0: + SPG_lmi = asarray(buf[:p1]) + SPG_lmi *= SI[ia].conj() + SPG_lm_aoGs = SPG_lmi.dot(aokG) + rho = SPG_lm_aoGs.dot(dm).dot(SPG_lm_aoGs.conj().T).real.get() + p1 = 0 + for l, proj in enumerate(pp[5:]): + rl, nl, hl = proj + if nl > 0: + nf = l * 2 + 1 + p0, p1 = p1, p1+nl*nf + hl = np.asarray(hl) + rho_sub = rho[p0:p1,p0:p1].reshape(nl, nf, nl, nf) + vppnl += np.einsum('ij,jmim->', hl, rho_sub) + return vppnl / (nkpts*vol) + + disp = max(1e-5, (cell.precision*.1)**.5) + out = np.empty((3, 3)) + for i in range(3): + for j in range(3): + cell1, cell2 = _finite_diff_cells(cell, i, j, disp) + e1 = eval_pp_nonloc(cell1) + e2 = eval_pp_nonloc(cell2) + out[i,j] = (e1 - e2) / (2*disp) + return out + class TestCrossBasisIntegrals(unittest.TestCase): """Test GPU _int_vnl_gpu against CPU _int_vnl for each element.""" @@ -373,43 +459,6 @@ def _fd_check(self, cell, atom_id=1, cart_id=0, places=5): def test_carbon_fd(self): self._fd_check(cell_c, atom_id=1, cart_id=0, places=5) - def _strain_fd_check(self, cell, scaled_kpts): - from pyscf.pbc.gto.pseudo.pp_int import get_pp_nl - from gpu4pyscf.pbc.grad.pp import _get_pp_nonloc_strain_derivatives - from gpu4pyscf.pbc.grad.rhf import _finite_diff_cells - - kpts = cell.get_abs_kpts(np.asarray(scaled_kpts)) - rng = np.random.default_rng(19) - dm = rng.normal(size=(len(kpts), cell.nao, cell.nao)) - if not gamma_point(kpts): - dm = dm + 1j * rng.normal(size=dm.shape) - dm = (dm + dm.transpose(0, 2, 1).conj()) / cell.nao - analytical = _get_pp_nonloc_strain_derivatives(cell, cell.mesh, dm, kpts) - reference = np.empty((3, 3)) - # Keep fractional k-points and AO density fixed. The CPU real-space - # PP energy independently checks all nine strain components, including - # the image displacement and k-point averaging in the GPU derivative. - step = 1e-5 - for x in range(3): - for y in range(3): - energies = [] - for strained in _finite_diff_cells(cell, x, y, disp=step): - vpp = get_pp_nl(strained, strained.get_abs_kpts(scaled_kpts)) - energies.append(np.einsum('kij,kji->', vpp, dm).real / len(kpts)) - reference[x, y] = (energies[0] - energies[1]) / (2 * step) - np.testing.assert_allclose(analytical, reference, atol=2e-6, rtol=1e-5) - - def test_ppnl_strain_gamma(self): - self._strain_fd_check(cell_c, [[0., 0., 0.]]) - - def test_ppnl_strain_kpts(self): - self._strain_fd_check(cell_si, [[.13, .07, -.11], [-.21, .16, .09]]) - - @pytest.mark.slow - def test_ppnl_strain_iron(self): - # Fe exercises higher angular momentum and r^2/r^4 projectors. - self._strain_fd_check(cell_fe, [[.13, .07, -.11]]) - def test_silicon_fd(self): self._fd_check(cell_si, atom_id=0, cart_id=2, places=5) @@ -520,6 +569,69 @@ def get_pp_local_energy(cell): assert np.max(np.abs(numerical_gradient - analytical_gradient)) < 1e-8 + def test_ppnl_derivatives(self): + from gpu4pyscf.pbc.grad.rhf import _finite_diff_cells + cell = pyscf.M( + a = np.array([ + [3.18693029, 0.0, 0.0], + [1.593466157846262, 2.759963819342879, 0.0], + [1.5934664345206309, 0.9199872811273334, 2.6021185638285855], + ]), + atom = """ + Ga 0 -0 0 + Ga 0.25 0.75 0.25 + N 0.27 0.25 0.25 + N 0.75 0.25 0.75 + """, + unit = "Angstrom", + fractional = True, + basis=[[0, [1.2, .7, .2], [.5, .3, .8]], + [1, [.8, 1]], [2, [.7, 1]]], + pseudo = { + "Ga": """ + Ga GTH-PBE-q13 GTH-GGA-q13 + 2 1 10 0 + 0.49000018487159 0 + 3 + 0.41677483095310 3 10.48679119269639 -4.92176814704009 0.87070493953275 + 7.77018207637078 -2.24815160599927 + 1.78441528219626 + 0.56962661099353 2 1.77860037827899 0.19586036552562 + -0.23168154587648 + 0.23814730101676 1 -16.24818353736915 + """, + "N": """ + N GTH-PBE-q5 GTH-GGA-q5 + 2 3 0 0 + 0.28382600053810 2 -12.41517350030142 1.86813618209744 + 1 + 0.25541754972811 1 13.63124869974610 + """}, + ) + nao = cell.nao + kmesh = [3,2,1] + kpts = cell.make_kpts(kmesh) + nkpts = len(kpts) + cp.random.seed(11) + dm = cp.random.rand(nkpts, nao, nao) * .2 + expLk = fft_matrix(kmesh) + dm = cp.einsum('Lk,Lpq->kpq', expLk.conj(), dm) + dm = dm + dm.conj().transpose(0,2,1) + ref = vppnl_nuc_grad(cell, dm, kpts) / nkpts + dat = ppnl_derivatives(cell, dm, kpts) + assert abs(ref - dat[:-3]).max() < 1e-10 + + sigma = dat[-3:] + + disp = 1e-4 + for (i, j) in [(0, 0), (0, 1), (0, 2), (2, 0), (2, 2)]: + cell1, cell2 = _finite_diff_cells(cell, i, j, disp=disp) + v = pp_int.get_pp_nl_gpu(cell1, cell1.make_kpts(kmesh)) + e1 = cp.einsum('kpq,kqp->', dm, v).real / nkpts + v = pp_int.get_pp_nl_gpu(cell2, cell2.make_kpts(kmesh)) + e2 = cp.einsum('kpq,kqp->', dm, v).real / nkpts + assert abs(sigma[i, j] - (e1-e2)/2/disp) < 2e-7 + if __name__ == '__main__': unittest.main() diff --git a/gpu4pyscf/pbc/gto/pseudo/pp_int.py b/gpu4pyscf/pbc/gto/pseudo/pp_int.py index cd2a0c6ce..6daaabaf9 100644 --- a/gpu4pyscf/pbc/gto/pseudo/pp_int.py +++ b/gpu4pyscf/pbc/gto/pseudo/pp_int.py @@ -85,6 +85,7 @@ def get_pp_nl_gpu(cell, kpts=None): ppnl_half = _int_vnl_gpu(cell, fakecell, hl_blocks, kpts_lst) + is_gamma_point = gamma_point(kpts) dtype = np.float64 if is_gamma_point else np.complex128 nao = cell.nao ppnl = cp.zeros((nkpts, nao, nao), dtype=dtype) From 08633ec6c15e3790d00dad7341718746691b5aae Mon Sep 17 00:00:00 2001 From: Qiming Sun Date: Wed, 16 Sep 2026 16:42:24 -0700 Subject: [PATCH 3/7] Update relevant code in pbc.grad module --- gpu4pyscf/pbc/grad/krhf.py | 6 ++---- gpu4pyscf/pbc/grad/kuhf.py | 8 -------- gpu4pyscf/pbc/grad/rhf.py | 5 ++--- gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py | 5 +++++ 4 files changed, 9 insertions(+), 15 deletions(-) diff --git a/gpu4pyscf/pbc/grad/krhf.py b/gpu4pyscf/pbc/grad/krhf.py index 5953b9519..60642102f 100644 --- a/gpu4pyscf/pbc/grad/krhf.py +++ b/gpu4pyscf/pbc/grad/krhf.py @@ -34,7 +34,7 @@ from gpu4pyscf.pbc.gto import int1e from gpu4pyscf.pbc.scf.rsjk import PBCJKMatrixOpt from gpu4pyscf.pbc import tools as pbctools -from gpu4pyscf.pbc.grad.pp import vppnl_nuc_grad +from gpu4pyscf.pbc.grad.pp import ppnl_derivatives from gpu4pyscf.pbc.grad.rhf import contract_h1e_dm, _get_ejk_derivatives from gpu4pyscf.pbc.grad import rhf as pbchf_grad @@ -370,9 +370,7 @@ def grad_elec(self, mo_energy=None, mo_coeff=None, mo_occ=None): 'HF', dm0, kpts, spin=0, with_j=False, with_nuc=True) if cell._pseudo: - grad_sigma[:-3] += vppnl_nuc_grad(cell, dm0, kpts=kpts) / nkpts - grad_sigma[-3:] += _get_pp_nonloc_strain_derivatives( - cell, cell.mesh, dm0, kpts=kpts) + grad_sigma += ppnl_derivatives(cell, dm0, kpts) log.timer_debug1('gradients of 1e part', *t1) diff --git a/gpu4pyscf/pbc/grad/kuhf.py b/gpu4pyscf/pbc/grad/kuhf.py index a94ef30cb..c5a858751 100644 --- a/gpu4pyscf/pbc/grad/kuhf.py +++ b/gpu4pyscf/pbc/grad/kuhf.py @@ -17,15 +17,7 @@ Analytical nuclear gradients for RHF with kpoints sampling ''' -import numpy as np -import cupy as cp -from pyscf import lib -from gpu4pyscf.lib import logger from gpu4pyscf.pbc.grad import krhf as krhf_grad -from gpu4pyscf.lib.cupy_helper import contract -from gpu4pyscf.pbc.grad.pp import vppnl_nuc_grad -from gpu4pyscf.pbc.dft import multigrid -from gpu4pyscf.pbc.gto import int1e __all__ = ['Gradients'] diff --git a/gpu4pyscf/pbc/grad/rhf.py b/gpu4pyscf/pbc/grad/rhf.py index d6c760365..94c6fd8eb 100644 --- a/gpu4pyscf/pbc/grad/rhf.py +++ b/gpu4pyscf/pbc/grad/rhf.py @@ -30,7 +30,7 @@ from gpu4pyscf.pbc.df import aft_jk, AFTDF, GDF from gpu4pyscf.pbc.gto import int1e from gpu4pyscf.pbc.dft import KohnShamDFT, BeckeGrids -from gpu4pyscf.pbc.grad.pp import vppnl_nuc_grad +from gpu4pyscf.pbc.grad.pp import ppnl_derivatives from gpu4pyscf.gto.mole import groupby __all__ = ['Gradients'] @@ -201,8 +201,7 @@ def grad_elec(self, mo_energy=None, mo_coeff=None, mo_occ=None, atmlst=None): 'HF', dm0, spin=0, with_j=False, with_nuc=True) if cell._pseudo: - grad_sigma[:-3] += vppnl_nuc_grad(cell, dm0) - grad_sigma[-3:] += _get_pp_nonloc_strain_derivatives(cell, cell.mesh, dm0) + grad_sigma += ppnl_derivatives(cell, dm0) t1 = log.timer_debug1('gradients of 1e part', *t1) dme0 = self.make_rdm1e(mo_energy, mo_coeff, mo_occ) diff --git a/gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py b/gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py index 62538f11a..d65112793 100644 --- a/gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py +++ b/gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py @@ -201,7 +201,12 @@ def vppnl_nuc_grad(cell, dm, kpts=None): return grad.real def _get_pp_nonloc_strain_derivatives(cell, mesh, dm_kpts, kpts=None): + from pyscf import lib, gto + from pyscf.pbc.gto import pseudo from gpu4pyscf.pbc.grad.rhf import _finite_diff_cells + from gpu4pyscf.lib.cupy_helper import asarray + from gpu4pyscf.pbc.df import ft_ao + from gpu4pyscf.pbc.df.aft import get_SI if kpts is None: assert dm_kpts.ndim == 2 dm_kpts = dm_kpts[None,:,:] From 30603c44e298320c050c92e876c3d1e6b9e52156 Mon Sep 17 00:00:00 2001 From: Qiming Sun Date: Wed, 16 Sep 2026 21:19:38 -0700 Subject: [PATCH 4/7] Missing file --- gpu4pyscf/lib/pbc/overlap.cuh | 67 +++++++++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) create mode 100644 gpu4pyscf/lib/pbc/overlap.cuh diff --git a/gpu4pyscf/lib/pbc/overlap.cuh b/gpu4pyscf/lib/pbc/overlap.cuh new file mode 100644 index 000000000..710091326 --- /dev/null +++ b/gpu4pyscf/lib/pbc/overlap.cuh @@ -0,0 +1,67 @@ +/* + * 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 + +#define PI_POW_1_5 5.568327996831707845 +#define REMOTE_THRESHOLD 50 + +__inline__ __device__ +void vrr_hrr(double *gx, double *rjri, double ai, double aj, double cicj, + int li, int lj, int gout_id, int gout_stride, int nsp_per_block) +{ + int stride_j = li + 1; + int g_size = (li + 1) * (lj + 1); + int gx_len = g_size * nsp_per_block; + double aij = ai + aj; + double aj_aij = aj / aij; + if (gout_id == 0) { + double theta = ai * aj_aij; + double theta_rr = theta * rjri[3*nsp_per_block]; + gx[gx_len*2] = cicj / (aij*sqrt(aij)) * exp(-theta_rr); + } + int lij = li + lj; + if (lij > 0) { + __syncthreads(); + double s0x, s1x, s2x; + double b = .5 / aij; + for (int n = gout_id; n < 3; n += gout_stride) { + double *_gx = gx + n * gx_len; + double xjxi = rjri[n*nsp_per_block]; + double xpa = xjxi * aj_aij; + s0x = _gx[0]; + s1x = xpa * s0x; + _gx[nsp_per_block] = s1x; + for (int i = 1; i < lij; ++i) { + s2x = xpa * s1x + i * b * s0x; + _gx[(i+1)*nsp_per_block] = s2x; + s0x = s1x; + s1x = s2x; + } + for (int j = 0; j < lj; ++j) { + int ij = (lij-j) + j*stride_j; + s1x = _gx[ij*nsp_per_block]; + for (--ij; ij >= j*stride_j; --ij) { + s0x = _gx[ij*nsp_per_block]; + _gx[(ij+stride_j)*nsp_per_block] = s1x - xjxi * s0x; + s1x = s0x; + } + } + } + } + __syncthreads(); +} + From d2b09963375450782ca28b5c54694c75c39a90be Mon Sep 17 00:00:00 2001 From: Qiming Sun Date: Fri, 18 Sep 2026 08:19:19 -0700 Subject: [PATCH 5/7] bugfix --- gpu4pyscf/pbc/grad/pp.py | 5 ++--- gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/gpu4pyscf/pbc/grad/pp.py b/gpu4pyscf/pbc/grad/pp.py index a4bf720c9..0e2d9f422 100644 --- a/gpu4pyscf/pbc/grad/pp.py +++ b/gpu4pyscf/pbc/grad/pp.py @@ -95,8 +95,8 @@ def ppnl_derivatives(cell, dm, kpts=None): opt = int1e.CrossInt1e(pcell, sorted_cell, bvk_kmesh) for rn in range(hl_dim): kern, deriv = derivative_kernels[rn] - as_dm = weights[rn].reshape(nkpts, n_hl*nd, nao) - grad_sigma += _derivatives_intor(opt, as_dm, kern, deriv) + as_dm = weights[rn].reshape(-1, n_hl*nd, nao) + grad_sigma += _derivatives_intor(opt, as_dm, kern, deriv) / nkpts return grad_sigma def _derivatives_intor(cross_int1e, dm, kern, deriv): @@ -142,5 +142,4 @@ def _derivatives_intor(cross_int1e, dm, kern, deriv): 'fakecell for ppnl must have the same number of atoms as the AO cell') grad = (grad[:natm] + grad[natm:]).get() grad_sigma = np.vstack([grad, sigma.get()]) - grad_sigma /= nkpts return grad_sigma diff --git a/gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py b/gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py index d65112793..948c012f3 100644 --- a/gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py +++ b/gpu4pyscf/pbc/grad/tests/test_pbc_pp_grad.py @@ -131,7 +131,7 @@ def _cpu_vppnl_nuc_grad(cell, dm, kpts=None): grad[ia] -= np.einsum('kdpq,kqp->d', dppnl[:, :, p0:p1, :], dm_dmH[:, :, p0:p1]) - return grad.real + return grad.real / nkpts def vppnl_nuc_grad(cell, dm, kpts=None): '''Nuclear gradients of the non-local part of the GTH pseudo potential, From c845e00a0012538e01015a03543840231fc7939c Mon Sep 17 00:00:00 2001 From: Qiming Sun Date: Fri, 18 Sep 2026 08:55:18 -0700 Subject: [PATCH 6/7] lint --- gpu4pyscf/pbc/grad/pp.py | 1 - 1 file changed, 1 deletion(-) diff --git a/gpu4pyscf/pbc/grad/pp.py b/gpu4pyscf/pbc/grad/pp.py index 0e2d9f422..8777634ea 100644 --- a/gpu4pyscf/pbc/grad/pp.py +++ b/gpu4pyscf/pbc/grad/pp.py @@ -105,7 +105,6 @@ def _derivatives_intor(cross_int1e, dm, kern, deriv): cell = cross_int1e.cell assert dm.ndim == 3 assert dm.dtype == np.float64 - nkpts = len(dm) tmp = cross_int1e.cell2.apply_C_dot(dm, axis=2) dm = cross_int1e.cell1.apply_C_dot(tmp, axis=1) From a596e3e74efc50884eb18a8e80fa995a8b6a554f Mon Sep 17 00:00:00 2001 From: Qiming Sun Date: Fri, 18 Sep 2026 13:52:45 -0700 Subject: [PATCH 7/7] Bugfixes --- gpu4pyscf/pbc/grad/tests/test_pbc_krks_stress.py | 6 ++---- gpu4pyscf/pbc/grad/tests/test_pbc_rks_stress.py | 4 ++-- gpu4pyscf/pbc/gto/pseudo/pp_int.py | 2 +- 3 files changed, 5 insertions(+), 7 deletions(-) diff --git a/gpu4pyscf/pbc/grad/tests/test_pbc_krks_stress.py b/gpu4pyscf/pbc/grad/tests/test_pbc_krks_stress.py index 634b29ede..34c5c65fe 100644 --- a/gpu4pyscf/pbc/grad/tests/test_pbc_krks_stress.py +++ b/gpu4pyscf/pbc/grad/tests/test_pbc_krks_stress.py @@ -304,8 +304,7 @@ def test_get_nuc(self): ni = MultiGridNumInt(cell) ni.allow_mesh_reduction = False dat = ni.energy_derivatives(xc, dm, spin=0, kpts=kpts, with_nuc=True, with_j=False)[-3:] - dat += pp_grad._get_pp_nonloc_strain_derivatives( - cell, cell.mesh, cp.array(dm), kpts) + dat += pp_grad.ppnl_derivatives(cell, cp.array(dm), kpts)[-3:] 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-4) @@ -337,8 +336,7 @@ def test_get_pp(self): ni = MultiGridNumInt(cell) ni.allow_mesh_reduction = False dat = ni.energy_derivatives(xc, dm, spin=0, kpts=kpts, with_nuc=True, with_j=False)[-3:] - dat += pp_grad._get_pp_nonloc_strain_derivatives( - cell, cell.mesh, cp.array(dm), kpts) + dat += pp_grad.ppnl_derivatives(cell, cp.array(dm), kpts)[-3:] 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-4) diff --git a/gpu4pyscf/pbc/grad/tests/test_pbc_rks_stress.py b/gpu4pyscf/pbc/grad/tests/test_pbc_rks_stress.py index a76b608aa..66cdf14d6 100644 --- a/gpu4pyscf/pbc/grad/tests/test_pbc_rks_stress.py +++ b/gpu4pyscf/pbc/grad/tests/test_pbc_rks_stress.py @@ -226,7 +226,7 @@ def test_get_nuc(self): assert abs(dat[i,j] - de/2e-4) < 2e-7 def test_get_pp(self): - from gpu4pyscf.pbc.grad.pp import _get_pp_nonloc_strain_derivatives + from gpu4pyscf.pbc.grad.pp import ppnl_derivatives a = np.eye(3) * 5 np.random.seed(5) a += np.random.rand(3, 3) - .5 @@ -240,7 +240,7 @@ def test_get_pp(self): ni = MultiGridNumInt(cell) ni.allow_mesh_reduction = False dat = ni.energy_derivatives(xc, dm, spin=0, with_nuc=True, with_j=False)[-3:] - dat += _get_pp_nonloc_strain_derivatives(cell, cell.mesh, cp.array(dm)) + dat += ppnl_derivatives(cell, cp.array(dm))[-3:] ni = NumInt() kpt = np.zeros(3) for (i, j) in [(0, 0), (0, 1), (0, 2), (2, 1), (2, 2)]: diff --git a/gpu4pyscf/pbc/gto/pseudo/pp_int.py b/gpu4pyscf/pbc/gto/pseudo/pp_int.py index 6daaabaf9..e315d11dd 100644 --- a/gpu4pyscf/pbc/gto/pseudo/pp_int.py +++ b/gpu4pyscf/pbc/gto/pseudo/pp_int.py @@ -85,7 +85,7 @@ def get_pp_nl_gpu(cell, kpts=None): ppnl_half = _int_vnl_gpu(cell, fakecell, hl_blocks, kpts_lst) - is_gamma_point = gamma_point(kpts) + is_gamma_point = gamma_point(kpts_lst) dtype = np.float64 if is_gamma_point else np.complex128 nao = cell.nao ppnl = cp.zeros((nkpts, nao, nao), dtype=dtype)