Skip to content

Add GETAWAY 3D descriptors - #358

Merged
scal444 merged 8 commits into
NVIDIA-BioNeMo:mainfrom
scal444:getaway-descriptors
Oct 7, 2026
Merged

scal444 merged 8 commits into
NVIDIA-BioNeMo:mainfrom
scal444:getaway-descriptors

Conversation

@scal444

@scal444 scal444 commented Oct 6, 2026

Copy link
Copy Markdown
Collaborator

No description provided.

groupAllReduceSum shuffled with __activemask() as its mask, which does not
guarantee the group's lanes are converged; it now reduces with
cooperative_groups::reduce over an explicit thread_block_tile, so every
lane of the tile must participate. USR's closest/farthest-atom search uses
the same reduction with its tie-breaking comparison. The pairwise kernel's
combine across a warp's four groups (lanes with equal lane % kGroupSize),
which no tile covers, moves into one helper, sumAcrossGroups.

All 3D properties are bit-identical to before in both precisions on 484
ChEMBL + planar conformers, compute-sanitizer synccheck reports no errors,
and kernel times are unchanged.
RDKit's 273-value GETAWAY vector (CalcGETAWAY order) in a new Getaway
family, one warp per conformer:
- Leverages from H = X (X^T X)^+ X^T with RDKit's 1e-3 pseudo-inverse
  cutoff, in float64 (they are rounded and clustered for ITH/ISH).
- Pair loop over the atom-pair triangle with a bond-distance search capped
  at 8 (shared searchBondDepths, now templated on its cap): H, HATS, R and
  R+ per lag 1-8 and 7 channels; farther or disconnected pairs enter the
  totals only, as in RDKit.
- REIG by power iteration on the nonnegative symmetric R (its Perron
  eigenvalue); the first product gives row sums for RARS and RCON.
- HIC uses the PBF of the molecule's default conformer, HGM a sequential
  leverage product, and ITH/ISH a literal port of RDKit's %.*g rounding
  and digit-string clustering. GetawayOptions.precision (default 2) is
  RDKit's precision argument, limited to 1-6.

FULL matches RDKit exactly on all 120,000+ values of 439 single-fragment
ChEMBL + planar conformers; SINGLE is within one 0.001 rounding unit.

RDKit's lag loop runs to its 1e8 distance for disconnected atoms, so it
does not finish on salts; the benchmark's RDKit GETAWAY reference skips
multi-fragment molecules and reports how many.

24,200 conformers, SINGLE: all properties 333 -> 531 ms (GETAWAY kernel
194 ms); requests without GETAWAY are unchanged. Serial RDKit GETAWAY is
~26 ms per conformer.
The group maximum reduces with cooperative_groups::reduce over the group's
tile instead of __activemask() shuffles; the cross-group sums and maximum
use sumAcrossGroups and a new maxAcrossGroups beside it. GETAWAY output is
unchanged (FULL still matches RDKit exactly) and synccheck is clean.
ITH/ISH ran on one lane per conformer, rebuilding decimal text for every
IsClose2 comparison during an insertion sort and greedy scan. The lanes now
round each heavy-atom leverage and reduce its printed text to the integers
IsClose2 reads (leading zeros, length, leading integer, whether it is all
digits), rank the heavy atoms in parallel, and lane 0 runs only the greedy
scan on those keys.

GETAWAY kernel on 24,200 conformers, SINGLE: 192 -> 150 ms. FULL still
matches RDKit exactly, including ITH and ISH.
Bond distances depend only on the molecule, but both kernels searched them
per conformer. A per-molecule table (one byte per atom pair, depth 10) is
now built once per batch and read by both; batches with molecules above
3072 atoms, or past a 256 MB budget, keep the per-row searches.

Outputs are bit-identical. SINGLE, all properties, 24,200 conformers:
487 -> 409 ms end to end (getaway3DKernel 152 -> 111 ms, pairwise3DKernel
122 -> 62 ms, table build 13 ms). FULL is unchanged within noise.
Compute the lag-0 sums after the pair loop and write the H and R blocks
before REIG, so neither set of accumulators is live through the other's
loop. float32: 168 -> 128 registers (3 -> 4 blocks per SM); float64:
218 -> 186. Outputs are bit-identical.

GETAWAY only, 24,200 conformers: SINGLE kernel 106.7 -> 93.1 ms (end to
end 178.8 -> 164.3 ms); FULL unchanged (2276 -> 2273 ms).
@greptile-apps

greptile-apps Bot commented Oct 6, 2026 •

Copy link
Copy Markdown
Contributor

RetriggerConfidence Score: 5/5

[Medium risk] Adds GETAWAY 3D molecular descriptor computation.

The PR appears safe to merge; all three previous findings are addressed.

What we checked:

  • Smaller tables keep reads safe: Both the writer and reader check that a molecule's full range fits. Molecules that do not fit use a separate bond search.
Summary

Adds GETAWAY with 273 values per conformer, Python and C++ options, and a shared bond-distance table for GETAWAY and AUTOCORR3D.

  • Adds RDKit comparisons and benchmark support.
  • Fixes all three previous findings.
  • No new actionable issues were found in the changes since the last review.

Reviews (2) · Last reviewed commit: "Address PR 358 review comments" · Reviewed by Greptile

Comment thread src/descriptors3d_getaway.cuh
Comment thread nvmolkit/descriptors3d.py Outdated
Comment thread src/descriptors3d_topology.cuh Outdated
Comment thread nvmolkit/descriptors3d.py Outdated

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.

The docstring here currently says the is3D flag is preserved for PBF, but GETAWAY's HIC now uses it as well. It may also be worth mentioning that, while the RDKit-molecule path follows RDKit in using the first conformer's flatness for HIC, device-coordinate inputs use each conformer's own geometry.

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.

Done

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

Changes look good to me.

- ITH and ISH are NaN when a heavy-atom leverage is not finite, instead of
  clustering unwritten sort keys; non-finite values skip significant-digit
  rounding.
- GETAWAY precision is validated and passed only when GETAWAY is
  requested, so unused options are ignored as documented.
- The bond-distance table is sized by the batch's actual atom pairs
  (Property3DDeviceInputs::moleculeAtomPairs) rather than the largest
  molecule times the molecule count.
- Document that GETAWAY's HIC reads is3D and takes its flatness from each
  molecule's first conformer, or each row's own with device coordinates.
@scal444
scal444 merged commit 5db48c4 into NVIDIA-BioNeMo:main Oct 7, 2026
16 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