Fix OneDFT gradient task ordering and MPI distribution

[DCM-Uni-Paderborn]

Summary

• preserve the association between OneDFT grid-point outputs and GauXC tasks before workload reordering
• implement distributed OneDFT nuclear gradients for host and device execution
• add a finite-difference regression with QZVPP g functions and host/device coverage

Problem

OneDFT assembles primitive fields in atom/task order and returns the exchange-correlation energy density in the same point order. The host gradient path subsequently sorted tasks by workload and then only restored their parent-atom order. That does not recover the original order of multiple tasks belonging to the same atom. Consequently, energy-density values could be contracted with partition-weight derivatives at different quadrature points. Energies and potentials were unaffected, while nuclear gradients could be strongly wrong and asymmetric. High-angular-momentum bases exposed the problem particularly clearly because screening produces more heterogeneous task sizes.

The existing device ordering correction in #215 does not fix this host-side intra-atom mismatch and distributed OneDFT gradients were still disabled.

Implementation

The gradient path now attaches every model energy-density block to its task before any workload sorting. For MPI execution, model inputs are gathered and arranged atom by atom as before; model energy densities and point-coordinate gradients are transformed back to rank order and scattered to the corresponding local tasks. Rank-local Pulay, quadrature-weight, point-motion, and direct atomic-coordinate derivatives are then summed collectively. The device path additionally gathers and reorders the unpartitioned quadrature weights needed by the model. Explicit layout checks detect incomplete or inconsistent point consumption.

Validation

• host serial OneDFT tests: 23 assertions in 4 test cases
• host MPI (4 ranks): 29 assertions in 4 test cases on every rank
• H100 serial host/device gradient tests: 29 assertions in 2 test cases
• H100 MPI (2 ranks): finite-difference test and both host/device sections pass on every rank
• H2O/QZVPP high-l derivative: analytical 0.32143510 versus central finite difference 0.3214351157 a.u.
• downstream CP2K H2O/GAPW-AE/QZVPP checks after removing its high-l fallback:
  • PBE-model force residual: about 1e-8 a.u.
  • molecular XC virial residual: 4.16e-9 a.u.
  • Skala-1.1 force residual with two MPI ranks: 4.28e-6 a.u.

[DCM-Uni-Paderborn]

Dear @wavefunction91,

Could you please take a look at this OneDFT gradient fix when you have a chance, or suggest the appropriate reviewer? It preserves the task-to-grid-output association before workload sorting and adds the distributed gradient path. Host/device and MPI finite-difference validation is documented in the PR description.

A review of the task ordering and gather/scatter logic would be particularly valuable. Please let me know if further validation or changes are needed before this can be considered for merge.

Many thanks!
Greetings,
Thomas

[wavefunction91]

Hi @DCM-Uni-Paderborn ,

The appropriate reviewer for Skala functionality is @awvwgk or someone on that team.

[awvwgk]

Thanks, looks good to me.