Skip to content

Improve DF task granularity for large fitting bases #970

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zc62 wants to merge 2 commits into
base: master
Choose a base branch
from
+60 −17
Open

Improve DF task granularity for large fitting bases #970

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a5ee202
df: bound CUDA tasks for integral builds, gradients and Hessians
zc62 Oct 16, 2026
7967987
df: subdivide whole auxiliary groups in Coulomb-only contractions
zc62 Oct 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 8 additions & 3 deletions gpu4pyscf/df/grad/rhf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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']
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
20 changes: 15 additions & 5 deletions gpu4pyscf/df/grad/tdrhf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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']
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
4 changes: 3 additions & 1 deletion gpu4pyscf/df/grad/uhf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
16 changes: 13 additions & 3 deletions gpu4pyscf/df/hessian/rhf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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
Expand Down
5 changes: 4 additions & 1 deletion gpu4pyscf/df/hessian/uhf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
21 changes: 17 additions & 4 deletions gpu4pyscf/df/int3c2e_bdiv.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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())

Expand Down
Loading