Skip to content

Add all-in-one ETKDG force field kernels - #375

Merged
scal444 merged 2 commits into
NVIDIA-BioNeMo:mainfrom
scal444:feat/aio-etkdg-kernels
Oct 23, 2026
Merged

scal444 merged 2 commits into
NVIDIA-BioNeMo:mainfrom
scal444:feat/aio-etkdg-kernels

Conversation

@scal444

@scal444 scal444 commented Oct 22, 2026

Copy link
Copy Markdown
Collaborator

Batched force field for RDKit's all-in-one ETKDG refinement (useLegacyImplementation=false): the 4D distance-geometry terms plus harmonic 1-2/1-3 distances, linear angles, experimental torsions and planarity, evaluated by block-per-molecule kernels in single and double precision like the other force fields. The all-in-one terms reuse the DG/ETK device helpers; only the planarity term is new. The DG gradient kernel's shared-memory accumulation is shared with the new kernel.

Batched force field for RDKit's all-in-one ETKDG refinement
(useLegacyImplementation=false): the 4D distance-geometry terms plus
harmonic 1-2/1-3 distances, linear angles, experimental torsions and
planarity, evaluated by block-per-molecule kernels in single and double
precision like the other force fields. The all-in-one terms reuse the
DG/ETK device helpers; only the planarity term is new. The DG gradient
kernel's shared-memory accumulation is shared with the new kernel.

constructAllInOneForceFieldContribs mirrors RDKit's
DistGeom::constructAllInOneForceField plus addTorsionTerms and takes its
force constants from CrystalFFDetails::forceConsts.

Tests compare each term family and the whole force field against RDKit
in both minimization phases and both precisions. RDKit 2026.09's torsion
gradient mis-evaluates the six-fold term (fixed in rdkit/rdkit#9699), so
the torsion gradient reference is a central difference of RDKit's energy.
@scal444

scal444 commented Oct 22, 2026

Copy link
Copy Markdown
Collaborator Author

PR 1/2 for #287

@greptile-apps

greptile-apps Bot commented Oct 22, 2026 •

Copy link
Copy Markdown
Contributor

RetriggerView in GreptileConfidence Score: 5/5

[Medium impact] The latest change appears safe to merge and fixes the previous performance concern.

Summary

Adds batched all-in-one ETKDG energy and gradient kernels in single and double precision.

  • Batched CUDA calls now evaluate RDKit’s all-in-one ETKDG force field for each molecule.
  • Distance-geometry and all-in-one gradients use the same accumulation path.

Reviews (2) · Last reviewed commit: "Update src/forcefields/dist_geom.cu" · Reviewed by Greptile

Comment thread src/forcefields/dist_geom.cu Resolved
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
@scal444
scal444 requested a review from evasnow1992 October 22, 2026 22:22

#include "versions.h"

#if RDKIT_AIO_ETKDG_API

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just wanted to confirm that these new tests aren't currently enabled in our CI pipeline, since our CI uses older RDKit versions. They're only exercised when running tests locally with RDKit 2026.03.4 or later, correct?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's right, but I'll be updating CI shortly

@evasnow1992 evasnow1992 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just one clarification question. Changes look good to me.

@scal444
scal444 merged commit a8294b8 into NVIDIA-BioNeMo:main Oct 23, 2026
14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants