Repository navigation
Conversation
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.
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Three-center CUDA tasks can contain a large orbital-pair range and a whole homogeneous auxiliary-shell group. With augmented fitting bases, a few long tasks keep running after other workers become idle. This PR subdivides those tasks to distribute the existing integral work more evenly.
Changes
The forward integral kernel uses one persistent worker block per SM. Workers claim tasks from a shared queue; a large indivisible task can extend the tail of a launch. A water10/aug-cc-pVTZ CDERI kernel trace measured a longest task of 186 ms and 13% worker busy time with the original partition, versus 8 ms and 84% with the proposed partition.
B/4 and the 16-pair floor are heuristic task-size choices. The floor becomes active for B below 64; for the g-g pair in aug-cc-pVQZ, B=32 gives a 16-pair range. The auxiliary-shell limit and pair-count floor bound different dimensions. Capacity sweeps show diminishing whole-build improvements beyond B/4 and low sensitivity to pair floors from 4 to 32.
Performance
RTX 5070 Ti (16 GB, 70 SMs), Ryzen 7 9700X, eight OpenMP/OpenBLAS threads, CUDA runtime 13.0, CuPy 13.6.0 with the CuPy contraction backend. Corresponding JKFIT auxiliary bases and level-3 grids are used throughout. Auxiliary-basis response is enabled for gradients and set to 2 for Hessians. Each entry is the median of three synchronized serial complete-operation calls following one warmup. The baseline commit is 9af748a.
All four configurations share the same base commit; the #967 columns include its four commits. Gradient and Hessian timings use identical saved converged orbitals; TD gradients use the same first singlet state. Derivative state preparation is excluded. Complete SCF kernel calls start from identical saved initial densities and include setup and CDERI construction through the standard hybrid-functional path. The shared RHF gradient and DF-Hessian helpers use a uniform 50% available-memory allowance across all four trees for the local CuPy backend; the remaining helpers retain their default allowance.
Among the small def2-SVP controls, the largest median increase for this PR versus the baseline is 1.8% in the PBE ground-state gradient (0.146 to 0.149 s).
Among the small def2-SVP controls, the largest median increase for adding this PR to #967 is 1.2% in the PBE ground-state gradient (0.147 to 0.149 s).
CDERI construction
Complete SCF / PBE0
Ground-State Gradient / PBE
Ground-State Gradient / PBE0
Hessian / PBE
Hessian / PBE0
TDDFT Excitation Energies / PBE
TDDFT Gradient / PBE0
All times are seconds. CDERI construction is a complete build. Pure-functional TDDFT excitation timings cover three singlet roots using direct DF-J response with multiple density sets and zero stored CDERI; TD gradients cover the first root. Other CDERI consumers share the construction change. Default single-density pure-functional SCF remains on the existing MD direct-J path.
AI Usage Disclaimer
OpenAI Codex with GPT-6.1 assisted with implementation, benchmark execution and analysis, and drafting this PR.