From 017d3c701e5a3f59cf90e238f63284cee605e5ce Mon Sep 17 00:00:00 2001 From: Abhishek Bagusetty Date: Thu, 3 Sep 2026 14:11:38 -0500 Subject: [PATCH 1/6] cleanup cp vs np usage, early return with warp-divergence in ft_ao.cu --- gpu4pyscf/df/tests/test_df_hessian.py | 2 +- gpu4pyscf/df/tests/test_df_rks_grad.py | 4 +- gpu4pyscf/df/tests/test_df_uhf.py | 4 +- gpu4pyscf/dft/tests/test_numint.py | 3 +- gpu4pyscf/lib/cupy_helper.py | 2 +- gpu4pyscf/lib/pbc/ft_ao.cu | 42 ++++++++++++++----- gpu4pyscf/pbc/df/ft_ao.py | 2 +- gpu4pyscf/pbc/df/int3c2e.py | 2 +- gpu4pyscf/pbc/df/rsdf_builder.py | 15 +++++-- gpu4pyscf/pbc/df/tests/test_pbc_int3c2e.py | 2 +- gpu4pyscf/pbc/dft/multigrid_v2.py | 2 +- gpu4pyscf/pbc/dft/tests/test_pbc_numint.py | 2 +- .../pbc/scf/tests/test_pbc_scf_smearing.py | 4 +- gpu4pyscf/pbc/tools/k2gamma.py | 3 +- gpu4pyscf/qmmm/external_field.py | 4 +- 15 files changed, 61 insertions(+), 32 deletions(-) diff --git a/gpu4pyscf/df/tests/test_df_hessian.py b/gpu4pyscf/df/tests/test_df_hessian.py index a5a365cf7..8061b876a 100644 --- a/gpu4pyscf/df/tests/test_df_hessian.py +++ b/gpu4pyscf/df/tests/test_df_hessian.py @@ -358,7 +358,7 @@ def test_unstable_j2c(self): mo_occ = mf.mo_occ test_hessian_round2 = hobj.partial_hess_elec(mo_energy, mo_coeff, mo_occ) - assert np.max(np.abs(test_hessian_round1 - test_hessian_round2)) < 2e-7 + assert abs(test_hessian_round1 - test_hessian_round2).max().item() < 2e-7 if __name__ == "__main__": print("Full Tests for DF Hessian") diff --git a/gpu4pyscf/df/tests/test_df_rks_grad.py b/gpu4pyscf/df/tests/test_df_rks_grad.py index c69a0b63a..9a4b26d2b 100644 --- a/gpu4pyscf/df/tests/test_df_rks_grad.py +++ b/gpu4pyscf/df/tests/test_df_rks_grad.py @@ -94,8 +94,8 @@ def _check_grad(mol, grid_response=False, xc=xc0, disp=disp0, tol=1e-6): grad_fd = np.array(grad_fd).reshape(-1,3) print('finite difference gradient:') print(grad_fd) - print('difference between analytical and finite difference gradient:', cupy.linalg.norm(g_analy - grad_fd)) - assert(cupy.linalg.norm(g_analy - grad_fd) < tol) + print('difference between analytical and finite difference gradient:', np.linalg.norm(g_analy - grad_fd)) + assert(np.linalg.norm(g_analy - grad_fd) < tol) def _vs_cpu(mol, grid_response=False, xc=xc0, disp=disp0, tol=1e-9): mf = rks.RKS(mol, xc=xc).density_fit(auxbasis=auxbasis0) diff --git a/gpu4pyscf/df/tests/test_df_uhf.py b/gpu4pyscf/df/tests/test_df_uhf.py index cdb3dd2c1..7032f31a7 100644 --- a/gpu4pyscf/df/tests/test_df_uhf.py +++ b/gpu4pyscf/df/tests/test_df_uhf.py @@ -81,8 +81,8 @@ def _check_grad(mol, tol=1e-5, disp=None): grad_fd = np.array(grad_fd).reshape(-1,3) print('finite difference gradient:') print(grad_fd) - print('difference between analytical and finite difference gradient:', cupy.linalg.norm(g_analy - grad_fd)) - assert(cupy.linalg.norm(g_analy - grad_fd) < tol) + print('difference between analytical and finite difference gradient:', np.linalg.norm(g_analy - grad_fd)) + assert(np.linalg.norm(g_analy - grad_fd) < tol) class KnownValues(unittest.TestCase): ''' diff --git a/gpu4pyscf/dft/tests/test_numint.py b/gpu4pyscf/dft/tests/test_numint.py index de51ba2ef..f2a0f77f1 100644 --- a/gpu4pyscf/dft/tests/test_numint.py +++ b/gpu4pyscf/dft/tests/test_numint.py @@ -272,7 +272,8 @@ def test_sparse_index(self): i1 = min(i0+numint.MIN_BLK_SIZE, ngrids) ref = numint._sparse_index( opt._sorted_mol, grids.coords[i0:i1], opt.l_ctr_offsets, ao_loc, opt) - assert all(np.array_equal(r, x) for r, x in zip(ref[1:], dat[i][1:])) + assert all(r.shape == x.shape and bool((r == x).all()) + for r, x in zip(ref[1:], dat[i][1:])) def test_scale_ao(self): ao = cupy.random.rand(1, 3, 256) diff --git a/gpu4pyscf/lib/cupy_helper.py b/gpu4pyscf/lib/cupy_helper.py index 7239603f5..b1234b857 100644 --- a/gpu4pyscf/lib/cupy_helper.py +++ b/gpu4pyscf/lib/cupy_helper.py @@ -1265,7 +1265,7 @@ def batched_vec3_norm2(batched_vec3): ''' einsum('gx,gx->g', vec3, vec3) for the (N,3)-array vec3 ''' - assert type(batched_vec3) is cupy.ndarray + assert isinstance(batched_vec3, cupy.ndarray) assert batched_vec3.dtype == cupy.float64 assert batched_vec3.ndim == 2 assert batched_vec3.shape[0] == 3 or batched_vec3.shape[1] == 3 diff --git a/gpu4pyscf/lib/pbc/ft_ao.cu b/gpu4pyscf/lib/pbc/ft_ao.cu index b4a09ba0f..dbdda8e97 100644 --- a/gpu4pyscf/lib/pbc/ft_ao.cu +++ b/gpu4pyscf/lib/pbc/ft_ao.cu @@ -49,16 +49,15 @@ void ft_ao_bdiv_kernel(double *out, RysIntEnvVars envs, int nGv, double *Gv) int sh_id_in_block = threadIdx.y; int Gv_id_in_block = threadIdx.x; int sh_id = sh_block_id * nsh_per_block + sh_id_in_block; - if (sh_id >= envs.nbas) { - return; - } + int valid = sh_id < envs.nbas; + int sh_id_clamped = valid ? sh_id : envs.nbas - 1; int *atm = envs.atm; int *bas = envs.bas; double *env = envs.env; - int li = bas[sh_id*BAS_SLOTS+ANG_OF]; + int li = bas[sh_id_clamped*BAS_SLOTS+ANG_OF]; int nfi = c_nf[li]; - int iprim = bas[sh_id*BAS_SLOTS+NPRIM_OF]; + int iprim = bas[sh_id_clamped*BAS_SLOTS+NPRIM_OF]; int Gv_id = Gv_block_id * NG_PER_BLOCK + Gv_id_in_block; double kx = 0; double ky = 0; @@ -72,6 +71,7 @@ void ft_ao_bdiv_kernel(double *out, RysIntEnvVars envs, int nGv, double *Gv) int gx_len = (AUXL+1) * FT_AO_THREADS; __shared__ double g[(AUXL+1)*FT_AO_THREADS * 6]; + __shared__ int block_iprim[FT_AO_THREADS/NG_PER_BLOCK]; double *gxR = g + (AUXL+1) * NG_PER_BLOCK * sh_id_in_block + Gv_id_in_block; double *gxI = gxR + gx_len; double *gyR = gxR + gx_len*2; @@ -95,12 +95,29 @@ void ft_ao_bdiv_kernel(double *out, RysIntEnvVars envs, int nGv, double *Gv) double s0zR, s1zR, s2zR; double s0zI, s1zI, s2zI; - int ia = bas[sh_id*BAS_SLOTS+ATOM_OF]; - double *expi = env + bas[sh_id*BAS_SLOTS+PTR_EXP]; - double *ci = env + bas[sh_id*BAS_SLOTS+PTR_COEFF]; + int ia = bas[sh_id_clamped*BAS_SLOTS+ATOM_OF]; + double *expi = env + bas[sh_id_clamped*BAS_SLOTS+PTR_EXP]; + double *ci = env + bas[sh_id_clamped*BAS_SLOTS+PTR_COEFF]; double *ri = env + atm[ia*ATM_SLOTS+PTR_COORD]; - for (int ip = 0; ip < iprim; ++ip) { + // The primitive loop below calls __syncthreads() every iteration, so its + // trip count must be uniform across the whole block. SortedGTO groups + // shells by (l, nprim) but does not align those groups to + // nsh_per_block boundaries, so a block routinely spans two groups with + // different nprim. Loop to the block-wide max instead of this lane's + // own iprim, and guard the per-iteration work so a lane with fewer + // primitives simply does nothing on the extra iterations. + if (Gv_id_in_block == 0) { + block_iprim[sh_id_in_block] = iprim; + } + __syncthreads(); + int max_iprim = 0; +#pragma unroll + for (int i = 0; i < FT_AO_THREADS/NG_PER_BLOCK; ++i) { + max_iprim = max(max_iprim, block_iprim[i]); + } + for (int ip = 0; ip < max_iprim; ++ip) { __syncthreads(); + if (ip < iprim) { double ai = expi[ip]; double xi = ri[0]; double yi = ri[1]; @@ -167,7 +184,9 @@ void ft_ao_bdiv_kernel(double *out, RysIntEnvVars envs, int nGv, double *Gv) s1zI = s2zI; } } + } __syncthreads(); + if (ip < iprim) { #pragma unroll for (int n = 0; n < aux_nf; ++n) { if (n >= nfi) break; @@ -185,11 +204,12 @@ void ft_ao_bdiv_kernel(double *out, RysIntEnvVars envs, int nGv, double *Gv) goutR[n] += xyR * zR - xyI * zI; goutI[n] += xyR * zI + xyI * zR; } + } } - if (Gv_id < nGv) { + if (valid && Gv_id < nGv) { size_t stride = (size_t)nGv * OF_COMPLEX; - double *aft_tensor = out + ((size_t)envs.ao_loc[sh_id] * nGv + Gv_id) * OF_COMPLEX; + double *aft_tensor = out + ((size_t)envs.ao_loc[sh_id_clamped] * nGv + Gv_id) * OF_COMPLEX; #pragma unroll for (int n = 0; n < aux_nf; ++n) { if (n >= nfi) break; diff --git a/gpu4pyscf/pbc/df/ft_ao.py b/gpu4pyscf/pbc/df/ft_ao.py index 6c23e5a97..2f49f27c5 100644 --- a/gpu4pyscf/pbc/df/ft_ao.py +++ b/gpu4pyscf/pbc/df/ft_ao.py @@ -362,7 +362,7 @@ def ft_evaluator(self, batch_size=None, compressing=True, cart=None, dims, tmp = np.empty_like(dims), dims dims[cell.sorted_idx] = tmp ao_loc = cp.asarray(np.append(0, np.cumsum(dims.ravel()))) - ao_loc = np.append(ao_loc[cell.sorted_idx], nao) + ao_loc = cp.append(ao_loc[cell.sorted_idx], nao) ao_loc = cp.asarray(ao_loc, dtype=np.int32) if batch_size is None: diff --git a/gpu4pyscf/pbc/df/int3c2e.py b/gpu4pyscf/pbc/df/int3c2e.py index c1cc7d11d..e84df4645 100644 --- a/gpu4pyscf/pbc/df/int3c2e.py +++ b/gpu4pyscf/pbc/df/int3c2e.py @@ -128,7 +128,7 @@ def sr_aux_e2(cell, auxcell, omega, kpts=None, bvk_kmesh=None, j_only=False): out = cp.empty((nkpts,nkpts,naux,nao,nao), dtype=np.complex128) kk_conserv = double_translation_indices(int3c2e_opt.bvk_kmesh) for k in range(nkpts): - ki_idx, kj_idx = np.where(kk_conserv == k) + ki_idx, kj_idx = cp.where(kk_conserv == k) out[k] = _unpack_cderi_v2(j3c[k], pair_address, kj_idx, conj_mapping, expLk, nao, axis) j3c = None diff --git a/gpu4pyscf/pbc/df/rsdf_builder.py b/gpu4pyscf/pbc/df/rsdf_builder.py index 209ba5a67..c8e94dab1 100644 --- a/gpu4pyscf/pbc/df/rsdf_builder.py +++ b/gpu4pyscf/pbc/df/rsdf_builder.py @@ -404,8 +404,14 @@ def proc(): ctypes.c_int(naux), ctypes.c_int(nao_pairs), ctypes.c_int(p0), ctypes.c_int(p1)) j3c = None - if num_devices > 1: - stream.synchronize() + # store_col_segment enqueues its device->host write asynchronously + # and does not itself block. The next reader of `cderi` is a + # plain host-side read (mydf.loop()) with no ordering relationship + # to this stream, so this segment must be drained here every + # time -- not only when num_devices > 1. Skipping this on a + # single device happened to be safe on this stream/allocator + # combination but is not guaranteed in general. + stream.synchronize() t1 = log.timer_debug1(f'store int3c2e on Device {device_id}', *t1) multi_gpu.run(proc, non_blocking=True) @@ -572,8 +578,9 @@ def proc(): if err != 0: raise RuntimeError('store_col_segment kernel failed') j3c = None - if num_devices > 1: - stream.synchronize() + # See the matching comment in compressed_cderi_j_only: this must + # run unconditionally, not only when num_devices > 1. + stream.synchronize() t1 = log.timer_debug1(f'store int3c2e on Device {device_id}', *t1) multi_gpu.run(proc, non_blocking=True) diff --git a/gpu4pyscf/pbc/df/tests/test_pbc_int3c2e.py b/gpu4pyscf/pbc/df/tests/test_pbc_int3c2e.py index ffb4c26e9..05624b902 100644 --- a/gpu4pyscf/pbc/df/tests/test_pbc_int3c2e.py +++ b/gpu4pyscf/pbc/df/tests/test_pbc_int3c2e.py @@ -324,7 +324,7 @@ def test_contract_dm_kpts(): np.random.seed(9) auxvec = np.random.rand(auxcell.nao) vj = opt.contract_auxvec(opt.auxcell.apply_C_dot(auxvec), kpts=kpts) - ref = cp.einsum('kpqr,r->kpq', j3c, auxvec) + ref = cp.einsum('kpqr,r->kpq', j3c, cp.asarray(auxvec)) # auxvec is host data assert abs(vj - ref).max() < 1e-10 def test_int3c2e_batch_evaluation(): diff --git a/gpu4pyscf/pbc/dft/multigrid_v2.py b/gpu4pyscf/pbc/dft/multigrid_v2.py index b9b5d2df1..664e06ade 100644 --- a/gpu4pyscf/pbc/dft/multigrid_v2.py +++ b/gpu4pyscf/pbc/dft/multigrid_v2.py @@ -84,7 +84,7 @@ def ifft_in_place(x): def unique_with_sort(x): # This function does the same thing as cp.unique(x, return_inverse=True). # It's not super optimized, but for whatever reason, cp.unique is very slow, so this one is better. - assert type(x) is cp.ndarray and (x.dtype == cp.int32 or x.dtype == cp.int64) and x.ndim == 1 + assert isinstance(x, cp.ndarray) and (x.dtype == cp.int32 or x.dtype == cp.int64) and x.ndim == 1 n = x.shape[0] if n <= 1: return x, cp.zeros(n) diff --git a/gpu4pyscf/pbc/dft/tests/test_pbc_numint.py b/gpu4pyscf/pbc/dft/tests/test_pbc_numint.py index 02f4bc0e4..e32a8e9dc 100644 --- a/gpu4pyscf/pbc/dft/tests/test_pbc_numint.py +++ b/gpu4pyscf/pbc/dft/tests/test_pbc_numint.py @@ -366,7 +366,7 @@ def test_uniform_grid_division_mode(self): test_coords = [] for frag_grids in grids.loop_grids(): test_coords.append(frag_grids.coords) - assert np.max(np.abs(frag_grids.weights - ref_weight)) < 1e-14 + assert cp.max(cp.abs(frag_grids.weights - ref_weight)) < 1e-14 test_coords = cp.vstack(test_coords).get() ref_coords = ref_coords[np.lexsort((ref_coords[:, 2], ref_coords[:, 1], ref_coords[:, 0])), :] diff --git a/gpu4pyscf/pbc/scf/tests/test_pbc_scf_smearing.py b/gpu4pyscf/pbc/scf/tests/test_pbc_scf_smearing.py index cada001a1..c7c49b57b 100644 --- a/gpu4pyscf/pbc/scf/tests/test_pbc_scf_smearing.py +++ b/gpu4pyscf/pbc/scf/tests/test_pbc_scf_smearing.py @@ -40,7 +40,7 @@ def test_krhf_smearing(self): mf = cell.KRHF(kpts=cell.make_kpts([2,1,1])).to_gpu() mf = mf.smearing(0.1, 'fermi') nkpts = len(mf.kpts) - mo_energy_kpts = cp.array([cp.arange(nao)*.2+cp.cos(i+.5)*.1 for i in range(nkpts)]) + mo_energy_kpts = cp.array([cp.arange(nao)*.2+np.cos(i+.5)*.1 for i in range(nkpts)]) mf.get_occ(mo_energy_kpts) self.assertAlmostEqual(mf.entropy, 6.1656394960533021/2, 9) @@ -56,7 +56,7 @@ def test_kuhf_smearing(self): mf = cell.KUHF(kpts=cell.make_kpts([2,1,1])).to_gpu() mf = mf.smearing(0.1, 'fermi') nkpts = len(mf.kpts) - mo_energy_kpts = cp.array([cp.arange(nao)*.2+cp.cos(i+.5)*.1 for i in range(nkpts)]) + mo_energy_kpts = cp.array([cp.arange(nao)*.2+np.cos(i+.5)*.1 for i in range(nkpts)]) mo_energy_kpts = cp.array([mo_energy_kpts, mo_energy_kpts+cp.cos(mo_energy_kpts)*.02]) mf.get_occ(mo_energy_kpts) self.assertAlmostEqual(mf.entropy, 6.1803390081500869/2, 9) diff --git a/gpu4pyscf/pbc/tools/k2gamma.py b/gpu4pyscf/pbc/tools/k2gamma.py index f27685bfa..5d0215b34 100644 --- a/gpu4pyscf/pbc/tools/k2gamma.py +++ b/gpu4pyscf/pbc/tools/k2gamma.py @@ -82,7 +82,8 @@ def double_translation_indices(kmesh): tz = cp.array(translation_map(kmesh[2]), dtype=np.int32) idx = cp.ravel_multi_index([tx[:,None,None,:,None,None], ty[None,:,None,None,:,None], - tz[None,None,:,None,None,:]], kmesh) + tz[None,None,:,None,None,:]], + tuple(int(n) for n in kmesh)) nk = np.prod(kmesh) return idx.reshape(nk, nk) diff --git a/gpu4pyscf/qmmm/external_field.py b/gpu4pyscf/qmmm/external_field.py index 96c85a865..642d62713 100644 --- a/gpu4pyscf/qmmm/external_field.py +++ b/gpu4pyscf/qmmm/external_field.py @@ -53,7 +53,7 @@ def __init__(self, method, electric_field=None, origin=None): if electric_field is not None: electric_field = cp.asarray(electric_field) - assert type(electric_field) is cp.ndarray + assert isinstance(electric_field, cp.ndarray) assert electric_field.shape == (3,) self.electric_field = electric_field @@ -61,7 +61,7 @@ def __init__(self, method, electric_field=None, origin=None): origin = np.zeros(3) else: origin = cp.asarray(origin).get() - assert type(origin) is np.ndarray + assert isinstance(origin, np.ndarray) assert origin.shape == (3,) self.origin = origin From f60fbaa39f0c0978b4ff54b13fcce5efcf38fe93 Mon Sep 17 00:00:00 2001 From: Abhishek Bagusetty Date: Fri, 4 Sep 2026 17:07:18 +0000 Subject: [PATCH 2/6] Fix a bug related to finding the dims for single kpts --- gpu4pyscf/pbc/dft/multigrid_v2.py | 4 ++-- gpu4pyscf/pbc/dft/multigrid_v3.py | 4 ++-- gpu4pyscf/pbc/dft/tests/test_multigrid_v2.py | 16 ++++++++++++++++ gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py | 16 ++++++++++++++++ 4 files changed, 36 insertions(+), 4 deletions(-) diff --git a/gpu4pyscf/pbc/dft/multigrid_v2.py b/gpu4pyscf/pbc/dft/multigrid_v2.py index 664e06ade..00ffb5524 100644 --- a/gpu4pyscf/pbc/dft/multigrid_v2.py +++ b/gpu4pyscf/pbc/dft/multigrid_v2.py @@ -1398,7 +1398,7 @@ def convert_xc_on_g_mesh_to_fock_gradient( def get_nuc(ni, kpts=None): if ni.sorted_gaussian_pairs is None: ni.build() - is_single_kpt = kpts is not None and kpts.ndim == 1 + is_single_kpt = kpts is None or kpts.ndim == 1 if kpts is None: kpts = np.zeros((1, 3)) else: @@ -1417,7 +1417,7 @@ def get_pp(ni, kpts=None): """Get the periodic pseudopotential nuc-el AO matrix, with G=0 removed.""" if ni.sorted_gaussian_pairs is None: ni.build() - is_single_kpt = kpts is not None and kpts.ndim == 1 + is_single_kpt = kpts is None or kpts.ndim == 1 if kpts is None: kpts = np.zeros((1, 3)) else: diff --git a/gpu4pyscf/pbc/dft/multigrid_v3.py b/gpu4pyscf/pbc/dft/multigrid_v3.py index c9aab23e6..bdf6e032d 100644 --- a/gpu4pyscf/pbc/dft/multigrid_v3.py +++ b/gpu4pyscf/pbc/dft/multigrid_v3.py @@ -1535,7 +1535,7 @@ def get_rho(ni, dm_kpts, kpts=None): def get_nuc(ni, kpts=None): cell = ni.cell - is_single_kpt = kpts is not None and kpts.ndim == 1 + is_single_kpt = kpts is None or kpts.ndim == 1 if kpts is None: kpts = np.zeros((1, 3)) else: @@ -1561,7 +1561,7 @@ def get_pp(ni, kpts=None): log = logger.new_logger(cell) t0 = log.init_timer() - is_single_kpt = kpts is not None and kpts.ndim == 1 + is_single_kpt = kpts is None or kpts.ndim == 1 if kpts is None: kpts = np.zeros((1, 3)) else: diff --git a/gpu4pyscf/pbc/dft/tests/test_multigrid_v2.py b/gpu4pyscf/pbc/dft/tests/test_multigrid_v2.py index c7548e9d8..18c616bde 100644 --- a/gpu4pyscf/pbc/dft/tests/test_multigrid_v2.py +++ b/gpu4pyscf/pbc/dft/tests/test_multigrid_v2.py @@ -112,6 +112,22 @@ def test_get_nuc_kpts_nonorth(self): self.assertEqual(out.shape, ref.shape) self.assertAlmostEqual(abs(ref-out).max(), 0, 8) + def test_get_nuc_get_pp_single_kpt_ndim(self): + nao = cell_orth.nao + ni = multigrid.MultiGridNumInt(cell_orth) + self.assertEqual(ni.get_nuc().ndim, 2) + self.assertEqual(ni.get_nuc().shape, (nao, nao)) + self.assertEqual(ni.get_pp().ndim, 2) + self.assertEqual(ni.get_pp().shape, (nao, nao)) + + single_kpt = np.zeros(3) + self.assertEqual(ni.get_nuc(single_kpt).ndim, 2) + self.assertEqual(ni.get_pp(single_kpt).ndim, 2) + + multi_kpts = np.zeros((1, 3)) + self.assertEqual(ni.get_nuc(multi_kpts).ndim, 3) + self.assertEqual(ni.get_pp(multi_kpts).ndim, 3) + def test_get_rho(self): nao = cell_orth.nao np.random.seed(2) diff --git a/gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py b/gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py index 64b505f30..01680d3c9 100644 --- a/gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py +++ b/gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py @@ -225,6 +225,22 @@ def test_get_nuc_kpts_nonorth(self): self.assertEqual(out.shape, ref.shape) self.assertAlmostEqual(abs(ref-out).max(), 0, 7) + def test_get_nuc_get_pp_single_kpt_ndim(self): + nao = cell_orth.nao + ni = multigrid.MultiGridNumInt(cell_orth) + self.assertEqual(ni.get_nuc().ndim, 2) + self.assertEqual(ni.get_nuc().shape, (nao, nao)) + self.assertEqual(ni.get_pp().ndim, 2) + self.assertEqual(ni.get_pp().shape, (nao, nao)) + + single_kpt = np.zeros(3) + self.assertEqual(ni.get_nuc(single_kpt).ndim, 2) + self.assertEqual(ni.get_pp(single_kpt).ndim, 2) + + multi_kpts = np.zeros((1, 3)) + self.assertEqual(ni.get_nuc(multi_kpts).ndim, 3) + self.assertEqual(ni.get_pp(multi_kpts).ndim, 3) + def test_get_rho(self): nao = cell_orth.nao np.random.seed(2) From 1cd5adbc211fd56528308f86271b7101b800950b Mon Sep 17 00:00:00 2001 From: Abhishek Bagusetty Date: Fri, 4 Sep 2026 19:07:29 +0000 Subject: [PATCH 3/6] Unify CI on pyscf 2.14.0 and compare multigrid get_nuc/get_pp by value, not shape --- .github/workflows/unittest.yml | 2 +- gpu4pyscf/pbc/dft/tests/test_multigrid_v2.py | 18 +++++++++--------- gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py | 18 +++++++++--------- 3 files changed, 19 insertions(+), 19 deletions(-) diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml index e6958a3aa..40909c309 100644 --- a/.github/workflows/unittest.yml +++ b/.github/workflows/unittest.yml @@ -59,4 +59,4 @@ jobs: -v $GITHUB_WORKSPACE:/workspace \ -v ~/.cache/pip:/root/.cache/pip \ pyscf/gpu4pyscf-devel:pyscf-2.14 \ - /bin/bash -c "cd /workspace && pip3 install -r requirements.txt && pip3 install pyscf==2.8 scipy==1.17 && source build.sh && pytest -m 'not slow and not benchmark and not special' --cov=/workspace --durations=50 && rm -rf .pytest_cache" + /bin/bash -c "cd /workspace && pip3 install -r requirements.txt && pip3 install MCFun && source build.sh && pytest -m 'not slow and not benchmark and not special' --cov=/workspace --durations=50 && rm -rf .pytest_cache" diff --git a/gpu4pyscf/pbc/dft/tests/test_multigrid_v2.py b/gpu4pyscf/pbc/dft/tests/test_multigrid_v2.py index 18c616bde..620e42474 100644 --- a/gpu4pyscf/pbc/dft/tests/test_multigrid_v2.py +++ b/gpu4pyscf/pbc/dft/tests/test_multigrid_v2.py @@ -83,21 +83,21 @@ def tearDownModule(): class KnownValues(unittest.TestCase): def test_get_pp(self): - ref = MultiGridNumInt_cpu(cell_orth).get_pp() - out = multigrid.MultiGridNumInt(cell_orth).get_pp().get() - # self.assertEqual(out.shape, ref.shape) + nao = cell_orth.nao + ref = MultiGridNumInt_cpu(cell_orth).get_pp().reshape(nao, nao) + out = multigrid.MultiGridNumInt(cell_orth).get_pp().get().reshape(nao, nao) self.assertAlmostEqual(abs(ref-out).max(), 0, 8) def test_get_nuc(self): - ref = MultiGridNumInt_cpu(cell_orth).get_nuc() - out = multigrid.MultiGridNumInt(cell_orth).get_nuc().get() - # self.assertEqual(out.shape, ref.shape) + nao = cell_orth.nao + ref = MultiGridNumInt_cpu(cell_orth).get_nuc().reshape(nao, nao) + out = multigrid.MultiGridNumInt(cell_orth).get_nuc().get().reshape(nao, nao) self.assertAlmostEqual(abs(ref-out).max(), 0, 8) def test_get_nuc_nonorth(self): - ref = MultiGridNumInt_cpu(cell_nonorth).get_nuc() - out = multigrid.MultiGridNumInt(cell_nonorth).get_nuc().get() - # self.assertEqual(out.shape, ref.shape) + nao = cell_nonorth.nao + ref = MultiGridNumInt_cpu(cell_nonorth).get_nuc().reshape(nao, nao) + out = multigrid.MultiGridNumInt(cell_nonorth).get_nuc().get().reshape(nao, nao) self.assertAlmostEqual(abs(ref-out).max(), 0, 8) def test_get_nuc_kpts(self): diff --git a/gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py b/gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py index 01680d3c9..11acc134b 100644 --- a/gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py +++ b/gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py @@ -196,21 +196,21 @@ def eval_nucG_SI_gradient(cell, mesh, rho_g): class KnownValues(unittest.TestCase): def test_get_pp(self): - ref = MultiGridNumInt_cpu(cell_orth).get_pp() - out = multigrid.MultiGridNumInt(cell_orth).get_pp().get() - self.assertEqual(out.shape, ref.shape) + nao = cell_orth.nao + ref = MultiGridNumInt_cpu(cell_orth).get_pp().reshape(nao, nao) + out = multigrid.MultiGridNumInt(cell_orth).get_pp().get().reshape(nao, nao) self.assertAlmostEqual(abs(ref-out).max(), 0, 8) def test_get_nuc(self): - ref = MultiGridNumInt_cpu(cell_orth).get_nuc() - out = multigrid.MultiGridNumInt(cell_orth).get_nuc().get() - self.assertEqual(out.shape, ref.shape) + nao = cell_orth.nao + ref = MultiGridNumInt_cpu(cell_orth).get_nuc().reshape(nao, nao) + out = multigrid.MultiGridNumInt(cell_orth).get_nuc().get().reshape(nao, nao) self.assertAlmostEqual(abs(ref-out).max(), 0, 8) def test_get_nuc_nonorth(self): - ref = MultiGridNumInt_cpu(cell_nonorth).get_nuc() - out = multigrid.MultiGridNumInt(cell_nonorth).get_nuc().get() - self.assertEqual(out.shape, ref.shape) + nao = cell_nonorth.nao + ref = MultiGridNumInt_cpu(cell_nonorth).get_nuc().reshape(nao, nao) + out = multigrid.MultiGridNumInt(cell_nonorth).get_nuc().get().reshape(nao, nao) self.assertAlmostEqual(abs(ref-out).max(), 0, 7) def test_get_nuc_kpts(self): From 90e57af816cafa7a19988e9bb5281b9362e174be Mon Sep 17 00:00:00 2001 From: Abhishek Bagusetty Date: Tue, 8 Sep 2026 17:58:00 +0000 Subject: [PATCH 4/6] Revert local is_single_kpt/CI-pin changes; superseded by upstream fix --- .github/workflows/unittest.yml | 2 +- gpu4pyscf/pbc/dft/multigrid_v2.py | 6 ++-- gpu4pyscf/pbc/dft/multigrid_v3.py | 4 +-- gpu4pyscf/pbc/dft/tests/test_multigrid_v2.py | 34 ++++++-------------- gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py | 34 ++++++-------------- 5 files changed, 24 insertions(+), 56 deletions(-) diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml index 40909c309..e6958a3aa 100644 --- a/.github/workflows/unittest.yml +++ b/.github/workflows/unittest.yml @@ -59,4 +59,4 @@ jobs: -v $GITHUB_WORKSPACE:/workspace \ -v ~/.cache/pip:/root/.cache/pip \ pyscf/gpu4pyscf-devel:pyscf-2.14 \ - /bin/bash -c "cd /workspace && pip3 install -r requirements.txt && pip3 install MCFun && source build.sh && pytest -m 'not slow and not benchmark and not special' --cov=/workspace --durations=50 && rm -rf .pytest_cache" + /bin/bash -c "cd /workspace && pip3 install -r requirements.txt && pip3 install pyscf==2.8 scipy==1.17 && source build.sh && pytest -m 'not slow and not benchmark and not special' --cov=/workspace --durations=50 && rm -rf .pytest_cache" diff --git a/gpu4pyscf/pbc/dft/multigrid_v2.py b/gpu4pyscf/pbc/dft/multigrid_v2.py index 00ffb5524..b9b5d2df1 100644 --- a/gpu4pyscf/pbc/dft/multigrid_v2.py +++ b/gpu4pyscf/pbc/dft/multigrid_v2.py @@ -84,7 +84,7 @@ def ifft_in_place(x): def unique_with_sort(x): # This function does the same thing as cp.unique(x, return_inverse=True). # It's not super optimized, but for whatever reason, cp.unique is very slow, so this one is better. - assert isinstance(x, cp.ndarray) and (x.dtype == cp.int32 or x.dtype == cp.int64) and x.ndim == 1 + assert type(x) is cp.ndarray and (x.dtype == cp.int32 or x.dtype == cp.int64) and x.ndim == 1 n = x.shape[0] if n <= 1: return x, cp.zeros(n) @@ -1398,7 +1398,7 @@ def convert_xc_on_g_mesh_to_fock_gradient( def get_nuc(ni, kpts=None): if ni.sorted_gaussian_pairs is None: ni.build() - is_single_kpt = kpts is None or kpts.ndim == 1 + is_single_kpt = kpts is not None and kpts.ndim == 1 if kpts is None: kpts = np.zeros((1, 3)) else: @@ -1417,7 +1417,7 @@ def get_pp(ni, kpts=None): """Get the periodic pseudopotential nuc-el AO matrix, with G=0 removed.""" if ni.sorted_gaussian_pairs is None: ni.build() - is_single_kpt = kpts is None or kpts.ndim == 1 + is_single_kpt = kpts is not None and kpts.ndim == 1 if kpts is None: kpts = np.zeros((1, 3)) else: diff --git a/gpu4pyscf/pbc/dft/multigrid_v3.py b/gpu4pyscf/pbc/dft/multigrid_v3.py index bdf6e032d..c9aab23e6 100644 --- a/gpu4pyscf/pbc/dft/multigrid_v3.py +++ b/gpu4pyscf/pbc/dft/multigrid_v3.py @@ -1535,7 +1535,7 @@ def get_rho(ni, dm_kpts, kpts=None): def get_nuc(ni, kpts=None): cell = ni.cell - is_single_kpt = kpts is None or kpts.ndim == 1 + is_single_kpt = kpts is not None and kpts.ndim == 1 if kpts is None: kpts = np.zeros((1, 3)) else: @@ -1561,7 +1561,7 @@ def get_pp(ni, kpts=None): log = logger.new_logger(cell) t0 = log.init_timer() - is_single_kpt = kpts is None or kpts.ndim == 1 + is_single_kpt = kpts is not None and kpts.ndim == 1 if kpts is None: kpts = np.zeros((1, 3)) else: diff --git a/gpu4pyscf/pbc/dft/tests/test_multigrid_v2.py b/gpu4pyscf/pbc/dft/tests/test_multigrid_v2.py index 620e42474..c7548e9d8 100644 --- a/gpu4pyscf/pbc/dft/tests/test_multigrid_v2.py +++ b/gpu4pyscf/pbc/dft/tests/test_multigrid_v2.py @@ -83,21 +83,21 @@ def tearDownModule(): class KnownValues(unittest.TestCase): def test_get_pp(self): - nao = cell_orth.nao - ref = MultiGridNumInt_cpu(cell_orth).get_pp().reshape(nao, nao) - out = multigrid.MultiGridNumInt(cell_orth).get_pp().get().reshape(nao, nao) + ref = MultiGridNumInt_cpu(cell_orth).get_pp() + out = multigrid.MultiGridNumInt(cell_orth).get_pp().get() + # self.assertEqual(out.shape, ref.shape) self.assertAlmostEqual(abs(ref-out).max(), 0, 8) def test_get_nuc(self): - nao = cell_orth.nao - ref = MultiGridNumInt_cpu(cell_orth).get_nuc().reshape(nao, nao) - out = multigrid.MultiGridNumInt(cell_orth).get_nuc().get().reshape(nao, nao) + ref = MultiGridNumInt_cpu(cell_orth).get_nuc() + out = multigrid.MultiGridNumInt(cell_orth).get_nuc().get() + # self.assertEqual(out.shape, ref.shape) self.assertAlmostEqual(abs(ref-out).max(), 0, 8) def test_get_nuc_nonorth(self): - nao = cell_nonorth.nao - ref = MultiGridNumInt_cpu(cell_nonorth).get_nuc().reshape(nao, nao) - out = multigrid.MultiGridNumInt(cell_nonorth).get_nuc().get().reshape(nao, nao) + ref = MultiGridNumInt_cpu(cell_nonorth).get_nuc() + out = multigrid.MultiGridNumInt(cell_nonorth).get_nuc().get() + # self.assertEqual(out.shape, ref.shape) self.assertAlmostEqual(abs(ref-out).max(), 0, 8) def test_get_nuc_kpts(self): @@ -112,22 +112,6 @@ def test_get_nuc_kpts_nonorth(self): self.assertEqual(out.shape, ref.shape) self.assertAlmostEqual(abs(ref-out).max(), 0, 8) - def test_get_nuc_get_pp_single_kpt_ndim(self): - nao = cell_orth.nao - ni = multigrid.MultiGridNumInt(cell_orth) - self.assertEqual(ni.get_nuc().ndim, 2) - self.assertEqual(ni.get_nuc().shape, (nao, nao)) - self.assertEqual(ni.get_pp().ndim, 2) - self.assertEqual(ni.get_pp().shape, (nao, nao)) - - single_kpt = np.zeros(3) - self.assertEqual(ni.get_nuc(single_kpt).ndim, 2) - self.assertEqual(ni.get_pp(single_kpt).ndim, 2) - - multi_kpts = np.zeros((1, 3)) - self.assertEqual(ni.get_nuc(multi_kpts).ndim, 3) - self.assertEqual(ni.get_pp(multi_kpts).ndim, 3) - def test_get_rho(self): nao = cell_orth.nao np.random.seed(2) diff --git a/gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py b/gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py index 11acc134b..64b505f30 100644 --- a/gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py +++ b/gpu4pyscf/pbc/dft/tests/test_multigrid_v3.py @@ -196,21 +196,21 @@ def eval_nucG_SI_gradient(cell, mesh, rho_g): class KnownValues(unittest.TestCase): def test_get_pp(self): - nao = cell_orth.nao - ref = MultiGridNumInt_cpu(cell_orth).get_pp().reshape(nao, nao) - out = multigrid.MultiGridNumInt(cell_orth).get_pp().get().reshape(nao, nao) + ref = MultiGridNumInt_cpu(cell_orth).get_pp() + out = multigrid.MultiGridNumInt(cell_orth).get_pp().get() + self.assertEqual(out.shape, ref.shape) self.assertAlmostEqual(abs(ref-out).max(), 0, 8) def test_get_nuc(self): - nao = cell_orth.nao - ref = MultiGridNumInt_cpu(cell_orth).get_nuc().reshape(nao, nao) - out = multigrid.MultiGridNumInt(cell_orth).get_nuc().get().reshape(nao, nao) + ref = MultiGridNumInt_cpu(cell_orth).get_nuc() + out = multigrid.MultiGridNumInt(cell_orth).get_nuc().get() + self.assertEqual(out.shape, ref.shape) self.assertAlmostEqual(abs(ref-out).max(), 0, 8) def test_get_nuc_nonorth(self): - nao = cell_nonorth.nao - ref = MultiGridNumInt_cpu(cell_nonorth).get_nuc().reshape(nao, nao) - out = multigrid.MultiGridNumInt(cell_nonorth).get_nuc().get().reshape(nao, nao) + ref = MultiGridNumInt_cpu(cell_nonorth).get_nuc() + out = multigrid.MultiGridNumInt(cell_nonorth).get_nuc().get() + self.assertEqual(out.shape, ref.shape) self.assertAlmostEqual(abs(ref-out).max(), 0, 7) def test_get_nuc_kpts(self): @@ -225,22 +225,6 @@ def test_get_nuc_kpts_nonorth(self): self.assertEqual(out.shape, ref.shape) self.assertAlmostEqual(abs(ref-out).max(), 0, 7) - def test_get_nuc_get_pp_single_kpt_ndim(self): - nao = cell_orth.nao - ni = multigrid.MultiGridNumInt(cell_orth) - self.assertEqual(ni.get_nuc().ndim, 2) - self.assertEqual(ni.get_nuc().shape, (nao, nao)) - self.assertEqual(ni.get_pp().ndim, 2) - self.assertEqual(ni.get_pp().shape, (nao, nao)) - - single_kpt = np.zeros(3) - self.assertEqual(ni.get_nuc(single_kpt).ndim, 2) - self.assertEqual(ni.get_pp(single_kpt).ndim, 2) - - multi_kpts = np.zeros((1, 3)) - self.assertEqual(ni.get_nuc(multi_kpts).ndim, 3) - self.assertEqual(ni.get_pp(multi_kpts).ndim, 3) - def test_get_rho(self): nao = cell_orth.nao np.random.seed(2) From 883e0654dc0a285769ef5794523ef77bc9995c25 Mon Sep 17 00:00:00 2001 From: Abhishek Bagusetty Date: Tue, 8 Sep 2026 18:01:13 +0000 Subject: [PATCH 5/6] Align NG_PER_BLOCK to FT_AO_THREADS in ft_ao.cu, dropping the block_iprim workaround --- gpu4pyscf/lib/pbc/ft_ao.cu | 28 +++++----------------------- 1 file changed, 5 insertions(+), 23 deletions(-) diff --git a/gpu4pyscf/lib/pbc/ft_ao.cu b/gpu4pyscf/lib/pbc/ft_ao.cu index dbdda8e97..9e80496dd 100644 --- a/gpu4pyscf/lib/pbc/ft_ao.cu +++ b/gpu4pyscf/lib/pbc/ft_ao.cu @@ -30,8 +30,11 @@ #endif #define WARPS 8 #define THREADS 256 -#define NG_PER_BLOCK WARP_SIZE #define FT_AO_THREADS (WARP_SIZE*4) +// One shell per block (nsh_per_block == 1): every thread in the block then +// sees the same shell's iprim, so the primitive loop's __syncthreads() trip +// count is uniform without needing a per-block max-iprim workaround. +#define NG_PER_BLOCK FT_AO_THREADS #define GOUT_WIDTH 30 // pi^1.5 #define OVERLAP_FAC 5.56832799683170787 @@ -71,7 +74,6 @@ void ft_ao_bdiv_kernel(double *out, RysIntEnvVars envs, int nGv, double *Gv) int gx_len = (AUXL+1) * FT_AO_THREADS; __shared__ double g[(AUXL+1)*FT_AO_THREADS * 6]; - __shared__ int block_iprim[FT_AO_THREADS/NG_PER_BLOCK]; double *gxR = g + (AUXL+1) * NG_PER_BLOCK * sh_id_in_block + Gv_id_in_block; double *gxI = gxR + gx_len; double *gyR = gxR + gx_len*2; @@ -99,25 +101,8 @@ void ft_ao_bdiv_kernel(double *out, RysIntEnvVars envs, int nGv, double *Gv) double *expi = env + bas[sh_id_clamped*BAS_SLOTS+PTR_EXP]; double *ci = env + bas[sh_id_clamped*BAS_SLOTS+PTR_COEFF]; double *ri = env + atm[ia*ATM_SLOTS+PTR_COORD]; - // The primitive loop below calls __syncthreads() every iteration, so its - // trip count must be uniform across the whole block. SortedGTO groups - // shells by (l, nprim) but does not align those groups to - // nsh_per_block boundaries, so a block routinely spans two groups with - // different nprim. Loop to the block-wide max instead of this lane's - // own iprim, and guard the per-iteration work so a lane with fewer - // primitives simply does nothing on the extra iterations. - if (Gv_id_in_block == 0) { - block_iprim[sh_id_in_block] = iprim; - } - __syncthreads(); - int max_iprim = 0; -#pragma unroll - for (int i = 0; i < FT_AO_THREADS/NG_PER_BLOCK; ++i) { - max_iprim = max(max_iprim, block_iprim[i]); - } - for (int ip = 0; ip < max_iprim; ++ip) { + for (int ip = 0; ip < iprim; ++ip) { __syncthreads(); - if (ip < iprim) { double ai = expi[ip]; double xi = ri[0]; double yi = ri[1]; @@ -184,9 +169,7 @@ void ft_ao_bdiv_kernel(double *out, RysIntEnvVars envs, int nGv, double *Gv) s1zI = s2zI; } } - } __syncthreads(); - if (ip < iprim) { #pragma unroll for (int n = 0; n < aux_nf; ++n) { if (n >= nfi) break; @@ -204,7 +187,6 @@ void ft_ao_bdiv_kernel(double *out, RysIntEnvVars envs, int nGv, double *Gv) goutR[n] += xyR * zR - xyI * zI; goutI[n] += xyR * zI + xyI * zR; } - } } if (valid && Gv_id < nGv) { From c981e092835fcb971981535cc4739101444f9e4c Mon Sep 17 00:00:00 2001 From: Abhishek Bagusetty Date: Thu, 24 Sep 2026 15:11:08 -0500 Subject: [PATCH 6/6] Revert to upstream --- gpu4pyscf/lib/pbc/ft_ao.cu | 73 ++++++++++---------------------------- 1 file changed, 19 insertions(+), 54 deletions(-) diff --git a/gpu4pyscf/lib/pbc/ft_ao.cu b/gpu4pyscf/lib/pbc/ft_ao.cu index dbdda8e97..9d519f42a 100644 --- a/gpu4pyscf/lib/pbc/ft_ao.cu +++ b/gpu4pyscf/lib/pbc/ft_ao.cu @@ -20,25 +20,11 @@ #include #include "gvhf-rys/vhf.cuh" #include "gvhf-rys/rys_contract_k.cuh" +#include "ft_ao.cuh" -// WARP_SIZE: compile-time constant used for shared-memory sizing. -// `warpSize` (HIP/CUDA device-runtime built-in) is not constexpr, -// so we keep a literal here. Guarded so the build can override -// it (e.g. -DWARP_SIZE=64) for future wider-wavefront targets. -#ifndef WARP_SIZE -#define WARP_SIZE 32 -#endif -#define WARPS 8 #define THREADS 256 -#define NG_PER_BLOCK WARP_SIZE -#define FT_AO_THREADS (WARP_SIZE*4) #define GOUT_WIDTH 30 -// pi^1.5 -#define OVERLAP_FAC 5.56832799683170787 -#define OF_COMPLEX 2 #define POOL_SIZE 65536 -#define AUXL 6 -#define AUXNF ((AUXL+1)*(AUXL+2)/2) __global__ static void ft_ao_bdiv_kernel(double *out, RysIntEnvVars envs, int nGv, double *Gv) @@ -49,15 +35,16 @@ void ft_ao_bdiv_kernel(double *out, RysIntEnvVars envs, int nGv, double *Gv) int sh_id_in_block = threadIdx.y; int Gv_id_in_block = threadIdx.x; int sh_id = sh_block_id * nsh_per_block + sh_id_in_block; - int valid = sh_id < envs.nbas; - int sh_id_clamped = valid ? sh_id : envs.nbas - 1; + if (sh_id >= envs.nbas) { + return; + } int *atm = envs.atm; int *bas = envs.bas; double *env = envs.env; - int li = bas[sh_id_clamped*BAS_SLOTS+ANG_OF]; + int li = bas[sh_id*BAS_SLOTS+ANG_OF]; int nfi = c_nf[li]; - int iprim = bas[sh_id_clamped*BAS_SLOTS+NPRIM_OF]; + int iprim = bas[sh_id*BAS_SLOTS+NPRIM_OF]; int Gv_id = Gv_block_id * NG_PER_BLOCK + Gv_id_in_block; double kx = 0; double ky = 0; @@ -71,7 +58,6 @@ void ft_ao_bdiv_kernel(double *out, RysIntEnvVars envs, int nGv, double *Gv) int gx_len = (AUXL+1) * FT_AO_THREADS; __shared__ double g[(AUXL+1)*FT_AO_THREADS * 6]; - __shared__ int block_iprim[FT_AO_THREADS/NG_PER_BLOCK]; double *gxR = g + (AUXL+1) * NG_PER_BLOCK * sh_id_in_block + Gv_id_in_block; double *gxI = gxR + gx_len; double *gyR = gxR + gx_len*2; @@ -80,11 +66,10 @@ void ft_ao_bdiv_kernel(double *out, RysIntEnvVars envs, int nGv, double *Gv) double *gzI = gxR + gx_len*5; int *idx = _c_cartesian_lexical_xyz + lex_xyz_offset(li); - constexpr int aux_nf = (AUXL+1)*(AUXL+2)/2; - double goutR[aux_nf]; - double goutI[aux_nf]; + double goutR[AUXNF]; + double goutI[AUXNF]; #pragma unroll - for (int n = 0; n < aux_nf; ++n) { + for (int n = 0; n < AUXNF; ++n) { goutR[n] = 0.; goutI[n] = 0.; } @@ -95,29 +80,12 @@ void ft_ao_bdiv_kernel(double *out, RysIntEnvVars envs, int nGv, double *Gv) double s0zR, s1zR, s2zR; double s0zI, s1zI, s2zI; - int ia = bas[sh_id_clamped*BAS_SLOTS+ATOM_OF]; - double *expi = env + bas[sh_id_clamped*BAS_SLOTS+PTR_EXP]; - double *ci = env + bas[sh_id_clamped*BAS_SLOTS+PTR_COEFF]; + int ia = bas[sh_id*BAS_SLOTS+ATOM_OF]; + double *expi = env + bas[sh_id*BAS_SLOTS+PTR_EXP]; + double *ci = env + bas[sh_id*BAS_SLOTS+PTR_COEFF]; double *ri = env + atm[ia*ATM_SLOTS+PTR_COORD]; - // The primitive loop below calls __syncthreads() every iteration, so its - // trip count must be uniform across the whole block. SortedGTO groups - // shells by (l, nprim) but does not align those groups to - // nsh_per_block boundaries, so a block routinely spans two groups with - // different nprim. Loop to the block-wide max instead of this lane's - // own iprim, and guard the per-iteration work so a lane with fewer - // primitives simply does nothing on the extra iterations. - if (Gv_id_in_block == 0) { - block_iprim[sh_id_in_block] = iprim; - } - __syncthreads(); - int max_iprim = 0; -#pragma unroll - for (int i = 0; i < FT_AO_THREADS/NG_PER_BLOCK; ++i) { - max_iprim = max(max_iprim, block_iprim[i]); - } - for (int ip = 0; ip < max_iprim; ++ip) { + for (int ip = 0; ip < iprim; ++ip) { __syncthreads(); - if (ip < iprim) { double ai = expi[ip]; double xi = ri[0]; double yi = ri[1]; @@ -184,11 +152,9 @@ void ft_ao_bdiv_kernel(double *out, RysIntEnvVars envs, int nGv, double *Gv) s1zI = s2zI; } } - } __syncthreads(); - if (ip < iprim) { #pragma unroll - for (int n = 0; n < aux_nf; ++n) { + for (int n = 0; n < AUXNF; ++n) { if (n >= nfi) break; int addrx = idx[n*3+0] * NG_PER_BLOCK; int addry = idx[n*3+1] * NG_PER_BLOCK; @@ -204,14 +170,13 @@ void ft_ao_bdiv_kernel(double *out, RysIntEnvVars envs, int nGv, double *Gv) goutR[n] += xyR * zR - xyI * zI; goutI[n] += xyR * zI + xyI * zR; } - } } - if (valid && Gv_id < nGv) { + if (Gv_id < nGv) { size_t stride = (size_t)nGv * OF_COMPLEX; - double *aft_tensor = out + ((size_t)envs.ao_loc[sh_id_clamped] * nGv + Gv_id) * OF_COMPLEX; + double *aft_tensor = out + ((size_t)envs.ao_loc[sh_id] * nGv + Gv_id) * OF_COMPLEX; #pragma unroll - for (int n = 0; n < aux_nf; ++n) { + for (int n = 0; n < AUXNF; ++n) { if (n >= nfi) break; aft_tensor[n*stride ] = goutR[n]; aft_tensor[n*stride+1] = goutI[n]; @@ -1172,12 +1137,12 @@ while (1) { } extern "C" { -int build_ft_ao(double *out, RysIntEnvVars *envs, int ngrids, double *grids, int nbas) +int build_ft_ao(double *out, RysIntEnvVars *envs, int ngrids, double *grids) { int nsh_per_block = FT_AO_THREADS/NG_PER_BLOCK; dim3 threads(NG_PER_BLOCK, nsh_per_block); int nbatches_grids = (ngrids + NG_PER_BLOCK - 1) / NG_PER_BLOCK; - int nbatches_shls = (nbas + nsh_per_block - 1) / nsh_per_block; + int nbatches_shls = (envs->nbas + nsh_per_block - 1) / nsh_per_block; dim3 blocks(nbatches_grids, nbatches_shls); ft_ao_bdiv_kernel<<>>(out, *envs, ngrids, grids); cudaError_t err = cudaGetLastError();