diff --git a/gpu4pyscf/df/grad/rhf.py b/gpu4pyscf/df/grad/rhf.py index 6210d4d7e..890e9b64b 100644 --- a/gpu4pyscf/df/grad/rhf.py +++ b/gpu4pyscf/df/grad/rhf.py @@ -27,7 +27,7 @@ SHM_SIZE, LMAX, L_AUX_MAX, THREADS, libvhf_rys, Int3c2eOpt, int2c2e) from gpu4pyscf.df import df from gpu4pyscf.df.df_jk import factorize_dm, _DFHF -from gpu4pyscf.gto.mole import SortedMole +from gpu4pyscf.gto.mole import SortedMole, _split_l_ctr_groups from gpu4pyscf.scf.jk import _check_rsh_factors __all__ = ['Gradients'] @@ -147,8 +147,10 @@ def _jk_energy_per_atom(int3c2e_opt, dm, j_factor=1, k_factor=1, hermi=0, laux = auxmol.uniq_l_ctr[:,0].max() shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() + pair_batch_size = np.minimum(nsp_per_block[0]*4, + np.maximum(nsp_per_block[0]//4, 16)) bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( - int3c2e_opt.bas_ij_cache, nsp_per_block[0]*4) + int3c2e_opt.bas_ij_cache, pair_batch_size) ao_pair_loc = get_ao_pair_loc(mol.uniq_l_ctr[:,0], int3c2e_opt.bas_ij_cache) aux_loc = auxmol.ao_loc @@ -282,7 +284,10 @@ def _j_energy_per_atom(int3c2e_opt, dm, hermi=0, auxbasis_response=True, verbose shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( int3c2e_opt.bas_ij_cache, nsp_per_block[0]*16) - ksh_offsets_cpu = np.append(0, np.cumsum(auxmol.l_ctr_counts)) + # Split auxbasis into small batches for load balance + _, aux_counts = _split_l_ctr_groups( + auxmol.uniq_l_ctr, auxmol.l_ctr_counts, [16]*(laux+1)) + ksh_offsets_cpu = np.append(0, np.cumsum(aux_counts)) ksh_offsets_gpu = cp.asarray(ksh_offsets_cpu+mol.nbas, dtype=np.int32) int3c2e_envs = int3c2e_opt.int3c2e_envs diff --git a/gpu4pyscf/df/grad/tdrhf.py b/gpu4pyscf/df/grad/tdrhf.py index 1bbf11fec..f77d8e6e7 100644 --- a/gpu4pyscf/df/grad/tdrhf.py +++ b/gpu4pyscf/df/grad/tdrhf.py @@ -27,7 +27,7 @@ from gpu4pyscf.df import df from gpu4pyscf.tdscf import rhf as tdrhf from gpu4pyscf.grad import tdrhf as tdrhf_grad -from gpu4pyscf.gto.mole import SortedMole +from gpu4pyscf.gto.mole import SortedMole, _split_l_ctr_groups from gpu4pyscf.scf.jk import _check_rsh_factors __all__ = ['Gradients'] @@ -132,8 +132,10 @@ def _jk_energy_per_atom(int3c2e_opt, dms, j_factor=None, k_factor=None, hermi=0, laux = auxmol.uniq_l_ctr[:,0].max() shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() + pair_batch_size = np.minimum(nsp_per_block[0]*4, + np.maximum(nsp_per_block[0]//4, 16)) bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( - int3c2e_opt.bas_ij_cache, nsp_per_block[0]*4) + int3c2e_opt.bas_ij_cache, pair_batch_size) ao_pair_loc = get_ao_pair_loc(mol.uniq_l_ctr[:,0], int3c2e_opt.bas_ij_cache) aux_loc = auxmol.ao_loc @@ -268,7 +270,10 @@ def _j_energy_per_atom(int3c2e_opt, dms, j_factor, hermi=0, verbose=None): shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( int3c2e_opt.bas_ij_cache, nsp_per_block[0]*16) - ksh_offsets_cpu = np.append(0, np.cumsum(auxmol.l_ctr_counts)) + # Split auxbasis into small batches for load balance + _, aux_counts = _split_l_ctr_groups( + auxmol.uniq_l_ctr, auxmol.l_ctr_counts, [16]*(laux+1)) + ksh_offsets_cpu = np.append(0, np.cumsum(aux_counts)) ksh_offsets_gpu = cp.asarray(ksh_offsets_cpu+mol.nbas, dtype=np.int32) int3c2e_envs = int3c2e_opt.int3c2e_envs @@ -542,8 +547,10 @@ def _jk_energies_by_dm_factors(int3c2e_opt, dm_factors, j_factor, k_factor, laux = auxmol.uniq_l_ctr[:,0].max() shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() + pair_batch_size = np.minimum(nsp_per_block[0]*4, + np.maximum(nsp_per_block[0]//4, 16)) bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( - int3c2e_opt.bas_ij_cache, nsp_per_block[0]*4) + int3c2e_opt.bas_ij_cache, pair_batch_size) ao_pair_loc = get_ao_pair_loc(mol.uniq_l_ctr[:,0], int3c2e_opt.bas_ij_cache) aux_loc = auxmol.ao_loc @@ -691,7 +698,10 @@ def _j_energies_per_atom(int3c2e_opt, dm_pairs, j_factor, shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( int3c2e_opt.bas_ij_cache, nsp_per_block[0]*16) - ksh_offsets_cpu = np.append(0, np.cumsum(auxmol.l_ctr_counts)) + # Split auxbasis into small batches for load balance + _, aux_counts = _split_l_ctr_groups( + auxmol.uniq_l_ctr, auxmol.l_ctr_counts, [16]*(laux+1)) + ksh_offsets_cpu = np.append(0, np.cumsum(aux_counts)) ksh_offsets_gpu = cp.asarray(ksh_offsets_cpu+mol.nbas, dtype=np.int32) int3c2e_envs = int3c2e_opt.int3c2e_envs diff --git a/gpu4pyscf/df/grad/uhf.py b/gpu4pyscf/df/grad/uhf.py index 67da9e48f..af0ad29a4 100644 --- a/gpu4pyscf/df/grad/uhf.py +++ b/gpu4pyscf/df/grad/uhf.py @@ -124,8 +124,10 @@ def _jk_energy_per_atom(int3c2e_opt, dm, j_factor=1, k_factor=1, hermi=0, laux = auxmol.uniq_l_ctr[:,0].max() shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() + pair_batch_size = np.minimum(nsp_per_block[0]*4, + np.maximum(nsp_per_block[0]//4, 16)) bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( - int3c2e_opt.bas_ij_cache, nsp_per_block[0]*4) + int3c2e_opt.bas_ij_cache, pair_batch_size) ao_pair_loc = get_ao_pair_loc(mol.uniq_l_ctr[:,0], int3c2e_opt.bas_ij_cache) aux_loc = auxmol.ao_loc diff --git a/gpu4pyscf/df/hessian/rhf.py b/gpu4pyscf/df/hessian/rhf.py index 091f550e8..0dcb64069 100644 --- a/gpu4pyscf/df/hessian/rhf.py +++ b/gpu4pyscf/df/hessian/rhf.py @@ -40,6 +40,7 @@ from gpu4pyscf.hessian import rhf as rhf_hess from gpu4pyscf.hessian.rhf import _hcore_energy, _aggregate_to_atoms from gpu4pyscf.lib import multi_gpu +from gpu4pyscf.gto.mole import _split_l_ctr_groups num_devices = multi_gpu.num_devices @@ -188,8 +189,11 @@ def proc(aux_batch_iter): laux = auxmol.uniq_l_ctr[:,0].max() shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() + # Bound serial work in each derivative-integral task. + pair_batch_size = np.minimum(nsp_per_block[0]*4, + np.maximum(nsp_per_block[0]//4, 16)) bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( - int3c2e_opt.bas_ij_cache, nsp_per_block[0]*4) + int3c2e_opt.bas_ij_cache, pair_batch_size) ao_pair_loc = get_ao_pair_loc(mol.uniq_l_ctr[:,0], int3c2e_opt.bas_ij_cache) l_ctr_aux_offsets = np.append(0, np.cumsum(auxmol.l_ctr_counts)) @@ -586,10 +590,16 @@ def proc(): laux = auxmol.uniq_l_ctr[:,0].max() shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() + # Bound serial work in each derivative-integral task. + pair_batch_size = np.minimum(nsp_per_block[0]*4, + np.maximum(nsp_per_block[0]//4, 16)) bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( - int3c2e_opt.bas_ij_cache, nsp_per_block[0]*4) + int3c2e_opt.bas_ij_cache, pair_batch_size) ao_pair_loc = get_ao_pair_loc(mol.uniq_l_ctr[:,0], int3c2e_opt.bas_ij_cache) - ksh_offsets_cpu = np.append(0, np.cumsum(auxmol.l_ctr_counts)) + # Split auxbasis into small batches for load balance + _, aux_counts = _split_l_ctr_groups( + auxmol.uniq_l_ctr, auxmol.l_ctr_counts, [16]*(laux+1)) + ksh_offsets_cpu = np.append(0, np.cumsum(aux_counts)) ksh_offsets_gpu = cp.asarray(ksh_offsets_cpu+mol.nbas, dtype=np.int32) int3c2e_envs = int3c2e_opt.int3c2e_envs diff --git a/gpu4pyscf/df/hessian/uhf.py b/gpu4pyscf/df/hessian/uhf.py index 593407991..f31960888 100644 --- a/gpu4pyscf/df/hessian/uhf.py +++ b/gpu4pyscf/df/hessian/uhf.py @@ -163,8 +163,11 @@ def proc(): laux = auxmol.uniq_l_ctr[:,0].max() shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() + # Bound serial work in each derivative-integral task. + pair_batch_size = np.minimum(nsp_per_block[0]*4, + np.maximum(nsp_per_block[0]//4, 16)) bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( - int3c2e_opt.bas_ij_cache, nsp_per_block[0]*4) + int3c2e_opt.bas_ij_cache, pair_batch_size) ao_pair_loc = get_ao_pair_loc(mol.uniq_l_ctr[:,0], int3c2e_opt.bas_ij_cache) l_ctr_aux_offsets = np.append(0, np.cumsum(auxmol.l_ctr_counts)) diff --git a/gpu4pyscf/df/int3c2e_bdiv.py b/gpu4pyscf/df/int3c2e_bdiv.py index 7832375ef..4e0cbfb10 100644 --- a/gpu4pyscf/df/int3c2e_bdiv.py +++ b/gpu4pyscf/df/int3c2e_bdiv.py @@ -34,7 +34,7 @@ from gpu4pyscf.lib import multi_gpu from gpu4pyscf.gto.mole import ( PTR_BAS_COORD, SortedMole, RysIntEnvVars, PBCIntEnvVars, - extract_pgto_params, groupby) + extract_pgto_params, groupby, _split_l_ctr_groups) from gpu4pyscf.scf.jk import ( _nearest_power2, _scale_sp_ctr_coeff, _cache_q_cond_and_non0pairs, _check_rsh_factors, SHM_SIZE, libvhf_rys) @@ -175,8 +175,11 @@ def int3c2e_evaluator(self, ao_pair_batch_size=None, aux_batch_size=None, bas_ij_cache = {k: v for k, v in mol.iter_pair_by_l(self.bas_ij_cache)} else: bas_ij_cache = self.bas_ij_cache + # Bound pair work alongside the auxiliary-shell task subdivision. + pair_batch_size = np.minimum(nsp_per_block[0]*4, + np.maximum(nsp_per_block[0]//4, 16)) bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( - bas_ij_cache, nsp_per_block[0]*4) + bas_ij_cache, pair_batch_size) ao_pair_loc = get_ao_pair_loc(mol.uniq_l_ctr[:,0], bas_ij_cache, cart) nao_pair = ao_pair_loc[-1].get() @@ -210,6 +213,10 @@ def int3c2e_evaluator(self, ao_pair_batch_size=None, aux_batch_size=None, uniq_l_ctr_aux = auxmol.uniq_l_ctr aux_loc = auxmol.ao_loc if aux_batch_size is None: + # Follow df.j_engine_3c2e.Int3c2eOpt.contract_auxvec's 16-shell tasks. + uniq_l_ctr_aux, aux_counts = _split_l_ctr_groups( + auxmol.uniq_l_ctr, auxmol.l_ctr_counts, [16]*(laux+1)) + l_ctr_aux_offsets = np.append(0, np.cumsum(aux_counts)) ksh_offsets_cpu = l_ctr_aux_offsets aux_splits = [0, len(ksh_offsets_cpu)-1] else: @@ -410,7 +417,10 @@ def contract_auxvec(self, auxvec): self.bas_ij_cache, nsp_per_block[0]*4) gout_stride = cp.asarray(gout_stride, dtype=np.int32) - l_ctr_aux_offsets = np.append(0, np.cumsum(auxmol.l_ctr_counts)) + # Split auxbasis into small batches for load balance + _, aux_counts = _split_l_ctr_groups( + auxmol.uniq_l_ctr, auxmol.l_ctr_counts, [16]*(laux+1)) + l_ctr_aux_offsets = np.append(0, np.cumsum(aux_counts)) ksh_offsets = cp.asarray(l_ctr_aux_offsets + mol.nbas, dtype=np.int32) log.debug1('sp_blocks = %d, ksh_blocks = %d, shm_size = %d B', len(shl_pair_offsets)-1, len(ksh_offsets)-1, shm_size_max) @@ -785,8 +795,11 @@ def _int3c2e_ip1_evaluator(int3c2e_opt, scheme, batch_size, laux = auxmol.uniq_l_ctr[:,0].max() shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() + # Bound serial work in each derivative-integral task. + pair_batch_size = np.minimum(nsp_per_block[0]*4, + np.maximum(nsp_per_block[0]//4, 16)) bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( - int3c2e_opt.bas_ij_cache, nsp_per_block[0]*4) + int3c2e_opt.bas_ij_cache, pair_batch_size) ao_pair_loc = get_ao_pair_loc(mol.uniq_l_ctr[:,0], int3c2e_opt.bas_ij_cache) nao_pair = int(ao_pair_loc[-1].get())