Improve DF task granularity for large fitting bases (#970)

* df: bound CUDA tasks for integral builds, gradients and Hessians

Before this change, the three-center builders and derivative
contractions allow up to 4B orbital pairs per task, where B is the
kernel capacity with s-type auxiliary shells. Each task also spans an
auxiliary-shell group. Large groups and expensive angular channels
produce long indivisible tasks, leaving workers idle as a launch
finishes.

This commit changes pair limits to min(4B, max(B//4, 16)) in CDERI
builds, the shared first-derivative evaluator, restricted/unrestricted
ground-state and TD J/K gradients, and restricted/unrestricted DF
Hessian energy contractions. The smaller pair ranges expose more
independently schedulable work. Each path uses its own kernel capacity;
pair ordering and packed auxiliary layouts remain unchanged.

Whole-auxiliary CDERI builds previously assign an entire homogeneous
auxiliary group to each task, potentially spanning hundreds of shells.
This commit also subdivides those groups into at most 16 shells when
aux_batch_size is None, following
df.j_engine_3c2e.Int3c2eOpt.contract_auxvec. The subdivisions are
carried through auxiliary sorting and multi-device clone metadata.
Explicit auxiliary batches and derivative auxiliary grouping retain
their existing layouts.

B/4 is a heuristic balancing parallel work against task setup. For
B=128, 32 pairs times 16 auxiliary shells give 512 shell triples.
Kernel capacities of 8 to 128 process that task in 64 to 4 passes.
Pair subdivision also shortens derivative tasks that retain whole
auxiliary groups. Further subdivision increases queue claims and may
increase repeated setup and partially filled passes.

The 16-pair floor is a separate heuristic limiting small-task setup
cost, with the original 4B cap preserved when it is below 16. For the
g-g pair in aug-cc-pVQZ, B=32: the floor raises B/4 from 8 to 16 pairs.
Sweeps of floors 4, 8, 16 and 32 show low sensitivity in complete
builds and gradients, supporting 16 as an intermediate choice. The
16-shell auxiliary limit follows the existing direct-J implementation.

* df: subdivide whole auxiliary groups in Coulomb-only contractions

Before this change, the Coulomb-only derivative and Rys potential
contractions assign each CUDA block an orbital-pair range and a whole
homogeneous auxiliary-shell group. Large fitting bases lengthen the
serial auxiliary loop inside each block. A long auxiliary group stays
assigned to one block even when other SMs become idle. Subdividing that
group allows its work to run across multiple SMs.

This commit splits each auxiliary group into ranges of at most 16 shells
with _split_l_ctr_groups, following Int3c2eOpt.contract_auxvec in
df/j_engine_3c2e.py. It applies this subdivision to ground-state and TD
Coulomb gradients, pure-J Hessians and multi-density direct SCF J
builds.

Author: zc62

gpu4pyscf/df/grad/rhf.py

@@ -28,7 +28,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']
@@ -149,8 +149,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
 
@@ -286,7 +288,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

gpu4pyscf/df/grad/tdrhf.py

@@ -29,7 +29,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']
@@ -135,8 +135,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
 
@@ -271,7 +273,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
@@ -562,8 +567,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
 
@@ -710,7 +717,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

gpu4pyscf/df/grad/uhf.py

@@ -126,8 +126,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
 

gpu4pyscf/df/hessian/rhf.py

@@ -41,6 +41,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
 
@@ -189,8 +190,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))
@@ -594,10 +598,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

gpu4pyscf/df/hessian/uhf.py

@@ -164,8 +164,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))

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())