Repository navigation
Get rid of old multi-kernel DG and ETK energy approach - #378
Conversation
The double-precision DG energy issued a scratch memset, one term-parallel launch per term family and a reduction over the per-term energy buffer; the ETK energy did the same for its six term families. Evaluate both with the existing block-per-molecule energy kernels instead (one launch, one block reduction per molecule), which SINGLE precision and the per-molecule path already use. On sm_89 this takes the DG energy from about 1.18 to 0.90 ms per call and the ETK energy from about 1.05 to 0.71 ms per call on ChEMBL 60-80 batches of 1000 conformers. The block-per-molecule kernel writes each active system's energy and zero for inactive systems; every caller zeroes the outputs first, so results match up to summation order. ETK term subsets other than PLAIN without improper torsions keep the term-parallel kernels. The DG forcefield no longer allocates the per-term energy scratch buffer. The FULL gradients keep the term-parallel kernels: the existing block-per- molecule DG gradient measured 25% slower per call (2.19 vs 1.75 ms).
[Medium impact] The reviewed changes appear safe to merge; no blocking issues were found.
|
There was a problem hiding this comment.
The docstring still says both energyOuts and energyBuffer must be zeroed. With the new ALL/PLAIN fast path, energyOuts is overwritten directly and energyBuffer is no longer used, while the other terms still require zero initialization. May consider updating the documentation to clarify these two behaviors.
|
It looks like the old per-term DG energy kernels and related energy-buffer indexing logic are no longer used in production, but are still referenced by |
evasnow1992
left a comment
There was a problem hiding this comment.
Two comments, none blocker. Changes look good to me.
The DG energy path now runs only the block-per-molecule kernel, so drop the per-term DG energy launchers, the DG energy buffer and its index arrays, the unused DG allocateIntermediateBuffers and computeEnergyBlockPerMol, and the host argument setupDeviceBuffers no longer reads. DG per-term energy tests now go through computeEnergy, isolating terms with the chiral and fourth-dimension weights. Update the computeEnergyETK docstring for the overwrite and accumulate paths.
The block-per-mol approach that the e2e per_mol backend has is better, even for large mols that don't go through the full per_mol backend.