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.