Skip to content

Add device-side ETKDG conformer pruning - #265

Merged
scal444 merged 6 commits into
NVIDIA-BioNeMo:mainfrom
Matthew-Neba:device-side-pruning
Aug 30, 2026
Merged

Add device-side ETKDG conformer pruning#265
scal444 merged 6 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

Adds GPU-resident, RDKit-compatible ordered conformer pruning to ETKDG device output.

  • Builds heavy-atom and symmetry mappings through RDKit.
  • Performs greedy RMSD selection and coordinate compaction on the target GPU.
  • Preserves molecule ownership and renumbers retained conformers within each molecule.
  • Adds coverage for device pruning, symmetry, heavy-atom comparisons, interleaved results, and metadata.

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 Prepares RDKit atom mappings, launches pruning, and compacts retained coordinates while preserving device-result metadata; no follow-up-eligible issue was identified.
src/conformer_rmsd.cu Adds symmetry-aware aligned RMSD and ordered greedy pruning kernels; no follow-up-eligible issue was identified.
src/etkdg.cpp Enables pruning after device-output collection while leaving the host conformer path unchanged.
tests/test_etkdg_device_output.cu Expands coverage for device pruning behavior, ordering, symmetry, heavy atoms, compaction, and conformer-index renumbering.

Reviews (3): Last reviewed commit: "Merge official/main into pr-265" | 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.

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

Approved pending CI, thanks!

@Matthew-Neba

Copy link
Copy Markdown
Contributor Author

I am faling during the build phase of the CI/CD pipeline, but it seems it is due to a Github issue, A Github action is timing out. Can you please trigger a rerun?

@Matthew-Neba
Matthew-Neba requested a review from scal444 August 28, 2026 00:49
@scal444
scal444 merged commit 0cd9165 into NVIDIA-BioNeMo:main Aug 30, 2026
13 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.

Device output Pruning for ETKDG

2 participants