Skip to content

Add device-side ETKDG conformer pruning - #265

Open
Matthew-Neba wants to merge 5 commits into
NVIDIA-BioNeMo:mainfrom
Matthew-Neba:device-side-pruning
Open

Add device-side ETKDG conformer pruning#265
Matthew-Neba wants to merge 5 commits into
NVIDIA-BioNeMo:mainfrom
Matthew-Neba:device-side-pruning

Conversation

@Matthew-Neba

@Matthew-Neba Matthew-Neba commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Adds conformer pruning for ETKDG device output. Fixes Device output Pruning for ETKDG #160.
  • Keeps conformer coordinates on the target GPU instead of copying them to the CPU for pruning.
  • Matches RDKit's ordered greedy pruning behavior, including heavy-atom and symmetry-aware RMSD comparisons.

Signed-off-by: Matthew Neba <mattneba4343@gmail.com>
Signed-off-by: Matthew Neba <mattneba4343@gmail.com>
Signed-off-by: Matthew Neba <mattneba4343@gmail.com>
Signed-off-by: Matthew Neba <mattneba4343@gmail.com>
@greptile-apps

greptile-apps Bot commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

Device-resident ETKDG results now support ordered RMSD conformer pruning without copying coordinates to the host.

  • Adds GPU kernels for symmetry-aware and heavy-atom-aware conformer comparison and coordinate compaction.
  • Integrates pruning into the ETKDG device-output path and updates native linkage.
  • Adds device-output tests covering greedy ordering, symmetry mappings, heavy-atom selection, compaction, and metadata renumbering.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename Overview
src/conformer/device_conformer_pruning.cpp Builds RDKit-compatible atom maps, launches device pruning, and compacts retained coordinates and metadata.
src/conformer_rmsd.cu Adds warp-level aligned RMSD, ordered greedy pruning, and retained-coordinate compaction kernels.
src/etkdg.cpp Enables conformer pruning after device results are consolidated onto the target GPU.
src/CMakeLists.txt Adds the device-pruning implementation and RMSD library to the ETKDG target.
tests/test_etkdg_device_output.cu Expands device-output coverage for pruning behavior, symmetry, heavy atoms, interleaving, compaction, and conformer indices.

Reviews (2): Last reviewed commit: "Minor fixes + in place retained list" | Re-trigger Greptile

@scal444

scal444 commented Aug 19, 2026

Copy link
Copy Markdown
Collaborator

Thanks for closing another feature gap! I'll be able to review this later today or tomorrow. Don't worry about the python build failures, I need to fix our CI from our v0.6 version bump. C++ build is the important one to look at.

@scal444 scal444 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.

Looks great! A few small things, and one performance suggestion that can definitely be left to a followup, or handled here.

Comment thread src/etkdg.h Outdated
Comment thread src/conformer_rmsd.cu Outdated
Comment thread src/conformer_rmsd.cu Outdated
Comment thread src/conformer_rmsd.cu Outdated
Comment thread src/conformer_rmsd.cu Outdated
@Matthew-Neba

Copy link
Copy Markdown
Contributor Author

Awesome! Thanks for taking a look, will fix those issues and take another deep look before the weekend

@scal444

scal444 commented Aug 20, 2026

Copy link
Copy Markdown
Collaborator

Also our CI should be fixed now so if you rebase onto or merge in main, those unrelated failures should go away

Signed-off-by: Matthew Neba <mattneba4343@gmail.com>
@Matthew-Neba

Copy link
Copy Markdown
Contributor Author

Applied the fixes. Also applied your suggestion to use a list of the retained conformers so that we don't waste time on those no ops; although it was done in place with the list containing conformers ( groupedConfIds) + a pointer (retainedCount) instead of allocating another buffer for the list.

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.

Device output Pruning for ETKDG

2 participants