From 2ce2764768a82b84e65c5fa9799fce2878b49779 Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Tue, 6 Oct 2026 17:17:21 -0400 Subject: [PATCH 1/8] Use cooperative-groups tile reductions in the 3D descriptor kernels 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. --- src/descriptors3d_kernel.cuh | 23 +++++++++++++++++++---- src/descriptors3d_pairwise.cuh | 16 ++++++---------- src/descriptors3d_usr.cuh | 27 ++++++++++++--------------- 3 files changed, 37 insertions(+), 29 deletions(-) diff --git a/src/descriptors3d_kernel.cuh b/src/descriptors3d_kernel.cuh index c538094f..a837b3bd 100644 --- a/src/descriptors3d_kernel.cuh +++ b/src/descriptors3d_kernel.cuh @@ -4,6 +4,9 @@ #ifndef NVMOLKIT_DESCRIPTORS3D_KERNEL_CUH #define NVMOLKIT_DESCRIPTORS3D_KERNEL_CUH +#include +#include + #include #include @@ -56,11 +59,23 @@ __device__ __forceinline__ ConformerAtoms loadConformer(const DeviceCoordView& c return atoms; } -//! Sum over aligned groups of @p kWidth lanes (a conformer group by default, or the whole warp). +//! Aligned tile of @p kWidth lanes containing the calling thread. +template __device__ __forceinline__ cooperative_groups::thread_block_tile laneTile() { + return cooperative_groups::tiled_partition(cooperative_groups::this_thread_block()); +} + +//! Sum over aligned groups of @p kWidth lanes (a conformer group by default, or the whole warp), returned to +//! every lane. Every lane of the group must call it. template __device__ __forceinline__ Real groupAllReduceSum(Real value) { - const unsigned activeMask = __activemask(); - for (int offset = kWidth / 2; offset > 0; offset >>= 1) { - value += __shfl_xor_sync(activeMask, value, offset); + return cooperative_groups::reduce(laneTile(), value, cooperative_groups::plus()); +} + +//! Sum over the warp's lanes with equal `lane % kGroupSize`, i.e. across its kGroupsPerWarp groups, returned +//! to each of those lanes. Every lane of the warp must call it. No tile covers these strided lanes, so this +//! combines with shuffles across group offsets. +template __device__ __forceinline__ Real sumAcrossGroups(Real value) { + for (int offset = kGroupSize; offset < kWarpSize; offset <<= 1) { + value += __shfl_xor_sync(0xffffffffu, value, offset); } return value; } diff --git a/src/descriptors3d_pairwise.cuh b/src/descriptors3d_pairwise.cuh index 5c3dca28..61c82ef8 100644 --- a/src/descriptors3d_pairwise.cuh +++ b/src/descriptors3d_pairwise.cuh @@ -357,22 +357,18 @@ __global__ void pairwise3DKernel(const DeviceCoordView coordinates, // Combine the pair streams' partial sums; lanes with equal laneInGroup own the same bins. for (int slot = 0; slot < kPairwiseBinsPerLane; ++slot) { for (int channel = 0; channel < kNumPairwiseChannels; ++channel) { - for (int offset = kGroupSize; offset < kWarpSize; offset <<= 1) { - if constexpr (kRdf) { - rdfAcc[slot][channel] += __shfl_xor_sync(0xffffffffu, rdfAcc[slot][channel], offset); - } - if constexpr (kMorse) { - morseAcc[slot][channel] += __shfl_xor_sync(0xffffffffu, morseAcc[slot][channel], offset); - } + if constexpr (kRdf) { + rdfAcc[slot][channel] = sumAcrossGroups(rdfAcc[slot][channel]); + } + if constexpr (kMorse) { + morseAcc[slot][channel] = sumAcrossGroups(morseAcc[slot][channel]); } } } if constexpr (kAutocorr) { for (int slot = 0; slot < kAutocorrLagsPerLane; ++slot) { for (int channel = 0; channel < kNumAutocorrChannels; ++channel) { - for (int offset = kGroupSize; offset < kWarpSize; offset <<= 1) { - autocorrAcc[slot][channel] += __shfl_xor_sync(0xffffffffu, autocorrAcc[slot][channel], offset); - } + autocorrAcc[slot][channel] = sumAcrossGroups(autocorrAcc[slot][channel]); } } } diff --git a/src/descriptors3d_usr.cuh b/src/descriptors3d_usr.cuh index 0a497e87..6ace7fdb 100644 --- a/src/descriptors3d_usr.cuh +++ b/src/descriptors3d_usr.cuh @@ -39,10 +39,7 @@ __device__ __forceinline__ bool isBetter(const double distance, const int atomId * reference point's moments by up to 1.3. The selection reads the float64 coordinates, O(atoms) per call. */ template -__device__ __forceinline__ int extremeAtom(const ConformerAtoms& atoms, - const int laneInGroup, - const unsigned groupMask, - const double3 point) { +__device__ __forceinline__ int extremeAtom(const ConformerAtoms& atoms, const int laneInGroup, const double3 point) { // Starting from atom 0 keeps the result in range when no distance compares (NaN coordinates), which is // also the atom RDKit's scans return then. ExtremeAtom best{kLargest ? -INFINITY : INFINITY, 0}; @@ -55,14 +52,14 @@ __device__ __forceinline__ int extremeAtom(const ConformerAtoms& atoms, best = {distance, atomIdx}; } } - for (int offset = kGroupSize / 2; offset > 0; offset >>= 1) { - const double otherDistance = __shfl_xor_sync(groupMask, best.distance, offset); - const int otherIdx = __shfl_xor_sync(groupMask, best.atomIdx, offset); - if (isBetter(otherDistance, otherIdx, best)) { - best = {otherDistance, otherIdx}; - } - } - return best.atomIdx; + // The preferred atom of any two is unique (ties go to the lower index), so the result is independent of + // the reduction order. + return cooperative_groups::reduce(laneTile(), + best, + [](const ExtremeAtom& a, const ExtremeAtom& b) { + return isBetter(b.distance, b.atomIdx, a) ? b : a; + }) + .atomIdx; } //! Float64 position of atom @p atomIdx. @@ -189,9 +186,9 @@ __global__ void usr3DKernel(const DeviceCoordView coordinates, const uint8_t* classes = kUsrcat ? inputs.usrcatAtomClasses + inputs.moleculeAtomStarts[coordinates.molIndices[conformerIdx]] : nullptr; - const int closest = extremeAtom(atoms, laneInGroup, groupMask, centroid); - const int farthest = extremeAtom(atoms, laneInGroup, groupMask, centroid); - const int farthestFromFarthest = extremeAtom(atoms, laneInGroup, groupMask, atomPosition(atoms, farthest)); + const int closest = extremeAtom(atoms, laneInGroup, centroid); + const int farthest = extremeAtom(atoms, laneInGroup, centroid); + const int farthestFromFarthest = extremeAtom(atoms, laneInGroup, atomPosition(atoms, farthest)); // -1 is the centroid, the origin of the centered coordinates. const int referenceAtoms[kNumUsrReferencePoints] = {-1, closest, farthest, farthestFromFarthest}; From 6f9c8deab359b20436f22a767e0f14b0ff64257b Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Tue, 6 Oct 2026 12:39:50 -0400 Subject: [PATCH 2/8] Add the GETAWAY 3D descriptor 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. --- benchmarks/descriptors3d_bench.py | 15 + docs/api/nvmolkit.rst | 1 + nvmolkit/descriptors3d.cpp | 3 + nvmolkit/descriptors3d.py | 30 +- nvmolkit/tests/test_descriptors3d.py | 96 ++++ src/descriptors3d.cu | 25 +- src/descriptors3d.h | 40 +- src/descriptors3d_getaway.cuh | 768 +++++++++++++++++++++++++++ src/descriptors3d_kernel.cuh | 41 ++ src/descriptors3d_mol.cpp | 47 +- src/descriptors3d_pairwise.cuh | 53 +- tests/test_descriptors3d.cpp | 22 + 12 files changed, 1075 insertions(+), 66 deletions(-) create mode 100644 src/descriptors3d_getaway.cuh diff --git a/benchmarks/descriptors3d_bench.py b/benchmarks/descriptors3d_bench.py index 7f4b939f..768609f9 100644 --- a/benchmarks/descriptors3d_bench.py +++ b/benchmarks/descriptors3d_bench.py @@ -84,6 +84,11 @@ def _calc_rdkit_property(mol: Chem.Mol, conf_id: int, prop: Property3D) -> float return rdMolDescriptors.GetUSR(mol, confId=conf_id) if prop == Property3D.USRCAT: return rdMolDescriptors.GetUSRCAT(mol, confId=conf_id) + if prop == Property3D.GETAWAY: + if len(Chem.GetMolFrags(mol)) > 1: + # RDKit's GETAWAY loops up to its 1e8 disconnected-pair distance and would not finish. + return [math.nan] * 273 + return rdMolDescriptors.CalcGETAWAY(mol, confId=conf_id) return RDKIT_CALCULATORS[prop](mol, confId=conf_id, useAtomicMasses=True) @@ -151,6 +156,11 @@ def _validate( for result in results: for prop in properties: actual = result[prop.value].numpy() + if prop == Property3D.GETAWAY: + # Rows RDKit could not compute (multi-fragment molecules) are NaN and not compared. + computed = ~np.isnan(expected[prop]).all(axis=1) + np.testing.assert_allclose(actual[computed], expected[prop][computed], rtol=0, atol=max(atol, 1.1e-3)) + continue if prop not in (Property3D.USR, Property3D.USRCAT): np.testing.assert_allclose(actual, expected[prop], rtol=rtol, atol=atol) continue @@ -256,6 +266,11 @@ def run( } rdkit_timing = None + if not no_rdkit and Property3D.GETAWAY in properties: + skipped = sum(len(Chem.GetMolFrags(mol)) > 1 for mol in mols) + if skipped: + print(f"RDKit GETAWAY skips {skipped} multi-fragment molecules (it would not finish on them)") + row["rdkit_getaway_skipped_mols"] = skipped if not no_rdkit: rdkit_timing = time_it( lambda mols=mols, properties=properties: _calc_rdkit(mols, properties), diff --git a/docs/api/nvmolkit.rst b/docs/api/nvmolkit.rst index cdb088e1..037ce4a1 100644 --- a/docs/api/nvmolkit.rst +++ b/docs/api/nvmolkit.rst @@ -176,6 +176,7 @@ Conformer RMSD descriptors3d.Property3DOptions descriptors3d.MomentOptions descriptors3d.WhimOptions + descriptors3d.GetawayOptions Torsion Fingerprint Deviation (TFD) ----------------------------------- diff --git a/nvmolkit/descriptors3d.cpp b/nvmolkit/descriptors3d.cpp index edb76a03..f1cf04dd 100644 --- a/nvmolkit/descriptors3d.cpp +++ b/nvmolkit/descriptors3d.cpp @@ -89,6 +89,7 @@ bp::object calc3DProperties(const bp::list& mols, const bp::list& propertyNames, const bool useAtomicMasses, const double whimThreshold, + const unsigned int getawayPrecision, const bp::object& coordinates, const nvMolKit::PrecisionMode& precision, const int preprocessingThreads, @@ -100,6 +101,7 @@ bp::object calc3DProperties(const bp::list& mols, nvMolKit::Property3DOptions options; options.moments.useAtomicMasses = useAtomicMasses; options.whim.threshold = whimThreshold; + options.getaway.precision = getawayPrecision; std::vector properties; for (int i = 0; i < bp::len(propertyNames); ++i) { @@ -131,6 +133,7 @@ BOOST_PYTHON_MODULE(_descriptors3d) { bp::arg("properties"), bp::arg("useAtomicMasses"), bp::arg("whimThreshold"), + bp::arg("getawayPrecision"), bp::arg("coordinates"), bp::arg("precision"), bp::arg("preprocessingThreads"), diff --git a/nvmolkit/descriptors3d.py b/nvmolkit/descriptors3d.py index be4a3deb..987afede 100644 --- a/nvmolkit/descriptors3d.py +++ b/nvmolkit/descriptors3d.py @@ -40,6 +40,10 @@ class Property3D(Enum): Conformers with fewer than 3 atoms, which RDKit rejects, produce NaN. - ``USRCAT``: 60 values, ``USR`` for all atoms and then for RDKit's hydrophobic, aromatic, acceptor and donor atoms. Same atom-count rule. + - ``GETAWAY``: 273 values in RDKit's ``CalcGETAWAY`` order. RDKit counts + atom pairs in different fragments (salts, counterions) only in the + totals; nvMolKit returns the same values, which RDKit itself takes + impractically long to reach for multi-fragment molecules. """ PMI1 = "PMI1" @@ -59,6 +63,7 @@ class Property3D(Enum): AUTOCORR3D = "AUTOCORR3D" USR = "USR" USRCAT = "USRCAT" + GETAWAY = "GETAWAY" @dataclass(frozen=True) @@ -85,6 +90,18 @@ class WhimOptions: threshold: float = 0.001 +@dataclass(frozen=True) +class GetawayOptions: + """Options for :attr:`Property3D.GETAWAY`. + + Attributes: + precision: Significant digits the heavy-atom leverages are rounded to before ``ITH`` and ``ISH`` + group them, RDKit's default ``2``. Must be between 1 and 6. + """ + + precision: int = 2 + + @dataclass(frozen=True) class Property3DOptions: """Per-family options for :func:`Calc3DProperties`; each family reads only its own member. @@ -106,10 +123,12 @@ class Property3DOptions: Attributes: moments: Options for the moment-based properties. whim: Options for ``WHIM``. + getaway: Options for ``GETAWAY``. """ moments: MomentOptions = MomentOptions() whim: WhimOptions = WhimOptions() + getaway: GetawayOptions = GetawayOptions() @dataclass(frozen=True) @@ -314,8 +333,9 @@ def Calc3DProperties( Steps where float32 measurably loses accuracy always compute in float64: coordinate centering (conformers far from the origin), WHIM's PCA (inverse kurtosis of near-planar conformers), - MORSE's first scattering value (a large sum of positive terms) - and USR's reference-atom choice (near-equidistant atoms). + MORSE's first scattering value (a large sum of positive terms), + USR's reference-atom choice (near-equidistant atoms) and GETAWAY's + leverages (rounded and clustered for ITH and ISH). hardwareOptions: Only ``preprocessingThreads`` applies: the CPU threads used to extract coordinates and atom weights from the molecules (default ``-1``, all threads). @@ -379,6 +399,11 @@ def Calc3DProperties( options = Property3DOptions() elif not isinstance(options, Property3DOptions): raise TypeError(f"options must be a Property3DOptions or None, got {type(options).__name__}") + getaway_precision = options.getaway.precision + if isinstance(getaway_precision, bool) or not isinstance(getaway_precision, int): + raise TypeError(f"GETAWAY precision must be an int, got {type(getaway_precision).__name__}") + if not 1 <= getaway_precision <= 6: + raise ValueError(f"GETAWAY precision must be between 1 and 6 significant digits, got {getaway_precision}") if hardwareOptions is None: hardwareOptions = HardwareOptions() elif not isinstance(hardwareOptions, HardwareOptions): @@ -398,6 +423,7 @@ def Calc3DProperties( [prop.value for prop in normalized_properties], options.moments.useAtomicMasses, options.whim.threshold, + options.getaway.precision, coordinate_interfaces, precision, hardwareOptions.preprocessingThreads, diff --git a/nvmolkit/tests/test_descriptors3d.py b/nvmolkit/tests/test_descriptors3d.py index 56b6ebca..76deb69e 100644 --- a/nvmolkit/tests/test_descriptors3d.py +++ b/nvmolkit/tests/test_descriptors3d.py @@ -14,6 +14,7 @@ from nvmolkit.descriptors3d import ( Calc3DProperties, Device3DPropertyResult, + GetawayOptions, MomentOptions, Property3D, Property3DOptions, @@ -30,6 +31,7 @@ Property3D.AUTOCORR3D, Property3D.USR, Property3D.USRCAT, + Property3D.GETAWAY, ) SCALAR_PROPERTIES = tuple(prop for prop in Property3D if prop not in VECTOR_PROPERTIES) PAIRWISE_PROPERTIES = (Property3D.RDF, Property3D.MORSE, Property3D.AUTOCORR3D) @@ -547,6 +549,100 @@ def test_autocorr3d_non_finite_coordinates_match_rdkit(precision): ) +def _rdkit_getaway_rows(mols, precision_digits=2): + return np.asarray( + [ + rdMolDescriptors.CalcGETAWAY(mol, confId=conf.GetId(), precision=precision_digits) + for mol in mols + for conf in mol.GetConformers() + ] + ) + + +@pytest.mark.parametrize("precision", PRECISIONS) +def test_getaway_matches_rdkit(precision): + rng = np.random.default_rng(11) + steps = rng.normal(size=(60, 3)) + chain = np.cumsum(1.5 * steps / np.linalg.norm(steps, axis=1, keepdims=True), axis=0) + mols = [ + _embed("CC(=O)Nc1ccc(O)cc1", 3, 97), + # Planar: X^T X has a zero singular value, dropped from the pseudo-inverse, and HIC uses D = 2. + _embed("c1ccncc1", 2, 101), + _mol_with_conformers("C" * 60, [chain]), + ] + result = Calc3DProperties(mols, Property3D.GETAWAY, precision=precision) + assert result[Property3D.GETAWAY].torch().shape == (6, 273) + _assert_rounded_matches_rdkit(result[Property3D.GETAWAY].numpy(), _rdkit_getaway_rows(mols), precision) + + +@pytest.fixture(scope="module") +def getaway_salt(): + # RDKit's lag loop runs to its 1e8 distance for disconnected atoms: seconds even for this 4-atom salt. + mol = _mol_with_conformers("CCO.N", [[(0.0, 0.0, 0.0), (1.5, 0.1, 0.0), (2.1, 1.4, 0.2), (5.0, 1.0, -1.0)]]) + return mol, _rdkit_getaway_rows([mol]) + + +@pytest.mark.parametrize("precision", PRECISIONS) +def test_getaway_counts_disconnected_pairs_in_totals_like_rdkit(getaway_salt, precision): + mol, expected = getaway_salt + result = Calc3DProperties([mol], Property3D.GETAWAY, precision=precision) + _assert_rounded_matches_rdkit(result[Property3D.GETAWAY].numpy(), expected, precision) + + +@pytest.mark.parametrize("precision", PRECISIONS) +def test_getaway_hic_uses_default_conformer_dimension(precision): + # RDKit's HIC takes its 2D/3D choice from PBF(mol), i.e. the default (first) conformer: planar here. + planar = [(0.0, 0.0, 0.0), (1.5, 0.0, 0.0), (2.1, 1.4, 0.0), (3.6, 1.5, 0.0)] + bent = [(0.0, 0.0, 0.0), (1.5, 0.1, 0.3), (2.1, 1.4, -0.4), (3.6, 1.5, 0.6)] + mol = _mol_with_conformers("CCCO", [planar, bent]) + result = Calc3DProperties([mol], Property3D.GETAWAY, precision=precision) + _assert_rounded_matches_rdkit(result[Property3D.GETAWAY].numpy(), _rdkit_getaway_rows([mol]), precision) + + +@pytest.mark.parametrize("precision", PRECISIONS) +def test_getaway_tiny_leverage_matches_rdkit(precision): + # The central carbon sits 0.003 A from the centroid: its leverage (~2e-6) prints in scientific notation + # in RDKit's ITH/ISH digit-string clustering. + corner = 0.89 + coordinates = [ + (0.003, 0.0, 0.0), + (corner, corner, corner), + (corner, -corner, -corner), + (-corner, corner, -corner), + (-corner, -corner, corner), + ] + mol = _mol_with_conformers("C(C)(C)(C)C", [coordinates]) + result = Calc3DProperties([mol], Property3D.GETAWAY, precision=precision) + _assert_rounded_matches_rdkit(result[Property3D.GETAWAY].numpy(), _rdkit_getaway_rows([mol]), precision) + + +@pytest.mark.parametrize("precision_digits", [1, 3]) +def test_getaway_precision_option_matches_rdkit(precision_digits): + mols = [_embed("CC(=O)Nc1ccc(O)cc1", 2, 103)] + result = Calc3DProperties( + mols, + Property3D.GETAWAY, + options=Property3DOptions(getaway=GetawayOptions(precision=precision_digits)), + precision=PrecisionMode.FULL, + ) + _assert_rounded_matches_rdkit( + result[Property3D.GETAWAY].numpy(), _rdkit_getaway_rows(mols, precision_digits), PrecisionMode.FULL + ) + + +@pytest.mark.parametrize("precision_digits", [-1, 0, 7]) +def test_getaway_rejects_out_of_range_precision(precision_digits): + options = Property3DOptions(getaway=GetawayOptions(precision=precision_digits)) + with pytest.raises(ValueError, match="GETAWAY precision"): + Calc3DProperties([_embed("CCO", 1, 5)], Property3D.GETAWAY, options=options) + + +def test_getaway_rejects_non_integer_precision(): + options = Property3DOptions(getaway=GetawayOptions(precision=2.5)) + with pytest.raises(TypeError, match="GETAWAY precision"): + Calc3DProperties([_embed("CCO", 1, 5)], Property3D.GETAWAY, options=options) + + @pytest.mark.parametrize("precision", PRECISIONS) def test_autocorr3d_accepts_batches_without_bonds(precision): mols = [ diff --git a/src/descriptors3d.cu b/src/descriptors3d.cu index e9bc877e..06958bca 100644 --- a/src/descriptors3d.cu +++ b/src/descriptors3d.cu @@ -7,6 +7,7 @@ #include #include "src/descriptors3d.h" +#include "src/descriptors3d_getaway.cuh" #include "src/descriptors3d_kernel.cuh" #include "src/descriptors3d_moments.cuh" #include "src/descriptors3d_pairwise.cuh" @@ -52,6 +53,8 @@ std::string_view property3DName(const Property3D property) { return "USR"; case Property3D::USRCAT: return "USRCAT"; + case Property3D::GETAWAY: + return "GETAWAY"; } throw std::invalid_argument("Unknown Property3D value " + std::to_string(static_cast(property))); } @@ -77,6 +80,7 @@ using descriptors3d_detail::kGroupSize; using descriptors3d_detail::kGroupsPerWarp; using descriptors3d_detail::kWarpSize; using descriptors3d_detail::kWarpsPerBlock; +using descriptors3d_detail::launchGetawayProperties; using descriptors3d_detail::launchPairwiseProperties; using descriptors3d_detail::launchProjectionProperties; using descriptors3d_detail::launchUsrProperties; @@ -129,6 +133,7 @@ constexpr SharedStageSet directStages(const Property3D property) { case Property3D::AUTOCORR3D: case Property3D::USR: case Property3D::USRCAT: + case Property3D::GETAWAY: return 0; } return 0; @@ -293,8 +298,9 @@ Property3DResults calc3DPropertiesGpu(const DeviceCoordView& coord Real* pbfOutput = nullptr; Real* whimOutput = nullptr; PairwiseOutputs pairwiseOutputs; - Real* usrOutput = nullptr; - Real* usrcatOutput = nullptr; + Real* usrOutput = nullptr; + Real* usrcatOutput = nullptr; + Real* getawayOutput = nullptr; for (const Property3D property : properties) { // Bounds the work arrays, which hold one slot per known property. if (std::find(kAllProperty3D.begin(), kAllProperty3D.end(), property) == kAllProperty3D.end()) { @@ -321,6 +327,9 @@ Property3DResults calc3DPropertiesGpu(const DeviceCoordView& coord case Property3DFamily::Usr: (property == Property3D::USR ? usrOutput : usrcatOutput) = it->second.data(); break; + case Property3DFamily::Getaway: + getawayOutput = it->second.data(); + break; } } @@ -328,6 +337,9 @@ Property3DResults calc3DPropertiesGpu(const DeviceCoordView& coord if (whimOutput != nullptr && (!std::isfinite(options.whim.threshold) || options.whim.threshold < 0)) { throw std::invalid_argument("WHIM threshold must be finite and non-negative"); } + if (getawayOutput != nullptr && (options.getaway.precision < 1 || options.getaway.precision > 6)) { + throw std::invalid_argument("GETAWAY precision must be between 1 and 6 significant digits"); + } if (coordinates.numConformers == 0) { return results; @@ -338,9 +350,13 @@ Property3DResults calc3DPropertiesGpu(const DeviceCoordView& coord } const bool anyPairwise = pairwiseOutputs.rdf != nullptr || pairwiseOutputs.morse != nullptr || pairwiseOutputs.autocorr3D != nullptr; - if ((whimOutput != nullptr || anyPairwise) && inputs.atomPropertyWeights == nullptr) { + if ((whimOutput != nullptr || anyPairwise || getawayOutput != nullptr) && inputs.atomPropertyWeights == nullptr) { throw std::invalid_argument( - "Atom-property weights must not be null when WHIM, RDF, MORSE or AUTOCORR3D is requested"); + "Atom-property weights must not be null when WHIM, RDF, MORSE, AUTOCORR3D or GETAWAY is requested"); + } + if (getawayOutput != nullptr && + (inputs.bondNeighborStarts == nullptr || inputs.bondNeighbors == nullptr || inputs.heavyAtomFlags == nullptr)) { + throw std::invalid_argument("Bond adjacency and heavy-atom flags must not be null when GETAWAY is requested"); } if (pairwiseOutputs.autocorr3D != nullptr && (inputs.covalentRadiusWeights == nullptr || inputs.bondNeighborStarts == nullptr || @@ -362,6 +378,7 @@ Property3DResults calc3DPropertiesGpu(const DeviceCoordView& coord launchProjectionProperties(coordinates, inputs, options.whim, pbfOutput, whimOutput, stream); launchPairwiseProperties(coordinates, inputs, pairwiseOutputs, stream); launchUsrProperties(coordinates, inputs, usrOutput, usrcatOutput, stream); + launchGetawayProperties(coordinates, inputs, options.getaway, getawayOutput, stream); return results; } diff --git a/src/descriptors3d.h b/src/descriptors3d.h index 2705b416..c5ee7197 100644 --- a/src/descriptors3d.h +++ b/src/descriptors3d.h @@ -36,9 +36,10 @@ enum class Property3D : int { AUTOCORR3D = 14, USR = 15, USRCAT = 16, + GETAWAY = 17, }; -inline constexpr std::array kAllProperty3D = { +inline constexpr std::array kAllProperty3D = { Property3D::PMI1, Property3D::PMI2, Property3D::PMI3, @@ -56,6 +57,7 @@ inline constexpr std::array kAllProperty3D = { Property3D::AUTOCORR3D, Property3D::USR, Property3D::USRCAT, + Property3D::GETAWAY, }; inline constexpr int kNumWhimProperties = 114; @@ -64,6 +66,7 @@ inline constexpr int kNumMorseProperties = 224; //!< 7 atom-property chann inline constexpr int kNumAutocorr3DProperties = 80; //!< 8 atom-property channels x 10 topological lags. inline constexpr int kNumUsrProperties = 12; //!< 3 distance moments x 4 reference points. inline constexpr int kNumUsrcatProperties = 60; //!< USR of all atoms, then of each of 4 atom classes. +inline constexpr int kNumGetawayProperties = 273; //!< RDKit's GETAWAY vector (CalcGETAWAY order). //! Canonical name of @p property, e.g. "PMI1" or "RadiusOfGyration". std::string_view property3DName(Property3D property); @@ -86,6 +89,8 @@ constexpr int property3DWidth(const Property3D property) { return kNumUsrProperties; case Property3D::USRCAT: return kNumUsrcatProperties; + case Property3D::GETAWAY: + return kNumGetawayProperties; default: return 1; } @@ -98,6 +103,7 @@ enum class Property3DFamily : int { Projection, //!< Coordinate PCA and projections onto its axes: PBF and WHIM. Pairwise, //!< Sums over atom pairs weighted by atom-property pairs: RDF, MORSE and AUTOCORR3D. Usr, //!< Distance moments from four reference points: USR and USRCAT. + Getaway, //!< Leverage (molecular influence) matrix descriptors: GETAWAY. }; constexpr Property3DFamily property3DFamily(const Property3D property) { @@ -123,6 +129,8 @@ constexpr Property3DFamily property3DFamily(const Property3D property) { case Property3D::USR: case Property3D::USRCAT: return Property3DFamily::Usr; + case Property3D::GETAWAY: + return Property3DFamily::Getaway; } return Property3DFamily::Moments; } @@ -140,11 +148,19 @@ struct WhimOptions { double threshold = 0.001; }; +//! Options for GETAWAY. +struct GetawayOptions { + //! Significant digits the heavy-atom leverages are rounded to before ITH and ISH cluster them; RDKit's + //! default. Must be between 1 and 6. + unsigned int precision = 2; +}; + //! Per-family options; each family reads only its own member. PBF, RDF, MORSE, AUTOCORR3D, USR and //! USRCAT have no options. struct Property3DOptions { - MomentOptions moments; - WhimOptions whim; + MomentOptions moments; + WhimOptions whim; + GetawayOptions getaway; }; /** @@ -158,7 +174,7 @@ struct Property3DDeviceInputs { const int32_t* moleculeAtomStarts = nullptr; //! Moments: one weight per atom; null gives every atom unit weight. const double* momentWeights = nullptr; - //! WHIM, RDF, MORSE, AUTOCORR3D: six atom-property channels (RDKit's relative mass, van der Waals volume, + //! WHIM, RDF, MORSE, AUTOCORR3D, GETAWAY: six atom-property channels (RDKit's relative mass, van der Waals volume, //! electronegativity, polarizability and ionization potential, then I-state), channel-major with one //! value per atom. Required when any of them is requested. const double* atomPropertyWeights = nullptr; @@ -167,7 +183,7 @@ struct Property3DDeviceInputs { const double* iStateDragWeights = nullptr; //! AUTOCORR3D: RDKit's relative covalent radius (GetRelativeRcov), one value per atom. const double* covalentRadiusWeights = nullptr; - //! AUTOCORR3D: bond adjacency in CSR form, both required (with `bondNeighbors` non-null even when no + //! AUTOCORR3D, GETAWAY: bond adjacency in CSR form, both required (with `bondNeighbors` non-null even when no //! molecule has a bond). Atom g (indexed like @ref moleculeAtomStarts, `totalAtoms + 1` starts) has //! neighbors `bondNeighbors[bondNeighborStarts[g] .. bondNeighborStarts[g + 1])`, stored as atom indices //! within its molecule. @@ -176,7 +192,12 @@ struct Property3DDeviceInputs { //! USRCAT: per atom, bit c set when the atom is in RDKit's USRCAT class c (hydrophobic, aromatic, //! acceptor, donor). const uint8_t* usrcatAtomClasses = nullptr; - //! PBF: per-conformer RDKit is3D flags (one per coordinate row); null treats every row as 3D. + //! GETAWAY: per atom, 1 for heavy atoms (atomic number above 1), else 0. + const uint8_t* heavyAtomFlags = nullptr; + //! GETAWAY: per molecule, the coordinate row of its default (first) conformer, whose PBF decides HIC's + //! dimension as in RDKit. Null, or a negative entry, uses each row's own conformer. + const int32_t* defaultConformerRows = nullptr; + //! PBF, GETAWAY: per-conformer RDKit is3D flags (one per coordinate row); null treats every row as 3D. const int8_t* conformerIs3D = nullptr; //! WHIM: largest molecule atom count in the batch (host value); sizes the per-conformer symmetry-search //! scratch. Rows with more atoms produce NaN. @@ -193,15 +214,16 @@ template using Property3DResults = std::unordered_map diff --git a/src/descriptors3d_getaway.cuh b/src/descriptors3d_getaway.cuh new file mode 100644 index 00000000..1469c295 --- /dev/null +++ b/src/descriptors3d_getaway.cuh @@ -0,0 +1,768 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#ifndef NVMOLKIT_DESCRIPTORS3D_GETAWAY_CUH +#define NVMOLKIT_DESCRIPTORS3D_GETAWAY_CUH + +#include +#include + +#include "src/descriptors3d.h" +#include "src/descriptors3d_kernel.cuh" +#include "src/descriptors3d_projection.cuh" +#include "src/utils/cuda_error_check.h" +#include "src/utils/device_vector.h" +#include "src/utils/symmetric_eigenvalues_3x3.cuh" + +namespace nvMolKit::descriptors3d_detail { + +//! Unweighted channel plus the atom-property channels, in RDKit's output order (u, m, v, e, p, i, s). +constexpr int kNumGetawayChannels = kNumAtomPropertyChannels + 1; +//! Topological lags with their own H, HATS and R values (1-8); farther or disconnected pairs only enter the +//! totals. One lane of each kGroupSize-lane group owns each lag. +constexpr int kNumGetawayLags = 8; +static_assert(kNumGetawayLags == kGroupSize); +//! Pair streams per warp: each kGroupSize-lane group walks interleaved rows of the atom-pair triangle. +constexpr int kGetawayPairStreams = kWarpSize / kGroupSize; + +//! Values per channel in the H/HATS block: H0..H8, HT, HATS0..HATS8, HATST. +constexpr int kGetawayHBlock = 2 * (kNumGetawayLags + 2); +//! Values per channel in the R block: R1..R8, RT, R1+..R8+, RT+. +constexpr int kGetawayRBlock = 2 * (kNumGetawayLags + 1); +constexpr int kGetawayHStart = 4; +constexpr int kGetawayRconIndex = kGetawayHStart + kNumGetawayChannels * kGetawayHBlock; +constexpr int kGetawayRStart = kGetawayRconIndex + 3; +static_assert(kGetawayRStart + kNumGetawayChannels * kGetawayRBlock == kNumGetawayProperties); + +//! Singular values of X^T X at or below this are dropped from its pseudo-inverse, as in RDKit's GetPinv. +constexpr double kGetawayPinvTolerance = 1e-3; +//! PBF below this makes RDKit treat the molecule as two-dimensional for HIC. +constexpr double kGetawayPlanarPbf = 1e-5; +//! Upper bound on power-iteration steps for REIG. +constexpr int kGetawayMaxIterations = 1000; + +/** + * @brief Per-atom scratch row, written once per conformer as the compute type: scaled projections y (three + * values, H(j, k) = y_j . y_k), leverage h = H(j, j), sqrt(h), the atom-property channels, then the + * centered coordinates. + */ +constexpr int kGetawayRowY = 0; +constexpr int kGetawayRowLeverage = 3; +constexpr int kGetawayRowRootLev = 4; +constexpr int kGetawayRowWeights = 5; +constexpr int kGetawayRowCentered = kGetawayRowWeights + kNumAtomPropertyChannels; +constexpr int kGetawayRowStride = kGetawayRowCentered + 3; + +//! Exact powers of ten representable in float64. +__device__ __forceinline__ double exactPowerOfTen(const int exponent) { + constexpr double kPowers[] = {1e0, 1e1, 1e2, 1e3, 1e4, 1e5, 1e6, 1e7, 1e8, 1e9, 1e10, 1e11, + 1e12, 1e13, 1e14, 1e15, 1e16, 1e17, 1e18, 1e19, 1e20, 1e21, 1e22}; + return kPowers[exponent]; +} + +//! @p value * 10^@p exponent; exact scaling factors within float64's exact powers of ten. +__device__ __forceinline__ double scaleByPowerOfTen(double value, int exponent) { + while (exponent > 22) { + value *= 1e22; + exponent -= 22; + } + while (exponent < -22) { + value /= 1e22; + exponent += 22; + } + return exponent >= 0 ? value * exactPowerOfTen(exponent) : value / exactPowerOfTen(-exponent); +} + +/** + * @brief RDKit's round_to_n_digits_: `atof(sprintf("%.*g", digits, value))` for a non-negative @p value. + * + * Also returns the rounded value's significand (@p mantissa, exactly @p digits digits unless the value is 0) + * and decimal exponent, from which formatLikeOstream() rebuilds its text. Exact ties round to even, as glibc + * does for exactly representable halves; values within float64 rounding of a tie may differ from printf. + */ +__device__ __forceinline__ double roundToSignificantDigits(const double value, + const int digits, + int64_t& mantissa, + int& exponent) { + mantissa = 0; + exponent = 0; + if (!(value > 0.0)) { + return value; // 0 stays 0; NaN stays NaN. + } + exponent = static_cast(floor(log10(value))); + const double lowest = exactPowerOfTen(digits - 1); + const double highest = exactPowerOfTen(digits); + double scaled = scaleByPowerOfTen(value, digits - 1 - exponent); + if (scaled < lowest) { + --exponent; + scaled = scaleByPowerOfTen(value, digits - 1 - exponent); + } else if (scaled >= highest) { + ++exponent; + scaled = scaleByPowerOfTen(value, digits - 1 - exponent); + } + mantissa = llrint(scaled); + if (mantissa == static_cast(highest)) { + mantissa /= 10; + ++exponent; + } + return scaleByPowerOfTen(static_cast(mantissa), exponent - (digits - 1)); +} + +//! Short text buffer for RDKit's string comparisons. +struct DecimalText { + char chars[24]; + int length; +}; + +__device__ __forceinline__ void appendChar(DecimalText& text, const char character) { + text.chars[text.length++] = character; +} + +/** + * @brief The text `std::ostream << v` (default precision 6, %g style, trailing zeros removed) prints for + * v = @p mantissa * 10^(@p exponent - @p digits + 1), a value from roundToSignificantDigits() with + * @p digits <= 6 significant digits. + */ +__device__ __forceinline__ void formatLikeOstream(const int64_t mantissa, + const int exponent, + const int digits, + DecimalText& text) { + text.length = 0; + if (mantissa == 0) { + appendChar(text, '0'); + return; + } + char significand[8]; + int significandLength = 0; + int64_t remainder = mantissa; + for (int position = digits - 1; position >= 0; --position) { + significand[position] = static_cast('0' + remainder % 10); + remainder /= 10; + } + significandLength = digits; + while (significandLength > 1 && significand[significandLength - 1] == '0') { + --significandLength; + } + if (exponent < -4 || exponent >= 6) { + appendChar(text, significand[0]); + if (significandLength > 1) { + appendChar(text, '.'); + for (int position = 1; position < significandLength; ++position) { + appendChar(text, significand[position]); + } + } + appendChar(text, 'e'); + appendChar(text, exponent < 0 ? '-' : '+'); + const int magnitude = exponent < 0 ? -exponent : exponent; + if (magnitude >= 100) { + appendChar(text, static_cast('0' + magnitude / 100)); + } + appendChar(text, static_cast('0' + (magnitude / 10) % 10)); + appendChar(text, static_cast('0' + magnitude % 10)); + return; + } + if (exponent >= 0) { + for (int position = 0; position <= exponent; ++position) { + appendChar(text, position < significandLength ? significand[position] : '0'); + } + if (significandLength > exponent + 1) { + appendChar(text, '.'); + for (int position = exponent + 1; position < significandLength; ++position) { + appendChar(text, significand[position]); + } + } + return; + } + appendChar(text, '0'); + appendChar(text, '.'); + for (int zero = 0; zero < -exponent - 1; ++zero) { + appendChar(text, '0'); + } + for (int position = 0; position < significandLength; ++position) { + appendChar(text, significand[position]); + } +} + +//! `text.substr(text.find('.') + 1)`: the text after the first '.', or all of it when there is none. +__device__ __forceinline__ DecimalText textAfterPoint(const DecimalText& text) { + int start = 0; + for (int position = 0; position < text.length; ++position) { + if (text.chars[position] == '.') { + start = position + 1; + break; + } + } + DecimalText suffix; + suffix.length = 0; + for (int position = start; position < text.length; ++position) { + appendChar(suffix, text.chars[position]); + } + return suffix; +} + +//! RDKit's countZeros: leading '0' characters. +__device__ __forceinline__ int countLeadingZeros(const DecimalText& text) { + int zeros = 0; + while (zeros < text.length && text.chars[zeros] == '0') { + ++zeros; + } + return zeros; +} + +//! C atoi on the text: leading decimal digits, 0 if there are none. +__device__ __forceinline__ int64_t leadingInteger(const DecimalText& text) { + int64_t value = 0; + for (int position = 0; position < text.length && text.chars[position] >= '0' && text.chars[position] <= '9'; + ++position) { + value = value * 10 + (text.chars[position] - '0'); + } + return value; +} + +//! Rounded leverage with the decimal text RDKit's string comparisons see. +struct RoundedLeverage { + double value; + DecimalText fraction; //!< Text after the decimal point (textAfterPoint of the printed value). +}; + +__device__ __forceinline__ RoundedLeverage roundLeverage(const double leverage, const int digits) { + int64_t mantissa; + int exponent; + RoundedLeverage rounded; + rounded.value = roundToSignificantDigits(leverage, digits, mantissa, exponent); + DecimalText printed; + formatLikeOstream(mantissa, exponent, digits, printed); + rounded.fraction = textAfterPoint(printed); + return rounded; +} + +/** + * @brief RDKit's IsClose2: equal leading-zero counts pad the shorter fraction text once with '0'; texts of + * equal length are close when their leading integers differ by less than 2. + */ +__device__ __forceinline__ bool rdkitIsClose(const DecimalText& first, const DecimalText& second) { + DecimalText a = first; + DecimalText b = second; + if (countLeadingZeros(a) == countLeadingZeros(b)) { + if (a.length > b.length) { + appendChar(b, '0'); + } + if (a.length < b.length) { + appendChar(a, '0'); + } + } + if (a.length != b.length) { + return false; + } + const int64_t difference = leadingInteger(a) - leadingInteger(b); + return difference > -2 && difference < 2; +} + +/** + * @brief ITH and ISH from the heavy-atom leverages, porting RDKit's getGETAWAYDesc: each leverage rounded to + * @p digits significant digits, sorted in descending order and grouped by RDKit's clusterArray2 (a + * greedy scan with IsClose2 comparisons of the printed values), then + * ITH = n log2 n - sum(c log2 c) over the group sizes c and ISH = ITH / (n log2 n). + * + * Runs on one lane. @p sorted holds the conformer's rounded heavy-atom leverages on return. + */ +__device__ __forceinline__ void computeInformationIndices(const double* leverages, + const uint8_t* heavyAtomFlags, + const int numAtoms, + const int digits, + double* sorted, + double& ith, + double& ish) { + int numHeavy = 0; + for (int atomIdx = 0; atomIdx < numAtoms; ++atomIdx) { + if (heavyAtomFlags[atomIdx] != 0) { + int64_t mantissa; + int exponent; + const double rounded = roundToSignificantDigits(leverages[atomIdx], digits, mantissa, exponent); + int slot = numHeavy++; + while (slot > 0 && sorted[slot - 1] < rounded) { + sorted[slot] = sorted[slot - 1]; + --slot; + } + sorted[slot] = rounded; + } + } + const double heavy = numHeavy; + const double total0 = heavy * log(heavy) / log(2.0); + double total = total0; + auto store = [&](const int count) { + const double size = count; + total -= size * log(size) / log(2.0); + }; + int head = 0; + int count = 0; + while (head < numHeavy) { + const RoundedLeverage front = roundLeverage(sorted[head], digits); + ++head; + ++count; + const int remaining = numHeavy - head; + if (remaining == 0) { + store(count); + } + for (int offset = 0; offset < remaining; ++offset) { + const RoundedLeverage other = roundLeverage(sorted[head + offset], digits); + if (rdkitIsClose(front.fraction, other.fraction)) { + ++count; + } else { + store(count); + head += offset; + count = 0; + break; + } + if (offset == remaining - 1) { + store(count); + ++head; + count = 0; + } + } + } + ith = total; + ish = total / total0; +} + +//! Shuffle-combines a value across the warp's pair streams (lanes with equal laneInGroup). +template __device__ __forceinline__ Real sumOverPairStreams(Real value) { + for (int offset = kGroupSize; offset < kWarpSize; offset <<= 1) { + value += __shfl_xor_sync(0xffffffffu, value, offset); + } + return value; +} + +template __device__ __forceinline__ Real maxOverPairStreams(Real value) { + for (int offset = kGroupSize; offset < kWarpSize; offset <<= 1) { + const Real other = __shfl_xor_sync(0xffffffffu, value, offset); + value = other > value ? other : value; + } + return value; +} + +//! Maximum over a kGroupSize-lane group; callable with only that group active. +template __device__ __forceinline__ Real groupMax(Real value) { + const unsigned activeMask = __activemask(); + for (int offset = kGroupSize / 2; offset > 0; offset >>= 1) { + const Real other = __shfl_xor_sync(activeMask, value, offset); + value = other > value ? other : value; + } + return value; +} + +/** + * @brief Warp-collective. Leverages of the conformer's atoms: the diagonal of H = X (X^T X)^+ X^T for the + * centered coordinates X, with RDKit's pseudo-inverse (singular values at or below + * kGetawayPinvTolerance dropped). Writes each atom's scratch row and float64 leverage. + * + * FP64 required: leverages decide ITH and ISH through their rounding to `precision` significant digits and + * RDKit's digit-string clustering, so float32 leverages would move values across rounding boundaries; the + * pseudo-inverse cutoff decides which axes count for (near-)planar conformers. This is O(atoms) work. + */ +template +__device__ __forceinline__ void computeLeverages(const ConformerAtoms& atoms, + const double* propertyWeights, + const int totalAtoms, + const int lane, + Real* rows, + double* leverages) { + double sumX = 0; + double sumY = 0; + double sumZ = 0; + for (int atomIdx = lane; atomIdx < atoms.numAtoms; atomIdx += kWarpSize) { + sumX += atoms.positions[atomIdx * 3 + 0]; + sumY += atoms.positions[atomIdx * 3 + 1]; + sumZ += atoms.positions[atomIdx * 3 + 2]; + } + const double inverseAtoms = 1.0 / static_cast(atoms.numAtoms); + const double centroidX = groupAllReduceSum(sumX) * inverseAtoms; + const double centroidY = groupAllReduceSum(sumY) * inverseAtoms; + const double centroidZ = groupAllReduceSum(sumZ) * inverseAtoms; + + double moments[6] = {}; + for (int atomIdx = lane; atomIdx < atoms.numAtoms; atomIdx += kWarpSize) { + const double x = atoms.positions[atomIdx * 3 + 0] - centroidX; + const double y = atoms.positions[atomIdx * 3 + 1] - centroidY; + const double z = atoms.positions[atomIdx * 3 + 2] - centroidZ; + moments[0] += x * x; + moments[1] += x * y; + moments[2] += x * z; + moments[3] += y * y; + moments[4] += y * z; + moments[5] += z * z; + } + for (double& moment : moments) { + moment = groupAllReduceSum(moment); + } + double eigenvalues[3]; + double eigenvectors[9]; + symmetricEigensystemJacobi3x3(moments[0], + moments[1], + moments[2], + moments[3], + moments[4], + moments[5], + eigenvalues, + eigenvectors); + double inverseRoots[3]; + for (int axis = 0; axis < 3; ++axis) { + const double singular = fabs(eigenvalues[axis]); + inverseRoots[axis] = singular > kGetawayPinvTolerance ? 1.0 / sqrt(singular) : 0.0; + } + + for (int atomIdx = lane; atomIdx < atoms.numAtoms; atomIdx += kWarpSize) { + const double x = atoms.positions[atomIdx * 3 + 0] - centroidX; + const double y = atoms.positions[atomIdx * 3 + 1] - centroidY; + const double z = atoms.positions[atomIdx * 3 + 2] - centroidZ; + Real* const row = rows + atomIdx * kGetawayRowStride; + double leverage = 0; + for (int axis = 0; axis < 3; ++axis) { + const double scaled = + (x * eigenvectors[axis] + y * eigenvectors[3 + axis] + z * eigenvectors[6 + axis]) * inverseRoots[axis]; + row[kGetawayRowY + axis] = static_cast(scaled); + leverage += scaled * scaled; + } + leverages[atomIdx] = leverage; + row[kGetawayRowLeverage] = static_cast(leverage); + row[kGetawayRowRootLev] = static_cast(sqrt(leverage)); + for (int channel = 0; channel < kNumAtomPropertyChannels; ++channel) { + row[kGetawayRowWeights + channel] = + static_cast(propertyWeights[static_cast(channel) * totalAtoms + atomIdx]); + } + row[kGetawayRowCentered + 0] = static_cast(x); + row[kGetawayRowCentered + 1] = static_cast(y); + row[kGetawayRowCentered + 2] = static_cast(z); + } + __syncwarp(); +} + +//! Atom-property weight of channel @p channel (0 is the unweighted channel) from a scratch row. +template __device__ __forceinline__ Real getawayWeight(const Real* row, const int channel) { + return channel == 0 ? Real(1) : row[kGetawayRowWeights + channel - 1]; +} + +//! Distance between two atoms from their scratch rows' centered coordinates. +template __device__ __forceinline__ Real getawayDistance(const Real* first, const Real* second) { + const Real dx = first[kGetawayRowCentered + 0] - second[kGetawayRowCentered + 0]; + const Real dy = first[kGetawayRowCentered + 1] - second[kGetawayRowCentered + 1]; + const Real dz = first[kGetawayRowCentered + 2] - second[kGetawayRowCentered + 2]; + return sqrt(dx * dx + dy * dy + dz * dz); +} + +/** + * @brief Warp-collective. (R v)_i = sqrt(h_i) * sum_{k != i} sqrt(h_k) v_k / r_ik for every atom i into + * @p product; returns v^T R v and |R v|^2, reduced over the warp. Rows are computed on the fly. + */ +template +__device__ __forceinline__ void multiplyByInfluenceMatrix(const Real* rows, + const int numAtoms, + const int lane, + const Real* vector, + Real* product, + Real& rayleighNumerator, + Real& productNormSquared) { + Real numerator = 0; + Real norm = 0; + for (int atomIdx = lane; atomIdx < numAtoms; atomIdx += kWarpSize) { + const Real* rowI = rows + atomIdx * kGetawayRowStride; + Real sum = 0; + for (int other = 0; other < numAtoms; ++other) { + if (other != atomIdx) { + const Real* rowK = rows + other * kGetawayRowStride; + sum += rowK[kGetawayRowRootLev] * vector[other] / getawayDistance(rowI, rowK); + } + } + const Real value = rowI[kGetawayRowRootLev] * sum; + product[atomIdx] = value; + numerator += vector[atomIdx] * value; + norm += value * value; + } + rayleighNumerator = groupAllReduceSum(numerator); + productNormSquared = groupAllReduceSum(norm); + __syncwarp(); +} + +/** + * @brief RDKit's GETAWAY for every conformer; one warp per conformer. + * + * Per conformer: leverages and scratch rows (computeLeverages()); per-atom H0/HATS0 sums, HIC (dimension from + * the PBF of the molecule's default conformer, as RDKit's PBF(mol) call does), HGM (sequential product, so + * underflow matches RDKit) and ITH/ISH (computeInformationIndices()); then the pair loop: each kGroupSize-lane + * group walks interleaved rows j of the pair triangle, finds row j's bond distances (capped at + * kNumGetawayLags) with searchBondDepths() into its @p bondDepths array, and the lane owning a pair's lag + * accumulates its H, HATS and R terms; pairs farther apart or disconnected enter the totals only, as RDKit's + * lag loop reaches its 1e8 disconnected sentinel. REIG, the largest singular value of the nonnegative + * symmetric R, is its Perron eigenvalue, found by power iteration whose first product (with ones) gives the + * row sums for RARS and RCON. + */ +template +__global__ void getaway3DKernel(const DeviceCoordView coordinates, + const Property3DDeviceInputs inputs, + const int precisionDigits, + Real* __restrict__ rows, + double* __restrict__ leverageScratch, + double* __restrict__ sortScratch, + Real* __restrict__ iterateScratch, + uint8_t* __restrict__ bondDepths, + Real* __restrict__ output) { + const int lane = static_cast(threadIdx.x) % kWarpSize; + const int laneInGroup = lane % kGroupSize; + const int stream = lane / kGroupSize; + const int conformerIdx = blockIdx.x * kWarpsPerBlock + static_cast(threadIdx.x) / kWarpSize; + if (conformerIdx >= coordinates.numConformers) { + return; // Uniform across the warp. + } + Real* const row = output + static_cast(conformerIdx) * kNumGetawayProperties; + const ConformerAtoms atoms = loadConformer(coordinates, nullptr, inputs.moleculeAtomStarts, conformerIdx); + if (!atoms.valid) { + for (int valueIdx = lane; valueIdx < kNumGetawayProperties; valueIdx += kWarpSize) { + row[valueIdx] = static_cast(nan("")); + } + return; + } + const int numAtoms = atoms.numAtoms; + const int moleculeIdx = coordinates.molIndices[conformerIdx]; + const int moleculeStart = inputs.moleculeAtomStarts[moleculeIdx]; + const int totalAtoms = inputs.moleculeAtomStarts[coordinates.nMols]; + const size_t atomOffset = static_cast(atoms.positions - coordinates.positions) / 3; + Real* const atomRows = rows + atomOffset * kGetawayRowStride; + double* const leverages = leverageScratch + atomOffset; + computeLeverages(atoms, inputs.atomPropertyWeights + moleculeStart, totalAtoms, lane, atomRows, leverages); + + // HIC's dimension: RDKit's PBF of the molecule's default conformer below kGetawayPlanarPbf means 2D. + int defaultRow = inputs.defaultConformerRows != nullptr ? inputs.defaultConformerRows[moleculeIdx] : -1; + ConformerAtoms defaultAtoms = + defaultRow >= 0 ? loadConformer(coordinates, nullptr, inputs.moleculeAtomStarts, defaultRow) : ConformerAtoms{}; + if (!defaultAtoms.valid) { + defaultRow = conformerIdx; + defaultAtoms = atoms; + } + ProjectionState plane; + computeProjectionCentroid(defaultAtoms, laneInGroup, plane); + computeProjectionPca(defaultAtoms, nullptr, laneInGroup, plane); + const double pbf = computePbf(defaultAtoms, laneInGroup, plane); + const bool is3D = inputs.conformerIs3D == nullptr || inputs.conformerIs3D[defaultRow] != 0; + const double dimension = (defaultAtoms.numAtoms < 4 || !is3D ? 0.0 : pbf) < kGetawayPlanarPbf ? 2.0 : 3.0; + + // Lag 0: H0 = sum w^2 h and HATS0 = sum w^2 h^2 over atoms with h > 0; HIC from every leverage. + Real lagZeroH[kNumGetawayChannels] = {}; + Real lagZeroHats[kNumGetawayChannels] = {}; + double hic = 0; + for (int atomIdx = lane; atomIdx < numAtoms; atomIdx += kWarpSize) { + const Real* atomRow = atomRows + atomIdx * kGetawayRowStride; + const Real leverage = atomRow[kGetawayRowLeverage]; + const double exact = leverages[atomIdx]; + hic -= exact / dimension * log(exact / dimension) / log(2.0); + if (exact > 0.0) { + for (int channel = 0; channel < kNumGetawayChannels; ++channel) { + const Real weight = getawayWeight(atomRow, channel); + lagZeroH[channel] += weight * weight * leverage; + lagZeroHats[channel] += weight * weight * leverage * leverage; + } + } + } + hic = groupAllReduceSum(hic); + for (int channel = 0; channel < kNumGetawayChannels; ++channel) { + lagZeroH[channel] = groupAllReduceSum(lagZeroH[channel]); + lagZeroHats[channel] = groupAllReduceSum(lagZeroHats[channel]); + } + + // Pair loop: lane l owns lag l + 1; pairs beyond kNumGetawayLags (or disconnected) go to lane k % kGroupSize. + Real lagH[kNumGetawayChannels] = {}; + Real lagHats[kNumGetawayChannels] = {}; + Real lagR[kNumGetawayChannels] = {}; + Real lagRMax[kNumGetawayChannels] = {}; + Real farH[kNumGetawayChannels] = {}; + Real farHats[kNumGetawayChannels] = {}; + Real farR[kNumGetawayChannels] = {}; + { + const int32_t* neighborStarts = inputs.bondNeighborStarts + moleculeStart; + uint8_t* depth = bondDepths + atomOffset * kGetawayPairStreams + static_cast(stream) * numAtoms; + const unsigned groupMask = 0xffu << (stream * kGroupSize); + for (int j = stream; j < numAtoms - 1; j += kGetawayPairStreams) { + searchBondDepths(neighborStarts, + inputs.bondNeighbors, + numAtoms, + j, + laneInGroup, + groupMask, + depth); + const Real* rowJ = atomRows + j * kGetawayRowStride; + for (int k = j + 1; k < numAtoms; ++k) { + const int lag = depth[k]; + const bool near = lag != kUnreachedDepth; + const bool owner = near ? lag - 1 == laneInGroup : k % kGroupSize == laneInGroup; + if (!owner) { + continue; + } + const Real* rowK = atomRows + k * kGetawayRowStride; + const Real hatsTerm = rowJ[kGetawayRowLeverage] * rowK[kGetawayRowLeverage]; + const Real influence = rowJ[kGetawayRowY + 0] * rowK[kGetawayRowY + 0] + + rowJ[kGetawayRowY + 1] * rowK[kGetawayRowY + 1] + + rowJ[kGetawayRowY + 2] * rowK[kGetawayRowY + 2]; + const Real rTerm = rowJ[kGetawayRowRootLev] * rowK[kGetawayRowRootLev] / getawayDistance(rowJ, rowK); + for (int channel = 0; channel < kNumGetawayChannels; ++channel) { + const Real product = getawayWeight(rowJ, channel) * getawayWeight(rowK, channel); + if (near) { + lagHats[channel] += product * hatsTerm; + if (influence > Real(0)) { + lagH[channel] += product * influence; + } + const Real weighted = product * rTerm; + lagR[channel] += weighted; + if (weighted > lagRMax[channel]) { + lagRMax[channel] = weighted; + } + } else { + farHats[channel] += product * hatsTerm; + if (influence > Real(0)) { + farH[channel] += product * influence; + } + farR[channel] += product * rTerm; + } + } + } + } + } + for (int channel = 0; channel < kNumGetawayChannels; ++channel) { + lagH[channel] = sumOverPairStreams(lagH[channel]); + lagHats[channel] = sumOverPairStreams(lagHats[channel]); + lagR[channel] = sumOverPairStreams(lagR[channel]); + lagRMax[channel] = maxOverPairStreams(lagRMax[channel]); + farH[channel] = sumOverPairStreams(farH[channel]); + farHats[channel] = sumOverPairStreams(farHats[channel]); + farR[channel] = sumOverPairStreams(farR[channel]); + } + + // REIG by power iteration on R; the first product with ones gives R's row sums for RARS and RCON. + Real* const current = iterateScratch + atomOffset * 2; + Real* const next = current + numAtoms; + for (int atomIdx = lane; atomIdx < numAtoms; atomIdx += kWarpSize) { + current[atomIdx] = Real(1); + } + __syncwarp(); + Real numerator; + Real normSquared; + multiplyByInfluenceMatrix(atomRows, numAtoms, lane, current, next, numerator, normSquared); + Real rowSum = 0; + Real rcon = 0; + for (int atomIdx = lane; atomIdx < numAtoms; atomIdx += kWarpSize) { + rowSum += next[atomIdx]; + for (int edge = inputs.bondNeighborStarts[moleculeStart + atomIdx]; + edge < inputs.bondNeighborStarts[moleculeStart + atomIdx + 1]; + ++edge) { + const int neighbor = inputs.bondNeighbors[edge]; + if (neighbor > atomIdx) { + rcon += sqrt(next[atomIdx] * next[neighbor]); + } + } + } + const Real rars = groupAllReduceSum(rowSum) / static_cast(numAtoms); + rcon = groupAllReduceSum(rcon); + Real eigenvalue = numerator / static_cast(numAtoms); + constexpr Real kTolerance = sizeof(Real) == sizeof(double) ? Real(1e-13) : Real(2e-7); + for (int iteration = 0; iteration < kGetawayMaxIterations && normSquared > Real(0); ++iteration) { + const Real inverseNorm = Real(1) / sqrt(normSquared); + for (int atomIdx = lane; atomIdx < numAtoms; atomIdx += kWarpSize) { + current[atomIdx] = next[atomIdx] * inverseNorm; + } + __syncwarp(); + multiplyByInfluenceMatrix(atomRows, numAtoms, lane, current, next, numerator, normSquared); + const Real previous = eigenvalue; + eigenvalue = numerator; + if (fabs(eigenvalue - previous) <= kTolerance * fabs(eigenvalue)) { + break; + } + } + if (!(normSquared > Real(0))) { + eigenvalue = Real(0); + } + + if (lane == 0) { + double ith; + double ish; + computeInformationIndices(leverages, + inputs.heavyAtomFlags + moleculeStart, + numAtoms, + precisionDigits, + sortScratch + atomOffset, + ith, + ish); + double leverageProduct = 1.0; + for (int atomIdx = 0; atomIdx < numAtoms; ++atomIdx) { + leverageProduct *= leverages[atomIdx]; + } + const double hgm = 100.0 * pow(leverageProduct, 1.0 / numAtoms); + row[0] = static_cast(roundThousandths(ith)); + row[1] = static_cast(roundThousandths(ish)); + row[2] = static_cast(roundThousandths(hic)); + row[3] = static_cast(roundThousandths(hgm)); + row[kGetawayRconIndex + 0] = roundThousandths(rcon); + row[kGetawayRconIndex + 1] = roundThousandths(rars); + row[kGetawayRconIndex + 2] = roundThousandths(eigenvalue); + } + if (stream == 0) { + for (int channel = 0; channel < kNumGetawayChannels; ++channel) { + Real* const hBlock = row + kGetawayHStart + channel * kGetawayHBlock; + Real* const rBlock = row + kGetawayRStart + channel * kGetawayRBlock; + const Real lagHSum = groupAllReduceSum(lagH[channel]); + const Real lagHatsSum = groupAllReduceSum(lagHats[channel]); + const Real lagRSum = groupAllReduceSum(lagR[channel]); + const Real farHSum = groupAllReduceSum(farH[channel]); + const Real farHatsSum = groupAllReduceSum(farHats[channel]); + const Real farRSum = groupAllReduceSum(farR[channel]); + const Real rMax = groupMax(lagRMax[channel]); + const int lagIdx = laneInGroup + 1; + hBlock[lagIdx] = roundThousandths(lagH[channel]); + hBlock[kNumGetawayLags + 2 + lagIdx] = roundThousandths(lagHats[channel]); + rBlock[lagIdx - 1] = roundThousandths(lagR[channel]); + rBlock[kNumGetawayLags + lagIdx] = roundThousandths(lagRMax[channel]); + if (laneInGroup == 0) { + hBlock[0] = roundThousandths(lagZeroH[channel]); + hBlock[kNumGetawayLags + 1] = roundThousandths(lagZeroH[channel] + Real(2) * (lagHSum + farHSum)); + hBlock[kNumGetawayLags + 2] = roundThousandths(lagZeroHats[channel]); + hBlock[kGetawayHBlock - 1] = roundThousandths(lagZeroHats[channel] + Real(2) * (lagHatsSum + farHatsSum)); + rBlock[kNumGetawayLags] = roundThousandths(Real(2) * (lagRSum + farRSum)); + rBlock[kGetawayRBlock - 1] = roundThousandths(rMax); + } + } + } +} + +/** + * @brief Launches the GETAWAY kernel. Scratch per atom: one kGetawayRowStride row of @p Real, a float64 + * leverage and sort slot, two power-iteration values and one bond-depth byte per pair stream. + */ +template +void launchGetawayProperties(const DeviceCoordView& coordinates, + const Property3DDeviceInputs& inputs, + const GetawayOptions& options, + Real* output, + const cudaStream_t stream) { + const int numConformers = coordinates.numConformers; + if (numConformers == 0 || output == nullptr) { + return; + } + const size_t numAtoms = static_cast(coordinates.numAtoms); + AsyncDeviceVector rows(numAtoms * kGetawayRowStride, stream); + AsyncDeviceVector leverages(numAtoms, stream); + AsyncDeviceVector sorted(numAtoms, stream); + AsyncDeviceVector iterates(numAtoms * 2, stream); + AsyncDeviceVector depths(numAtoms * kGetawayPairStreams, stream); + const int numBlocks = (numConformers + kWarpsPerBlock - 1) / kWarpsPerBlock; + getaway3DKernel<<>>(coordinates, + inputs, + static_cast(options.precision), + rows.data(), + leverages.data(), + sorted.data(), + iterates.data(), + depths.data(), + output); + cudaCheckError(cudaGetLastError()); +} + +} // namespace nvMolKit::descriptors3d_detail + +#endif // NVMOLKIT_DESCRIPTORS3D_GETAWAY_CUH diff --git a/src/descriptors3d_kernel.cuh b/src/descriptors3d_kernel.cuh index a837b3bd..fce5bce4 100644 --- a/src/descriptors3d_kernel.cuh +++ b/src/descriptors3d_kernel.cuh @@ -137,6 +137,47 @@ writeCenteredConformer(const ConformerAtoms& atoms, const int laneInGroup, const return make_double3(centroidX, centroidY, centroidZ); } +//! Bond depth of an atom not reached within the search's depth cap. +constexpr uint8_t kUnreachedDepth = 0xFF; + +/** + * @brief Collective over the @p groupMask lanes. Breadth-first search over the molecule's bonds from atom + * @p source, writing each atom's bond-count distance to @p depth, or kUnreachedDepth beyond + * @p kMaxDepth bonds. Matches RDKit's MolOps::getDistanceMat(mol, false) up to that cap. + * + * Each level scans the atoms set at the previous level, O(kMaxDepth * atoms) per source; lanes that reach the + * same atom write the same value. + */ +template +__device__ __forceinline__ void searchBondDepths(const int32_t* neighborStarts, + const int32_t* neighbors, + const int numAtoms, + const int source, + const int laneInGroup, + const unsigned groupMask, + uint8_t* depth) { + static_assert(kMaxDepth < kUnreachedDepth); + // Every lane has finished reading the previous source's depths before they are overwritten. + __syncwarp(groupMask); + for (int atomIdx = laneInGroup; atomIdx < numAtoms; atomIdx += kGroupSize) { + depth[atomIdx] = atomIdx == source ? 0 : kUnreachedDepth; + } + __syncwarp(groupMask); + for (int level = 1; level <= kMaxDepth; ++level) { + for (int atomIdx = laneInGroup; atomIdx < numAtoms; atomIdx += kGroupSize) { + if (depth[atomIdx] == level - 1) { + for (int edge = neighborStarts[atomIdx]; edge < neighborStarts[atomIdx + 1]; ++edge) { + const int neighbor = neighbors[edge]; + if (depth[neighbor] == kUnreachedDepth) { + depth[neighbor] = static_cast(level); + } + } + } + } + __syncwarp(groupMask); + } +} + //! Round to three decimals as RDKit's WHIM, RDF and MORSE do: std::round(1000 * x) / 1000. template __device__ __forceinline__ Real roundThousandths(const Real value) { return round(value * Real(1000)) / Real(1000); diff --git a/src/descriptors3d_mol.cpp b/src/descriptors3d_mol.cpp index d6f77bdb..e5233e1f 100644 --- a/src/descriptors3d_mol.cpp +++ b/src/descriptors3d_mol.cpp @@ -33,6 +33,8 @@ struct DeviceDescriptorInputs { AsyncDeviceVector bondNeighborStarts; AsyncDeviceVector bondNeighbors; AsyncDeviceVector usrcatAtomClasses; + AsyncDeviceVector heavyAtomFlags; + AsyncDeviceVector defaultConformerRows; AsyncDeviceVector conformerIs3D; AsyncDeviceVector moleculeAtomStarts; }; @@ -45,6 +47,8 @@ struct DescriptorInputNeeds { bool covalentRadiusWeights = false; bool bondAdjacency = false; bool usrcatAtomClasses = false; + bool heavyAtomFlags = false; + bool defaultConformerRows = false; bool conformerFlags = false; }; @@ -136,7 +140,20 @@ DeviceDescriptorInputs uploadDescriptorInputs(const std::vector usrcatAtomClasses(needs.usrcatAtomClasses ? static_cast(totalAtoms) : 0); - std::vector conformerIs3D; + std::vector heavyAtomFlags(needs.heavyAtomFlags ? static_cast(totalAtoms) : 0); + // Coordinate rows follow each molecule's conformers in order, so a molecule's first row is its default + // conformer (RDKit's getConformer(-1)). + std::vector defaultConformerRows; + if (needs.defaultConformerRows) { + defaultConformerRows.resize(numMols); + int32_t row = 0; + for (int molIdx = 0; molIdx < numMols; ++molIdx) { + const int32_t numConformers = static_cast(mols[molIdx]->getNumConformers()); + defaultConformerRows[molIdx] = numConformers > 0 ? row : -1; + row += numConformers; + } + } + std::vector conformerIs3D; if (needs.conformerFlags) { for (const RDKit::ROMol* mol : mols) { for (auto conformer = mol->beginConformers(); conformer != mol->endConformers(); ++conformer) { @@ -161,7 +178,7 @@ DeviceDescriptorInputs uploadDescriptorInputs(const std::vector radii = descriptorData.GetRelativeRcov(mol); std::copy(radii.begin(), radii.end(), covalentRadiusWeights.begin() + atomStart); } + if (needs.heavyAtomFlags) { + for (const auto* atom : mol.atoms()) { + heavyAtomFlags[atomStart + atom->getIdx()] = atom->getAtomicNum() > 1 ? 1 : 0; + } + } if (needs.bondAdjacency) { writeBondAdjacency(mol, neighborOffsets[molIdx], bondNeighborStarts.data() + atomStart, bondNeighbors.data()); } @@ -233,6 +256,8 @@ DeviceDescriptorInputs uploadDescriptorInputs(const std::vector(bondNeighborStarts.size(), stream), AsyncDeviceVector(bondNeighbors.size(), stream), AsyncDeviceVector(usrcatAtomClasses.size(), stream), + AsyncDeviceVector(heavyAtomFlags.size(), stream), + AsyncDeviceVector(defaultConformerRows.size(), stream), AsyncDeviceVector(conformerIs3D.size(), stream), AsyncDeviceVector(atomStarts.size(), stream)}; if (!weights.empty()) { @@ -256,6 +281,12 @@ DeviceDescriptorInputs uploadDescriptorInputs(const std::vector calc3DProperties(const std::vector calc3DProperties(const std::vector(mol->getNumAtoms())); diff --git a/src/descriptors3d_pairwise.cuh b/src/descriptors3d_pairwise.cuh index 61c82ef8..9235f11b 100644 --- a/src/descriptors3d_pairwise.cuh +++ b/src/descriptors3d_pairwise.cuh @@ -28,45 +28,6 @@ constexpr int kNumAutocorrChannels = kNumPairwiseChannels + 1; constexpr int kNumAutocorrLags = kNumAutocorr3DProperties / kNumAutocorrChannels; constexpr int kAutocorrLagsPerLane = (kNumAutocorrLags + kGroupSize - 1) / kGroupSize; -//! Bond depth of an atom not reached within kNumAutocorrLags bonds of the search source. -constexpr uint8_t kUnreachedDepth = 0xFF; - -/** - * @brief Collective over the @p groupMask lanes. Breadth-first search over the molecule's bonds from atom - * @p source, writing each atom's bond-count distance to @p depth, or kUnreachedDepth beyond - * kNumAutocorrLags bonds. Matches RDKit's MolOps::getDistanceMat(mol, false) up to that cap. - * - * Each level scans the atoms set at the previous level, O(kNumAutocorrLags * atoms) per source; lanes that - * reach the same atom write the same value. - */ -__device__ __forceinline__ void searchBondDepths(const int32_t* neighborStarts, - const int32_t* neighbors, - const int numAtoms, - const int source, - const int laneInGroup, - const unsigned groupMask, - uint8_t* depth) { - // Every lane has finished reading the previous source's depths before they are overwritten. - __syncwarp(groupMask); - for (int atomIdx = laneInGroup; atomIdx < numAtoms; atomIdx += kGroupSize) { - depth[atomIdx] = atomIdx == source ? 0 : kUnreachedDepth; - } - __syncwarp(groupMask); - for (int level = 1; level <= kNumAutocorrLags; ++level) { - for (int atomIdx = laneInGroup; atomIdx < numAtoms; atomIdx += kGroupSize) { - if (depth[atomIdx] == level - 1) { - for (int edge = neighborStarts[atomIdx]; edge < neighborStarts[atomIdx + 1]; ++edge) { - const int neighbor = neighbors[edge]; - if (depth[neighbor] == kUnreachedDepth) { - depth[neighbor] = static_cast(level); - } - } - } - } - __syncwarp(groupMask); - } -} - //! Pairwise properties served by one kernel instantiation; unrequested ones compile out. enum PairwiseSet : unsigned { kPairwiseRdf = 1u, @@ -287,13 +248,13 @@ __global__ void pairwise3DKernel(const DeviceCoordView coordinates, const Real dragJ = rowJ[kScratchIStateDrag]; const Real radiusJ = rowJ[kScratchCovalentRadius]; if constexpr (kAutocorr) { - searchBondDepths(neighborStarts, - inputs.bondNeighbors, - atoms.numAtoms, - j, - laneInGroup, - 0xffu << (stream * kGroupSize), - depth); + searchBondDepths(neighborStarts, + inputs.bondNeighbors, + atoms.numAtoms, + j, + laneInGroup, + 0xffu << (stream * kGroupSize), + depth); } for (int k = j + 1; k < atoms.numAtoms; ++k) { const Real* const rowK = rows + k * kPairwiseScratchStride; diff --git a/tests/test_descriptors3d.cpp b/tests/test_descriptors3d.cpp index f8f7ae6c..65b786c5 100644 --- a/tests/test_descriptors3d.cpp +++ b/tests/test_descriptors3d.cpp @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -366,6 +367,27 @@ TEST(Descriptors3DPairwise, MatchesRdkitAutocorr3D) { } } +TEST(Descriptors3DGetaway, MatchesRdkitGetaway) { + auto mol = molWithConformers("CCCO", + { + {{-1.3, 0.2, 0.7}, {-0.2, -0.8, 0.1}, {0.9, 0.4, -0.6}, {1.7, 1.1, 0.9}}, + {{2.0, -1.0, 0.5}, {2.7, 0.3, -0.4}, {3.9, -0.2, 0.8}, {4.6, 1.0, -0.7}}, + }); + const std::vector mols = {mol.get()}; + auto results = nvMolKit::calc3DProperties(mols, {Property3D::GETAWAY}, {}, nullptr); + const auto values = toHost(results.properties.at(Property3D::GETAWAY)); + ASSERT_EQ(values.size(), 2u * nvMolKit::kNumGetawayProperties); + for (int confIdx = 0; confIdx < 2; ++confIdx) { + std::vector expected; + RDKit::Descriptors::GETAWAY(*mol, expected, confIdx); + ASSERT_EQ(expected.size(), static_cast(nvMolKit::kNumGetawayProperties)); + for (int valueIdx = 0; valueIdx < nvMolKit::kNumGetawayProperties; ++valueIdx) { + EXPECT_NEAR(values[confIdx * nvMolKit::kNumGetawayProperties + valueIdx], expected[valueIdx], 1.1e-3) + << "conformer " << confIdx << ", value " << valueIdx; + } + } +} + TEST(Descriptors3DUsr, MatchesRdkitUsrAndUsrcat) { auto mol = molWithConformers("CCCO", { From 71721487b3baa4b441f84f946842e7426f875a9d Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Tue, 6 Oct 2026 17:14:38 -0400 Subject: [PATCH 3/8] Remove RDKit-comparison narration from GETAWAY docs and comments --- benchmarks/descriptors3d_bench.py | 2 +- nvmolkit/descriptors3d.py | 5 +---- nvmolkit/tests/test_descriptors3d.py | 2 +- src/descriptors3d_getaway.cuh | 7 +++---- 4 files changed, 6 insertions(+), 10 deletions(-) diff --git a/benchmarks/descriptors3d_bench.py b/benchmarks/descriptors3d_bench.py index 768609f9..d673d6e1 100644 --- a/benchmarks/descriptors3d_bench.py +++ b/benchmarks/descriptors3d_bench.py @@ -86,7 +86,7 @@ def _calc_rdkit_property(mol: Chem.Mol, conf_id: int, prop: Property3D) -> float return rdMolDescriptors.GetUSRCAT(mol, confId=conf_id) if prop == Property3D.GETAWAY: if len(Chem.GetMolFrags(mol)) > 1: - # RDKit's GETAWAY loops up to its 1e8 disconnected-pair distance and would not finish. + # RDKit's GETAWAY does not finish in reasonable time on multi-fragment molecules. return [math.nan] * 273 return rdMolDescriptors.CalcGETAWAY(mol, confId=conf_id) return RDKIT_CALCULATORS[prop](mol, confId=conf_id, useAtomicMasses=True) diff --git a/nvmolkit/descriptors3d.py b/nvmolkit/descriptors3d.py index 987afede..9679ed36 100644 --- a/nvmolkit/descriptors3d.py +++ b/nvmolkit/descriptors3d.py @@ -40,10 +40,7 @@ class Property3D(Enum): Conformers with fewer than 3 atoms, which RDKit rejects, produce NaN. - ``USRCAT``: 60 values, ``USR`` for all atoms and then for RDKit's hydrophobic, aromatic, acceptor and donor atoms. Same atom-count rule. - - ``GETAWAY``: 273 values in RDKit's ``CalcGETAWAY`` order. RDKit counts - atom pairs in different fragments (salts, counterions) only in the - totals; nvMolKit returns the same values, which RDKit itself takes - impractically long to reach for multi-fragment molecules. + - ``GETAWAY``: 273 values in RDKit's ``CalcGETAWAY`` order. """ PMI1 = "PMI1" diff --git a/nvmolkit/tests/test_descriptors3d.py b/nvmolkit/tests/test_descriptors3d.py index 76deb69e..915a12be 100644 --- a/nvmolkit/tests/test_descriptors3d.py +++ b/nvmolkit/tests/test_descriptors3d.py @@ -577,7 +577,7 @@ def test_getaway_matches_rdkit(precision): @pytest.fixture(scope="module") def getaway_salt(): - # RDKit's lag loop runs to its 1e8 distance for disconnected atoms: seconds even for this 4-atom salt. + # RDKit's GETAWAY takes seconds on this 4-atom salt, so its reference is computed once. mol = _mol_with_conformers("CCO.N", [[(0.0, 0.0, 0.0), (1.5, 0.1, 0.0), (2.1, 1.4, 0.2), (5.0, 1.0, -1.0)]]) return mol, _rdkit_getaway_rows([mol]) diff --git a/src/descriptors3d_getaway.cuh b/src/descriptors3d_getaway.cuh index 1469c295..f016d8fe 100644 --- a/src/descriptors3d_getaway.cuh +++ b/src/descriptors3d_getaway.cuh @@ -491,10 +491,9 @@ __device__ __forceinline__ void multiplyByInfluenceMatrix(const Real* rows, * underflow matches RDKit) and ITH/ISH (computeInformationIndices()); then the pair loop: each kGroupSize-lane * group walks interleaved rows j of the pair triangle, finds row j's bond distances (capped at * kNumGetawayLags) with searchBondDepths() into its @p bondDepths array, and the lane owning a pair's lag - * accumulates its H, HATS and R terms; pairs farther apart or disconnected enter the totals only, as RDKit's - * lag loop reaches its 1e8 disconnected sentinel. REIG, the largest singular value of the nonnegative - * symmetric R, is its Perron eigenvalue, found by power iteration whose first product (with ones) gives the - * row sums for RARS and RCON. + * accumulates its H, HATS and R terms; pairs farther apart or disconnected enter the totals only. REIG, the largest + * singular value of the nonnegative symmetric R, is its Perron eigenvalue, found by power iteration whose first product + * (with ones) gives the row sums for RARS and RCON. */ template __global__ void getaway3DKernel(const DeviceCoordView coordinates, From 18fd7869675d6b3315419fc56441d656c925beb8 Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Tue, 6 Oct 2026 17:18:45 -0400 Subject: [PATCH 4/8] Use the shared tile and cross-group reductions in the GETAWAY kernel 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. --- src/descriptors3d_getaway.cuh | 61 +++++++++++------------------------ src/descriptors3d_kernel.cuh | 9 ++++++ 2 files changed, 27 insertions(+), 43 deletions(-) diff --git a/src/descriptors3d_getaway.cuh b/src/descriptors3d_getaway.cuh index f016d8fe..d20562e3 100644 --- a/src/descriptors3d_getaway.cuh +++ b/src/descriptors3d_getaway.cuh @@ -325,32 +325,6 @@ __device__ __forceinline__ void computeInformationIndices(const double* leverag ish = total / total0; } -//! Shuffle-combines a value across the warp's pair streams (lanes with equal laneInGroup). -template __device__ __forceinline__ Real sumOverPairStreams(Real value) { - for (int offset = kGroupSize; offset < kWarpSize; offset <<= 1) { - value += __shfl_xor_sync(0xffffffffu, value, offset); - } - return value; -} - -template __device__ __forceinline__ Real maxOverPairStreams(Real value) { - for (int offset = kGroupSize; offset < kWarpSize; offset <<= 1) { - const Real other = __shfl_xor_sync(0xffffffffu, value, offset); - value = other > value ? other : value; - } - return value; -} - -//! Maximum over a kGroupSize-lane group; callable with only that group active. -template __device__ __forceinline__ Real groupMax(Real value) { - const unsigned activeMask = __activemask(); - for (int offset = kGroupSize / 2; offset > 0; offset >>= 1) { - const Real other = __shfl_xor_sync(activeMask, value, offset); - value = other > value ? other : value; - } - return value; -} - /** * @brief Warp-collective. Leverages of the conformer's atoms: the diagonal of H = X (X^T X)^+ X^T for the * centered coordinates X, with RDKit's pseudo-inverse (singular values at or below @@ -625,13 +599,13 @@ __global__ void getaway3DKernel(const DeviceCoordView coordinates, } } for (int channel = 0; channel < kNumGetawayChannels; ++channel) { - lagH[channel] = sumOverPairStreams(lagH[channel]); - lagHats[channel] = sumOverPairStreams(lagHats[channel]); - lagR[channel] = sumOverPairStreams(lagR[channel]); - lagRMax[channel] = maxOverPairStreams(lagRMax[channel]); - farH[channel] = sumOverPairStreams(farH[channel]); - farHats[channel] = sumOverPairStreams(farHats[channel]); - farR[channel] = sumOverPairStreams(farR[channel]); + lagH[channel] = sumAcrossGroups(lagH[channel]); + lagHats[channel] = sumAcrossGroups(lagHats[channel]); + lagR[channel] = sumAcrossGroups(lagR[channel]); + lagRMax[channel] = maxAcrossGroups(lagRMax[channel]); + farH[channel] = sumAcrossGroups(farH[channel]); + farHats[channel] = sumAcrossGroups(farHats[channel]); + farR[channel] = sumAcrossGroups(farR[channel]); } // REIG by power iteration on R; the first product with ones gives R's row sums for RARS and RCON. @@ -703,16 +677,17 @@ __global__ void getaway3DKernel(const DeviceCoordView coordinates, } if (stream == 0) { for (int channel = 0; channel < kNumGetawayChannels; ++channel) { - Real* const hBlock = row + kGetawayHStart + channel * kGetawayHBlock; - Real* const rBlock = row + kGetawayRStart + channel * kGetawayRBlock; - const Real lagHSum = groupAllReduceSum(lagH[channel]); - const Real lagHatsSum = groupAllReduceSum(lagHats[channel]); - const Real lagRSum = groupAllReduceSum(lagR[channel]); - const Real farHSum = groupAllReduceSum(farH[channel]); - const Real farHatsSum = groupAllReduceSum(farHats[channel]); - const Real farRSum = groupAllReduceSum(farR[channel]); - const Real rMax = groupMax(lagRMax[channel]); - const int lagIdx = laneInGroup + 1; + Real* const hBlock = row + kGetawayHStart + channel * kGetawayHBlock; + Real* const rBlock = row + kGetawayRStart + channel * kGetawayRBlock; + const Real lagHSum = groupAllReduceSum(lagH[channel]); + const Real lagHatsSum = groupAllReduceSum(lagHats[channel]); + const Real lagRSum = groupAllReduceSum(lagR[channel]); + const Real farHSum = groupAllReduceSum(farH[channel]); + const Real farHatsSum = groupAllReduceSum(farHats[channel]); + const Real farRSum = groupAllReduceSum(farR[channel]); + const Real rMax = + cooperative_groups::reduce(laneTile(), lagRMax[channel], cooperative_groups::greater()); + const int lagIdx = laneInGroup + 1; hBlock[lagIdx] = roundThousandths(lagH[channel]); hBlock[kNumGetawayLags + 2 + lagIdx] = roundThousandths(lagHats[channel]); rBlock[lagIdx - 1] = roundThousandths(lagR[channel]); diff --git a/src/descriptors3d_kernel.cuh b/src/descriptors3d_kernel.cuh index fce5bce4..ea17cae0 100644 --- a/src/descriptors3d_kernel.cuh +++ b/src/descriptors3d_kernel.cuh @@ -80,6 +80,15 @@ template __device__ __forceinline__ Real sumAcrossGroups(Real va return value; } +//! Maximum over the warp's lanes with equal `lane % kGroupSize`; see sumAcrossGroups(). +template __device__ __forceinline__ Real maxAcrossGroups(Real value) { + for (int offset = kGroupSize; offset < kWarpSize; offset <<= 1) { + const Real other = __shfl_xor_sync(0xffffffffu, value, offset); + value = other > value ? other : value; + } + return value; +} + /** * @brief Coordinate position minus a centroid, converted to @p Real after the subtraction. * From ab0c2001c25bba3969db7a455a1c668b60394199 Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Tue, 6 Oct 2026 17:34:00 -0400 Subject: [PATCH 5/8] Compute GETAWAY's ITH and ISH across the warp 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. --- src/descriptors3d_getaway.cuh | 164 +++++++++++++++++++++++----------- 1 file changed, 110 insertions(+), 54 deletions(-) diff --git a/src/descriptors3d_getaway.cuh b/src/descriptors3d_getaway.cuh index d20562e3..f4d7dbdb 100644 --- a/src/descriptors3d_getaway.cuh +++ b/src/descriptors3d_getaway.cuh @@ -219,85 +219,134 @@ __device__ __forceinline__ int64_t leadingInteger(const DecimalText& text) { return value; } -//! Rounded leverage with the decimal text RDKit's string comparisons see. -struct RoundedLeverage { - double value; - DecimalText fraction; //!< Text after the decimal point (textAfterPoint of the printed value). +//! What RDKit's IsClose2 reads from a rounded leverage's printed fraction text (textAfterPoint()): its +//! leading '0' count, its length, its leading integer (C atoi) and whether it is all digits, in which case +//! appending '0' multiplies that integer by 10. +struct ClusterKey { + int zeros; + int length; + int64_t leading; + bool allDigits; }; -__device__ __forceinline__ RoundedLeverage roundLeverage(const double leverage, const int digits) { - int64_t mantissa; - int exponent; - RoundedLeverage rounded; - rounded.value = roundToSignificantDigits(leverage, digits, mantissa, exponent); +//! Packs a ClusterKey into 64 bits: leading integer (below 2^40) in the low bits, then length, zeros and the +//! all-digits flag. +__device__ __forceinline__ int64_t packClusterKey(const ClusterKey& key) { + return key.leading | (int64_t{key.length} << 40) | (int64_t{key.zeros} << 48) | + (int64_t{key.allDigits ? 1 : 0} << 56); +} + +__device__ __forceinline__ ClusterKey unpackClusterKey(const int64_t packed) { + return {static_cast((packed >> 48) & 0xff), + static_cast((packed >> 40) & 0xff), + packed & ((int64_t{1} << 40) - 1), + ((packed >> 56) & 1) != 0}; +} + +//! Rounds @p leverage to @p digits significant digits (into @p rounded) and returns the ClusterKey of its +//! printed text. +__device__ __forceinline__ ClusterKey clusterKey(const double leverage, const int digits, double& rounded) { + int64_t mantissa; + int exponent; + rounded = roundToSignificantDigits(leverage, digits, mantissa, exponent); DecimalText printed; formatLikeOstream(mantissa, exponent, digits, printed); - rounded.fraction = textAfterPoint(printed); - return rounded; + const DecimalText fraction = textAfterPoint(printed); + ClusterKey key{countLeadingZeros(fraction), fraction.length, leadingInteger(fraction), true}; + for (int position = 0; position < fraction.length; ++position) { + key.allDigits = key.allDigits && fraction.chars[position] >= '0' && fraction.chars[position] <= '9'; + } + return key; +} + +//! The effect of appending '0' to the text a ClusterKey describes. +__device__ __forceinline__ void appendZero(ClusterKey& key) { + ++key.length; + if (key.allDigits) { + key.leading *= 10; + } } /** - * @brief RDKit's IsClose2: equal leading-zero counts pad the shorter fraction text once with '0'; texts of - * equal length are close when their leading integers differ by less than 2. + * @brief RDKit's IsClose2 on the keys of two printed values: equal leading-zero counts pad the shorter text + * once with '0'; texts of equal length are close when their leading integers differ by less than 2. */ -__device__ __forceinline__ bool rdkitIsClose(const DecimalText& first, const DecimalText& second) { - DecimalText a = first; - DecimalText b = second; - if (countLeadingZeros(a) == countLeadingZeros(b)) { +__device__ __forceinline__ bool rdkitIsClose(ClusterKey a, ClusterKey b) { + if (a.zeros == b.zeros) { if (a.length > b.length) { - appendChar(b, '0'); + appendZero(b); } if (a.length < b.length) { - appendChar(a, '0'); + appendZero(a); } } if (a.length != b.length) { return false; } - const int64_t difference = leadingInteger(a) - leadingInteger(b); + const int64_t difference = a.leading - b.leading; return difference > -2 && difference < 2; } /** - * @brief ITH and ISH from the heavy-atom leverages, porting RDKit's getGETAWAYDesc: each leverage rounded to - * @p digits significant digits, sorted in descending order and grouped by RDKit's clusterArray2 (a - * greedy scan with IsClose2 comparisons of the printed values), then + * @brief Warp-collective. ITH and ISH from the heavy-atom leverages, porting RDKit's getGETAWAYDesc: each + * leverage rounded to @p digits significant digits, sorted in descending order and grouped by RDKit's + * clusterArray2 (a greedy scan with IsClose2 comparisons of the printed values), then * ITH = n log2 n - sum(c log2 c) over the group sizes c and ISH = ITH / (n log2 n). * - * Runs on one lane. @p sorted holds the conformer's rounded heavy-atom leverages on return. + * The lanes round, key and rank the heavy atoms (equal values print identically, so their relative order + * does not matter); lane 0 then runs the greedy scan on the sorted keys and alone receives @p ith and @p ish. + * @p heavyValues and @p unsortedKeys hold one entry and @p sortedKeys one entry per heavy atom. */ __device__ __forceinline__ void computeInformationIndices(const double* leverages, const uint8_t* heavyAtomFlags, const int numAtoms, const int digits, - double* sorted, + const int lane, + double* heavyValues, + int64_t* unsortedKeys, + int64_t* sortedKeys, double& ith, double& ish) { - int numHeavy = 0; - for (int atomIdx = 0; atomIdx < numAtoms; ++atomIdx) { - if (heavyAtomFlags[atomIdx] != 0) { - int64_t mantissa; - int exponent; - const double rounded = roundToSignificantDigits(leverages[atomIdx], digits, mantissa, exponent); - int slot = numHeavy++; - while (slot > 0 && sorted[slot - 1] < rounded) { - sorted[slot] = sorted[slot - 1]; - --slot; - } - sorted[slot] = rounded; + const auto warp = laneTile(); + int numHeavy = 0; + for (int base = 0; base < numAtoms; base += kWarpSize) { + const int atomIdx = base + lane; + const bool isHeavy = atomIdx < numAtoms && heavyAtomFlags[atomIdx] != 0; + const unsigned heavy = warp.ballot(isHeavy); + if (isHeavy) { + const int slot = numHeavy + __popc(heavy & ((1u << lane) - 1u)); + double rounded; + const auto key = clusterKey(leverages[atomIdx], digits, rounded); + heavyValues[slot] = rounded; + unsortedKeys[slot] = packClusterKey(key); } + numHeavy += __popc(heavy); + } + warp.sync(); + for (int slot = lane; slot < numHeavy; slot += kWarpSize) { + const double value = heavyValues[slot]; + int rank = 0; + for (int other = 0; other < numHeavy; ++other) { + const double otherValue = heavyValues[other]; + rank += otherValue > value || (otherValue == value && other < slot) ? 1 : 0; + } + sortedKeys[rank] = unsortedKeys[slot]; + } + warp.sync(); + if (lane != 0) { + return; } - const double heavy = numHeavy; - const double total0 = heavy * log(heavy) / log(2.0); - double total = total0; - auto store = [&](const int count) { + const double heavyCount = numHeavy; + const double total0 = heavyCount * log(heavyCount) / log(2.0); + double total = total0; + auto store = [&](const int count) { const double size = count; total -= size * log(size) / log(2.0); }; int head = 0; int count = 0; while (head < numHeavy) { - const RoundedLeverage front = roundLeverage(sorted[head], digits); + const ClusterKey front = unpackClusterKey(sortedKeys[head]); ++head; ++count; const int remaining = numHeavy - head; @@ -305,8 +354,7 @@ __device__ __forceinline__ void computeInformationIndices(const double* leverag store(count); } for (int offset = 0; offset < remaining; ++offset) { - const RoundedLeverage other = roundLeverage(sorted[head + offset], digits); - if (rdkitIsClose(front.fraction, other.fraction)) { + if (rdkitIsClose(front, unpackClusterKey(sortedKeys[head + offset]))) { ++count; } else { store(count); @@ -476,6 +524,7 @@ __global__ void getaway3DKernel(const DeviceCoordView coordinates, Real* __restrict__ rows, double* __restrict__ leverageScratch, double* __restrict__ sortScratch, + int64_t* __restrict__ keyScratch, Real* __restrict__ iterateScratch, uint8_t* __restrict__ bondDepths, Real* __restrict__ output) { @@ -652,16 +701,20 @@ __global__ void getaway3DKernel(const DeviceCoordView coordinates, eigenvalue = Real(0); } + double ith = 0; + double ish = 0; + int64_t* const keys = keyScratch + atomOffset * 2; + computeInformationIndices(leverages, + inputs.heavyAtomFlags + moleculeStart, + numAtoms, + precisionDigits, + lane, + sortScratch + atomOffset, + keys, + keys + numAtoms, + ith, + ish); if (lane == 0) { - double ith; - double ish; - computeInformationIndices(leverages, - inputs.heavyAtomFlags + moleculeStart, - numAtoms, - precisionDigits, - sortScratch + atomOffset, - ith, - ish); double leverageProduct = 1.0; for (int atomIdx = 0; atomIdx < numAtoms; ++atomIdx) { leverageProduct *= leverages[atomIdx]; @@ -706,7 +759,8 @@ __global__ void getaway3DKernel(const DeviceCoordView coordinates, /** * @brief Launches the GETAWAY kernel. Scratch per atom: one kGetawayRowStride row of @p Real, a float64 - * leverage and sort slot, two power-iteration values and one bond-depth byte per pair stream. + * leverage and rounded-leverage slot, two ITH/ISH sort keys, two power-iteration values and one + * bond-depth byte per pair stream. */ template void launchGetawayProperties(const DeviceCoordView& coordinates, @@ -722,6 +776,7 @@ void launchGetawayProperties(const DeviceCoordView& coordinates, AsyncDeviceVector rows(numAtoms * kGetawayRowStride, stream); AsyncDeviceVector leverages(numAtoms, stream); AsyncDeviceVector sorted(numAtoms, stream); + AsyncDeviceVector keys(numAtoms * 2, stream); AsyncDeviceVector iterates(numAtoms * 2, stream); AsyncDeviceVector depths(numAtoms * kGetawayPairStreams, stream); const int numBlocks = (numConformers + kWarpsPerBlock - 1) / kWarpsPerBlock; @@ -731,6 +786,7 @@ void launchGetawayProperties(const DeviceCoordView& coordinates, rows.data(), leverages.data(), sorted.data(), + keys.data(), iterates.data(), depths.data(), output); From 8d261816093bd5a750be0eda25060b2c938916bf Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Tue, 6 Oct 2026 17:55:35 -0400 Subject: [PATCH 6/8] Share one bond-distance table between AUTOCORR3D and GETAWAY 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. --- src/descriptors3d.cu | 10 +- src/descriptors3d_getaway.cuh | 59 ++++++++---- src/descriptors3d_pairwise.cuh | 62 +++++++++--- src/descriptors3d_topology.cuh | 166 +++++++++++++++++++++++++++++++++ tests/test_descriptors3d.cpp | 23 +++++ 5 files changed, 285 insertions(+), 35 deletions(-) create mode 100644 src/descriptors3d_topology.cuh diff --git a/src/descriptors3d.cu b/src/descriptors3d.cu index 06958bca..a34fa3c6 100644 --- a/src/descriptors3d.cu +++ b/src/descriptors3d.cu @@ -12,6 +12,7 @@ #include "src/descriptors3d_moments.cuh" #include "src/descriptors3d_pairwise.cuh" #include "src/descriptors3d_projection.cuh" +#include "src/descriptors3d_topology.cuh" #include "src/descriptors3d_usr.cuh" #include "src/utils/cuda_error_check.h" @@ -376,9 +377,14 @@ Property3DResults calc3DPropertiesGpu(const DeviceCoordView& coord const double* momentWeights = onlySpherocity ? nullptr : inputs.momentWeights; launchMomentProperties(coordinates, momentWeights, inputs.moleculeAtomStarts, work, separateSpherocityState, stream); launchProjectionProperties(coordinates, inputs, options.whim, pbfOutput, whimOutput, stream); - launchPairwiseProperties(coordinates, inputs, pairwiseOutputs, stream); + // Bond distances depend only on the molecule, so AUTOCORR3D and GETAWAY share one table for all conformers. + descriptors3d_detail::BondDistanceTableStorage bondTable; + if (pairwiseOutputs.autocorr3D != nullptr || getawayOutput != nullptr) { + descriptors3d_detail::buildBondDistanceTable(coordinates, inputs, bondTable, stream); + } + launchPairwiseProperties(coordinates, inputs, pairwiseOutputs, bondTable.table, stream); launchUsrProperties(coordinates, inputs, usrOutput, usrcatOutput, stream); - launchGetawayProperties(coordinates, inputs, options.getaway, getawayOutput, stream); + launchGetawayProperties(coordinates, inputs, options.getaway, bondTable.table, getawayOutput, stream); return results; } diff --git a/src/descriptors3d_getaway.cuh b/src/descriptors3d_getaway.cuh index f4d7dbdb..06c68d6c 100644 --- a/src/descriptors3d_getaway.cuh +++ b/src/descriptors3d_getaway.cuh @@ -10,6 +10,7 @@ #include "src/descriptors3d.h" #include "src/descriptors3d_kernel.cuh" #include "src/descriptors3d_projection.cuh" +#include "src/descriptors3d_topology.cuh" #include "src/utils/cuda_error_check.h" #include "src/utils/device_vector.h" #include "src/utils/symmetric_eigenvalues_3x3.cuh" @@ -53,6 +54,11 @@ constexpr int kGetawayRowWeights = 5; constexpr int kGetawayRowCentered = kGetawayRowWeights + kNumAtomPropertyChannels; constexpr int kGetawayRowStride = kGetawayRowCentered + 3; +// ITH and ISH group the heavy-atom leverages by their decimal digits, not their numeric values: RDKit rounds +// each leverage to `precision` significant digits (round_to_n_digits_), prints it with C++ stream output, and +// its IsClose2 compares substrings of that printed text. Reproducing RDKit's groups exactly therefore means +// reproducing that rounding and that text, which the helpers below do for the values leverages can take. + //! Exact powers of ten representable in float64. __device__ __forceinline__ double exactPowerOfTen(const int exponent) { constexpr double kPowers[] = {1e0, 1e1, 1e2, 1e3, 1e4, 1e5, 1e6, 1e7, 1e8, 1e9, 1e10, 1e11, @@ -77,7 +83,7 @@ __device__ __forceinline__ double scaleByPowerOfTen(double value, int exponent) * @brief RDKit's round_to_n_digits_: `atof(sprintf("%.*g", digits, value))` for a non-negative @p value. * * Also returns the rounded value's significand (@p mantissa, exactly @p digits digits unless the value is 0) - * and decimal exponent, from which formatLikeOstream() rebuilds its text. Exact ties round to even, as glibc + * and decimal exponent, from which rdkitDecimalText() rebuilds its printed text. Exact ties round to even, as glibc * does for exactly representable halves; values within float64 rounding of a tie may differ from printf. */ __device__ __forceinline__ double roundToSignificantDigits(const double value, @@ -108,7 +114,7 @@ __device__ __forceinline__ double roundToSignificantDigits(const double value, return scaleByPowerOfTen(static_cast(mantissa), exponent - (digits - 1)); } -//! Short text buffer for RDKit's string comparisons. +//! Decimal text of a printed value, as RDKit's IsClose2 compares it. struct DecimalText { char chars[24]; int length; @@ -119,14 +125,17 @@ __device__ __forceinline__ void appendChar(DecimalText& text, const char charact } /** - * @brief The text `std::ostream << v` (default precision 6, %g style, trailing zeros removed) prints for - * v = @p mantissa * 10^(@p exponent - @p digits + 1), a value from roundToSignificantDigits() with - * @p digits <= 6 significant digits. + * @brief The decimal text RDKit's IsClose2 compares for the rounded leverage + * v = @p mantissa * 10^(@p exponent - @p digits + 1) from roundToSignificantDigits(). + * + * IsClose2 prints with default C++ stream formatting: up to 6 significant digits with trailing zeros removed, + * fixed notation for decimal exponents -4 to 5 (e.g. "0.034", "1") and scientific notation otherwise + * (e.g. "1.8e-06"). With @p digits <= 6 every significant digit of v is printed. */ -__device__ __forceinline__ void formatLikeOstream(const int64_t mantissa, - const int exponent, - const int digits, - DecimalText& text) { +__device__ __forceinline__ void rdkitDecimalText(const int64_t mantissa, + const int exponent, + const int digits, + DecimalText& text) { text.length = 0; if (mantissa == 0) { appendChar(text, '0'); @@ -250,7 +259,7 @@ __device__ __forceinline__ ClusterKey clusterKey(const double leverage, const in int exponent; rounded = roundToSignificantDigits(leverage, digits, mantissa, exponent); DecimalText printed; - formatLikeOstream(mantissa, exponent, digits, printed); + rdkitDecimalText(mantissa, exponent, digits, printed); const DecimalText fraction = textAfterPoint(printed); ClusterKey key{countLeadingZeros(fraction), fraction.length, leadingInteger(fraction), true}; for (int position = 0; position < fraction.length; ++position) { @@ -527,6 +536,7 @@ __global__ void getaway3DKernel(const DeviceCoordView coordinates, int64_t* __restrict__ keyScratch, Real* __restrict__ iterateScratch, uint8_t* __restrict__ bondDepths, + const BondDistanceTable bondTable, Real* __restrict__ output) { const int lane = static_cast(threadIdx.x) % kWarpSize; const int laneInGroup = lane % kGroupSize; @@ -602,18 +612,27 @@ __global__ void getaway3DKernel(const DeviceCoordView coordinates, const int32_t* neighborStarts = inputs.bondNeighborStarts + moleculeStart; uint8_t* depth = bondDepths + atomOffset * kGetawayPairStreams + static_cast(stream) * numAtoms; const unsigned groupMask = 0xffu << (stream * kGroupSize); + const uint8_t* moleculeTable = moleculeBondDistances(bondTable, moleculeIdx); for (int j = stream; j < numAtoms - 1; j += kGetawayPairStreams) { - searchBondDepths(neighborStarts, - inputs.bondNeighbors, - numAtoms, - j, - laneInGroup, - groupMask, - depth); + // Row j's bond distances, indexed by k - lagBase: the molecule's table row, or a search into depth. + const uint8_t* lags = depth; + int lagBase = 0; + if (moleculeTable != nullptr) { + lags = moleculeTable + bondTableRowStart(j, numAtoms); + lagBase = j + 1; + } else { + searchBondDepths(neighborStarts, + inputs.bondNeighbors, + numAtoms, + j, + laneInGroup, + groupMask, + depth); + } const Real* rowJ = atomRows + j * kGetawayRowStride; for (int k = j + 1; k < numAtoms; ++k) { - const int lag = depth[k]; - const bool near = lag != kUnreachedDepth; + const int lag = lags[k - lagBase]; + const bool near = lag <= kNumGetawayLags; const bool owner = near ? lag - 1 == laneInGroup : k % kGroupSize == laneInGroup; if (!owner) { continue; @@ -766,6 +785,7 @@ template void launchGetawayProperties(const DeviceCoordView& coordinates, const Property3DDeviceInputs& inputs, const GetawayOptions& options, + const BondDistanceTable& bondTable, Real* output, const cudaStream_t stream) { const int numConformers = coordinates.numConformers; @@ -789,6 +809,7 @@ void launchGetawayProperties(const DeviceCoordView& coordinates, keys.data(), iterates.data(), depths.data(), + bondTable, output); cudaCheckError(cudaGetLastError()); } diff --git a/src/descriptors3d_pairwise.cuh b/src/descriptors3d_pairwise.cuh index 9235f11b..e43d205e 100644 --- a/src/descriptors3d_pairwise.cuh +++ b/src/descriptors3d_pairwise.cuh @@ -9,6 +9,7 @@ #include "src/descriptors3d.h" #include "src/descriptors3d_kernel.cuh" +#include "src/descriptors3d_topology.cuh" #include "src/utils/cuda_error_check.h" #include "src/utils/device_vector.h" @@ -156,6 +157,7 @@ __global__ void pairwise3DKernel(const DeviceCoordView coordinates, const Property3DDeviceInputs inputs, Real* __restrict__ scratch, uint8_t* __restrict__ bondDepths, + const BondDistanceTable bondTable, Real* __restrict__ rdfOutput, Real* __restrict__ morseOutput, Real* __restrict__ autocorrOutput) { @@ -219,6 +221,8 @@ __global__ void pairwise3DKernel(const DeviceCoordView coordinates, const int32_t* neighborStarts = nullptr; uint8_t* depth = nullptr; + const uint8_t* moleculeTable = + kAutocorr ? moleculeBondDistances(bondTable, coordinates.molIndices[conformerIdx]) : nullptr; if constexpr (kAutocorr) { neighborStarts = inputs.bondNeighborStarts + moleculeStart; depth = bondDepths + ((atoms.positions - coordinates.positions) / 3) * kPairStreams + @@ -247,14 +251,22 @@ __global__ void pairwise3DKernel(const DeviceCoordView coordinates, const PairwiseAtom first = loadPairwiseAtom(rowJ); const Real dragJ = rowJ[kScratchIStateDrag]; const Real radiusJ = rowJ[kScratchCovalentRadius]; + // Row j's bond distances, indexed by k - lagBase: the molecule's table row, or a search into depth. + const uint8_t* lags = depth; + int lagBase = 0; if constexpr (kAutocorr) { - searchBondDepths(neighborStarts, - inputs.bondNeighbors, - atoms.numAtoms, - j, - laneInGroup, - 0xffu << (stream * kGroupSize), - depth); + if (moleculeTable != nullptr) { + lags = moleculeTable + bondTableRowStart(j, atoms.numAtoms); + lagBase = j + 1; + } else { + searchBondDepths(neighborStarts, + inputs.bondNeighbors, + atoms.numAtoms, + j, + laneInGroup, + 0xffu << (stream * kGroupSize), + depth); + } } for (int k = j + 1; k < atoms.numAtoms; ++k) { const Real* const rowK = rows + k * kPairwiseScratchStride; @@ -283,7 +295,7 @@ __global__ void pairwise3DKernel(const DeviceCoordView coordinates, } } if constexpr (kAutocorr) { - const int lag = depth[k]; + const int lag = lags[k - lagBase]; if (lag != kUnreachedDepth && (lag - 1) % kGroupSize == laneInGroup) { // A select on the slot, not an index, keeps autocorrAcc in registers. const int lagSlot = (lag - 1) / kGroupSize; @@ -382,6 +394,7 @@ void launchPairwiseSet(const DeviceCoordView& coordinates, const Property3DDeviceInputs& inputs, Real* scratch, uint8_t* bondDepths, + const BondDistanceTable& bondTable, const PairwiseOutputs& outputs, const cudaStream_t stream) { const int numBlocks = (coordinates.numConformers + kWarpsPerBlock - 1) / kWarpsPerBlock; @@ -389,6 +402,7 @@ void launchPairwiseSet(const DeviceCoordView& coordinates, inputs, scratch, bondDepths, + bondTable, outputs.rdf, outputs.morse, outputs.autocorr3D); @@ -401,6 +415,7 @@ template void launchPairwiseProperties(const DeviceCoordView& coordinates, const Property3DDeviceInputs& inputs, const PairwiseOutputs& outputs, + const BondDistanceTable& bondTable, const cudaStream_t stream) { const unsigned set = (outputs.rdf != nullptr ? kPairwiseRdf : 0u) | (outputs.morse != nullptr ? kPairwiseMorse : 0u) | (outputs.autocorr3D != nullptr ? kPairwiseAutocorr3D : 0u); @@ -416,28 +431,47 @@ void launchPairwiseProperties(const DeviceCoordView& coordinates, uint8_t* const depths = bondDepthScratch.data(); switch (set) { case kPairwiseRdf: - launchPairwiseSet(coordinates, inputs, rows, depths, outputs, stream); + launchPairwiseSet(coordinates, inputs, rows, depths, bondTable, outputs, stream); break; case kPairwiseMorse: - launchPairwiseSet(coordinates, inputs, rows, depths, outputs, stream); + launchPairwiseSet(coordinates, inputs, rows, depths, bondTable, outputs, stream); break; case kPairwiseRdf | kPairwiseMorse: - launchPairwiseSet(coordinates, inputs, rows, depths, outputs, stream); + launchPairwiseSet(coordinates, + inputs, + rows, + depths, + bondTable, + outputs, + stream); break; case kPairwiseAutocorr3D: - launchPairwiseSet(coordinates, inputs, rows, depths, outputs, stream); + launchPairwiseSet(coordinates, inputs, rows, depths, bondTable, outputs, stream); break; case kPairwiseRdf | kPairwiseAutocorr3D: - launchPairwiseSet(coordinates, inputs, rows, depths, outputs, stream); + launchPairwiseSet(coordinates, + inputs, + rows, + depths, + bondTable, + outputs, + stream); break; case kPairwiseMorse | kPairwiseAutocorr3D: - launchPairwiseSet(coordinates, inputs, rows, depths, outputs, stream); + launchPairwiseSet(coordinates, + inputs, + rows, + depths, + bondTable, + outputs, + stream); break; default: launchPairwiseSet(coordinates, inputs, rows, depths, + bondTable, outputs, stream); break; diff --git a/src/descriptors3d_topology.cuh b/src/descriptors3d_topology.cuh new file mode 100644 index 00000000..064ec895 --- /dev/null +++ b/src/descriptors3d_topology.cuh @@ -0,0 +1,166 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#ifndef NVMOLKIT_DESCRIPTORS3D_TOPOLOGY_CUH +#define NVMOLKIT_DESCRIPTORS3D_TOPOLOGY_CUH + +#include +#include +#include + +#include "src/descriptors3d.h" +#include "src/descriptors3d_kernel.cuh" +#include "src/utils/cuda_error_check.h" +#include "src/utils/device_vector.h" + +namespace nvMolKit::descriptors3d_detail { + +//! Largest bond-count distance the table records; AUTOCORR3D reads up to 10, GETAWAY up to 8. +constexpr int kBondTableDepth = 10; +//! Upper bound on the bond-distance table; molecules past it fall back to per-row searches. +constexpr size_t kBondTableBudgetBytes = size_t{256} << 20; +//! Largest molecule whose per-row search arrays (one per group of a block) fit the default shared memory. +constexpr int kBondTableMaxAtoms = (48 << 10) / (kWarpsPerBlock * kGroupsPerWarp); + +/** + * @brief Per-molecule bond-count distances, shared by every conformer of a molecule and by AUTOCORR3D and + * GETAWAY. Molecule m's pairs (j < k) are stored row-major from `starts[m]`, one byte each: the + * distance up to kBondTableDepth, or kUnreachedDepth beyond it or between fragments. Molecules whose + * range ends past @ref capacity have no table entries; their kernels search per row instead. + */ +struct BondDistanceTable { + const uint8_t* distances = nullptr; + const int64_t* starts = nullptr; + int64_t capacity = 0; +}; + +//! Table entries of molecule @p moleculeIdx, or null when it has none. +__device__ __forceinline__ const uint8_t* moleculeBondDistances(const BondDistanceTable& table, const int moleculeIdx) { + if (table.distances == nullptr || table.starts[moleculeIdx + 1] > table.capacity) { + return nullptr; + } + return table.distances + table.starts[moleculeIdx]; +} + +//! Offset of pair (@p j, @p j + 1) within a molecule's table entries for @p numAtoms atoms. +__device__ __forceinline__ int64_t bondTableRowStart(const int j, const int numAtoms) { + return static_cast(j) * numAtoms - static_cast(j) * (j + 1) / 2; +} + +//! Pair count n (n - 1) / 2 of each molecule into `pairCounts[m + 1]`, and 0 into `pairCounts[0]`. +__global__ void moleculePairCountsKernel(const int32_t* moleculeAtomStarts, const int nMols, int64_t* pairCounts) { + const int moleculeIdx = blockIdx.x * blockDim.x + threadIdx.x; + if (moleculeIdx == 0) { + pairCounts[0] = 0; + } + if (moleculeIdx < nMols) { + const int64_t numAtoms = moleculeAtomStarts[moleculeIdx + 1] - moleculeAtomStarts[moleculeIdx]; + pairCounts[moleculeIdx + 1] = numAtoms * (numAtoms - 1) / 2; + } +} + +/** + * @brief Fills the table entries of every molecule that fits @p table's capacity; one warp per molecule. + * Each group searches interleaved rows j with searchBondDepths() into its own shared-memory row of + * @p maxAtoms bytes and writes the row's pairs (j, k > j). + */ +__global__ void bondDistanceTableKernel(const int32_t* moleculeAtomStarts, + const int32_t* bondNeighborStarts, + const int32_t* bondNeighbors, + const int nMols, + const int maxAtoms, + const int64_t* starts, + const int64_t capacity, + uint8_t* __restrict__ distances) { + extern __shared__ uint8_t searchRows[]; + const int lane = static_cast(threadIdx.x) % kWarpSize; + const int laneInGroup = lane % kGroupSize; + const int group = lane / kGroupSize; + const int warpInBlock = static_cast(threadIdx.x) / kWarpSize; + const int moleculeIdx = blockIdx.x * kWarpsPerBlock + warpInBlock; + if (moleculeIdx >= nMols || starts[moleculeIdx + 1] > capacity) { + return; // Uniform across the warp. + } + const int atomStart = moleculeAtomStarts[moleculeIdx]; + const int numAtoms = moleculeAtomStarts[moleculeIdx + 1] - atomStart; + uint8_t* const depth = searchRows + (warpInBlock * kGroupsPerWarp + group) * maxAtoms; + uint8_t* const entries = distances + starts[moleculeIdx]; + const unsigned groupMask = 0xffu << (group * kGroupSize); + for (int j = group; j < numAtoms - 1; j += kGroupsPerWarp) { + searchBondDepths(bondNeighborStarts + atomStart, + bondNeighbors, + numAtoms, + j, + laneInGroup, + groupMask, + depth); + uint8_t* const row = entries + bondTableRowStart(j, numAtoms); + for (int k = j + 1 + laneInGroup; k < numAtoms; k += kGroupSize) { + row[k - j - 1] = depth[k]; + } + } +} + +//! Owns a BondDistanceTable's device storage. +struct BondDistanceTableStorage { + AsyncDeviceVector starts; + AsyncDeviceVector distances; + BondDistanceTable table; +}; + +/** + * @brief Builds the bond-distance table for the batch's molecules on @p stream, without host + * synchronization. Capacity is the smaller of kBondTableBudgetBytes and the worst case for + * `inputs.maxMoleculeAtoms`; batches whose largest molecule exceeds kBondTableMaxAtoms (or whose + * maxMoleculeAtoms is unset) get an empty table, so every conformer searches per row. + */ +inline void buildBondDistanceTable(const DeviceCoordView& coordinates, + const Property3DDeviceInputs& inputs, + BondDistanceTableStorage& storage, + const cudaStream_t stream) { + const int nMols = coordinates.nMols; + const int maxAtoms = inputs.maxMoleculeAtoms; + if (nMols == 0 || maxAtoms < 2 || maxAtoms > kBondTableMaxAtoms) { + return; + } + const int64_t maxPairs = static_cast(maxAtoms) * (maxAtoms - 1) / 2; + const int64_t capacity = std::min(kBondTableBudgetBytes, maxPairs * nMols); + + storage.starts = AsyncDeviceVector(static_cast(nMols) + 1, stream); + moleculePairCountsKernel<<<(nMols + kBlockSize - 1) / kBlockSize, kBlockSize, 0, stream>>>(inputs.moleculeAtomStarts, + nMols, + storage.starts.data()); + cudaCheckError(cudaGetLastError()); + size_t tempBytes = 0; + cudaCheckError(cub::DeviceScan::InclusiveSum(nullptr, + tempBytes, + storage.starts.data() + 1, + storage.starts.data() + 1, + nMols, + stream)); + AsyncDeviceVector temp(tempBytes, stream); + cudaCheckError(cub::DeviceScan::InclusiveSum(temp.data(), + tempBytes, + storage.starts.data() + 1, + storage.starts.data() + 1, + nMols, + stream)); + + storage.distances = AsyncDeviceVector(static_cast(capacity), stream); + const int numBlocks = (nMols + kWarpsPerBlock - 1) / kWarpsPerBlock; + const size_t sharedBytes = static_cast(kWarpsPerBlock) * kGroupsPerWarp * maxAtoms; + bondDistanceTableKernel<<>>(inputs.moleculeAtomStarts, + inputs.bondNeighborStarts, + inputs.bondNeighbors, + nMols, + maxAtoms, + storage.starts.data(), + capacity, + storage.distances.data()); + cudaCheckError(cudaGetLastError()); + storage.table = {storage.distances.data(), storage.starts.data(), capacity}; +} + +} // namespace nvMolKit::descriptors3d_detail + +#endif // NVMOLKIT_DESCRIPTORS3D_TOPOLOGY_CUH diff --git a/tests/test_descriptors3d.cpp b/tests/test_descriptors3d.cpp index 65b786c5..cc8cf998 100644 --- a/tests/test_descriptors3d.cpp +++ b/tests/test_descriptors3d.cpp @@ -20,6 +20,7 @@ #include #include #include +#include #include #include "src/conformer/conformer_coord_upload.h" @@ -388,6 +389,28 @@ TEST(Descriptors3DGetaway, MatchesRdkitGetaway) { } } +TEST(Descriptors3DPairwise, BondSearchFallbackMatchesBondDistanceTable) { + // A 12-atom chain has bond distances up to 11, past the deepest lag either descriptor uses. + constexpr int kChainAtoms = 12; + std::vector> conformers(2); + for (int atomIdx = 0; atomIdx < kChainAtoms; ++atomIdx) { + conformers[0].push_back({1.25 * atomIdx, 0.8 * (atomIdx % 2), 0.3 * ((atomIdx / 2) % 2)}); + conformers[1].push_back({1.1 * atomIdx, 0.9 * ((atomIdx / 3) % 2), 0.5 * (atomIdx % 2)}); + } + auto small = molWithConformers("CCCCCCCCCCCO", conformers); + // More atoms than the bond-distance table supports (kBondTableMaxAtoms, 3072), so every molecule batched with + // this chain searches its bond distances per row instead. Without conformers it adds no results. + auto large = molWithConformers(std::string(3100, 'C').c_str(), {}); + + const std::vector properties = {Property3D::AUTOCORR3D, Property3D::GETAWAY}; + auto withTable = nvMolKit::calc3DProperties({small.get()}, properties, {}, nullptr); + auto searched = nvMolKit::calc3DProperties({small.get(), large.get()}, properties, {}, nullptr); + for (const Property3D property : properties) { + EXPECT_EQ(toHost(searched.properties.at(property)), toHost(withTable.properties.at(property))) + << nvMolKit::property3DName(property); + } +} + TEST(Descriptors3DUsr, MatchesRdkitUsrAndUsrcat) { auto mol = molWithConformers("CCCO", { From 9a8ae184b163c3b96ded7bd032de222296b5228e Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Tue, 6 Oct 2026 18:16:42 -0400 Subject: [PATCH 7/8] Shorten GETAWAY register live ranges 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). --- src/descriptors3d_getaway.cuh | 103 +++++++++++++++++----------------- 1 file changed, 53 insertions(+), 50 deletions(-) diff --git a/src/descriptors3d_getaway.cuh b/src/descriptors3d_getaway.cuh index 06c68d6c..8cd6dd9b 100644 --- a/src/descriptors3d_getaway.cuh +++ b/src/descriptors3d_getaway.cuh @@ -577,29 +577,6 @@ __global__ void getaway3DKernel(const DeviceCoordView coordinates, const bool is3D = inputs.conformerIs3D == nullptr || inputs.conformerIs3D[defaultRow] != 0; const double dimension = (defaultAtoms.numAtoms < 4 || !is3D ? 0.0 : pbf) < kGetawayPlanarPbf ? 2.0 : 3.0; - // Lag 0: H0 = sum w^2 h and HATS0 = sum w^2 h^2 over atoms with h > 0; HIC from every leverage. - Real lagZeroH[kNumGetawayChannels] = {}; - Real lagZeroHats[kNumGetawayChannels] = {}; - double hic = 0; - for (int atomIdx = lane; atomIdx < numAtoms; atomIdx += kWarpSize) { - const Real* atomRow = atomRows + atomIdx * kGetawayRowStride; - const Real leverage = atomRow[kGetawayRowLeverage]; - const double exact = leverages[atomIdx]; - hic -= exact / dimension * log(exact / dimension) / log(2.0); - if (exact > 0.0) { - for (int channel = 0; channel < kNumGetawayChannels; ++channel) { - const Real weight = getawayWeight(atomRow, channel); - lagZeroH[channel] += weight * weight * leverage; - lagZeroHats[channel] += weight * weight * leverage * leverage; - } - } - } - hic = groupAllReduceSum(hic); - for (int channel = 0; channel < kNumGetawayChannels; ++channel) { - lagZeroH[channel] = groupAllReduceSum(lagZeroH[channel]); - lagZeroHats[channel] = groupAllReduceSum(lagZeroHats[channel]); - } - // Pair loop: lane l owns lag l + 1; pairs beyond kNumGetawayLags (or disconnected) go to lane k % kGroupSize. Real lagH[kNumGetawayChannels] = {}; Real lagHats[kNumGetawayChannels] = {}; @@ -676,6 +653,59 @@ __global__ void getaway3DKernel(const DeviceCoordView coordinates, farR[channel] = sumAcrossGroups(farR[channel]); } + // Lag 0: H0 = sum w^2 h and HATS0 = sum w^2 h^2 over atoms with h > 0; HIC from every leverage. Computed + // after the pair loop, and the H and R blocks written before REIG, so neither set of sums holds registers + // through the other's loop (168 -> 128 registers in float32). + Real lagZeroH[kNumGetawayChannels] = {}; + Real lagZeroHats[kNumGetawayChannels] = {}; + double hic = 0; + for (int atomIdx = lane; atomIdx < numAtoms; atomIdx += kWarpSize) { + const Real* atomRow = atomRows + atomIdx * kGetawayRowStride; + const Real leverage = atomRow[kGetawayRowLeverage]; + const double exact = leverages[atomIdx]; + hic -= exact / dimension * log(exact / dimension) / log(2.0); + if (exact > 0.0) { + for (int channel = 0; channel < kNumGetawayChannels; ++channel) { + const Real weight = getawayWeight(atomRow, channel); + lagZeroH[channel] += weight * weight * leverage; + lagZeroHats[channel] += weight * weight * leverage * leverage; + } + } + } + hic = groupAllReduceSum(hic); + for (int channel = 0; channel < kNumGetawayChannels; ++channel) { + lagZeroH[channel] = groupAllReduceSum(lagZeroH[channel]); + lagZeroHats[channel] = groupAllReduceSum(lagZeroHats[channel]); + } + + if (stream == 0) { + for (int channel = 0; channel < kNumGetawayChannels; ++channel) { + Real* const hBlock = row + kGetawayHStart + channel * kGetawayHBlock; + Real* const rBlock = row + kGetawayRStart + channel * kGetawayRBlock; + const Real lagHSum = groupAllReduceSum(lagH[channel]); + const Real lagHatsSum = groupAllReduceSum(lagHats[channel]); + const Real lagRSum = groupAllReduceSum(lagR[channel]); + const Real farHSum = groupAllReduceSum(farH[channel]); + const Real farHatsSum = groupAllReduceSum(farHats[channel]); + const Real farRSum = groupAllReduceSum(farR[channel]); + const Real rMax = + cooperative_groups::reduce(laneTile(), lagRMax[channel], cooperative_groups::greater()); + const int lagIdx = laneInGroup + 1; + hBlock[lagIdx] = roundThousandths(lagH[channel]); + hBlock[kNumGetawayLags + 2 + lagIdx] = roundThousandths(lagHats[channel]); + rBlock[lagIdx - 1] = roundThousandths(lagR[channel]); + rBlock[kNumGetawayLags + lagIdx] = roundThousandths(lagRMax[channel]); + if (laneInGroup == 0) { + hBlock[0] = roundThousandths(lagZeroH[channel]); + hBlock[kNumGetawayLags + 1] = roundThousandths(lagZeroH[channel] + Real(2) * (lagHSum + farHSum)); + hBlock[kNumGetawayLags + 2] = roundThousandths(lagZeroHats[channel]); + hBlock[kGetawayHBlock - 1] = roundThousandths(lagZeroHats[channel] + Real(2) * (lagHatsSum + farHatsSum)); + rBlock[kNumGetawayLags] = roundThousandths(Real(2) * (lagRSum + farRSum)); + rBlock[kGetawayRBlock - 1] = roundThousandths(rMax); + } + } + } + // REIG by power iteration on R; the first product with ones gives R's row sums for RARS and RCON. Real* const current = iterateScratch + atomOffset * 2; Real* const next = current + numAtoms; @@ -747,33 +777,6 @@ __global__ void getaway3DKernel(const DeviceCoordView coordinates, row[kGetawayRconIndex + 1] = roundThousandths(rars); row[kGetawayRconIndex + 2] = roundThousandths(eigenvalue); } - if (stream == 0) { - for (int channel = 0; channel < kNumGetawayChannels; ++channel) { - Real* const hBlock = row + kGetawayHStart + channel * kGetawayHBlock; - Real* const rBlock = row + kGetawayRStart + channel * kGetawayRBlock; - const Real lagHSum = groupAllReduceSum(lagH[channel]); - const Real lagHatsSum = groupAllReduceSum(lagHats[channel]); - const Real lagRSum = groupAllReduceSum(lagR[channel]); - const Real farHSum = groupAllReduceSum(farH[channel]); - const Real farHatsSum = groupAllReduceSum(farHats[channel]); - const Real farRSum = groupAllReduceSum(farR[channel]); - const Real rMax = - cooperative_groups::reduce(laneTile(), lagRMax[channel], cooperative_groups::greater()); - const int lagIdx = laneInGroup + 1; - hBlock[lagIdx] = roundThousandths(lagH[channel]); - hBlock[kNumGetawayLags + 2 + lagIdx] = roundThousandths(lagHats[channel]); - rBlock[lagIdx - 1] = roundThousandths(lagR[channel]); - rBlock[kNumGetawayLags + lagIdx] = roundThousandths(lagRMax[channel]); - if (laneInGroup == 0) { - hBlock[0] = roundThousandths(lagZeroH[channel]); - hBlock[kNumGetawayLags + 1] = roundThousandths(lagZeroH[channel] + Real(2) * (lagHSum + farHSum)); - hBlock[kNumGetawayLags + 2] = roundThousandths(lagZeroHats[channel]); - hBlock[kGetawayHBlock - 1] = roundThousandths(lagZeroHats[channel] + Real(2) * (lagHatsSum + farHatsSum)); - rBlock[kNumGetawayLags] = roundThousandths(Real(2) * (lagRSum + farRSum)); - rBlock[kGetawayRBlock - 1] = roundThousandths(rMax); - } - } - } } /** From 608cc5cac04750fc0decf75b18a3a7d04561093e Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Wed, 7 Oct 2026 15:53:49 -0400 Subject: [PATCH 8/8] Address PR 358 review comments - 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. --- nvmolkit/descriptors3d.py | 19 ++++-- nvmolkit/tests/test_descriptors3d.py | 19 ++++++ src/descriptors3d.h | 8 ++- src/descriptors3d_getaway.cuh | 19 ++++-- src/descriptors3d_mol.cpp | 4 +- src/descriptors3d_topology.cuh | 8 ++- tests/CMakeLists.txt | 5 ++ tests/test_descriptors3d_topology.cu | 97 ++++++++++++++++++++++++++++ 8 files changed, 161 insertions(+), 18 deletions(-) create mode 100644 tests/test_descriptors3d_topology.cu diff --git a/nvmolkit/descriptors3d.py b/nvmolkit/descriptors3d.py index 9679ed36..235201e4 100644 --- a/nvmolkit/descriptors3d.py +++ b/nvmolkit/descriptors3d.py @@ -322,7 +322,10 @@ def Calc3DProperties( ``values``, or whose atom count differs from their molecule's produce NaN rather than an error, so no host synchronization is needed. Device coordinate rows are treated as three-dimensional; - molecule conformers preserve their RDKit ``is3D`` flag for PBF. + molecule conformers keep their RDKit ``is3D`` flag, which PBF and + GETAWAY's HIC read. HIC takes its flatness (PBF and ``is3D``) from + each molecule's first conformer for all of its conformers; with + ``coordinates``, each row uses its own geometry. options: Per-family options; defaults to :class:`Property3DOptions` (RDKit's defaults). precision: ``PrecisionMode.SINGLE`` (default) computes and returns @@ -396,11 +399,13 @@ def Calc3DProperties( options = Property3DOptions() elif not isinstance(options, Property3DOptions): raise TypeError(f"options must be a Property3DOptions or None, got {type(options).__name__}") - getaway_precision = options.getaway.precision - if isinstance(getaway_precision, bool) or not isinstance(getaway_precision, int): - raise TypeError(f"GETAWAY precision must be an int, got {type(getaway_precision).__name__}") - if not 1 <= getaway_precision <= 6: - raise ValueError(f"GETAWAY precision must be between 1 and 6 significant digits, got {getaway_precision}") + getaway_precision = GetawayOptions().precision + if Property3D.GETAWAY in normalized_properties: + getaway_precision = options.getaway.precision + if isinstance(getaway_precision, bool) or not isinstance(getaway_precision, int): + raise TypeError(f"GETAWAY precision must be an int, got {type(getaway_precision).__name__}") + if not 1 <= getaway_precision <= 6: + raise ValueError(f"GETAWAY precision must be between 1 and 6 significant digits, got {getaway_precision}") if hardwareOptions is None: hardwareOptions = HardwareOptions() elif not isinstance(hardwareOptions, HardwareOptions): @@ -420,7 +425,7 @@ def Calc3DProperties( [prop.value for prop in normalized_properties], options.moments.useAtomicMasses, options.whim.threshold, - options.getaway.precision, + getaway_precision, coordinate_interfaces, precision, hardwareOptions.preprocessingThreads, diff --git a/nvmolkit/tests/test_descriptors3d.py b/nvmolkit/tests/test_descriptors3d.py index 915a12be..f9f6b79b 100644 --- a/nvmolkit/tests/test_descriptors3d.py +++ b/nvmolkit/tests/test_descriptors3d.py @@ -643,6 +643,25 @@ def test_getaway_rejects_non_integer_precision(): Calc3DProperties([_embed("CCO", 1, 5)], Property3D.GETAWAY, options=options) +@pytest.mark.parametrize("precision", PRECISIONS) +def test_getaway_non_finite_coordinates_give_nan_information_indices(precision): + mols = [ + _mol_with_conformers("CCCO", [[(0.0, 0.0, 0.0), (1.5, float("nan"), 0.0), (2.0, 1.4, 0.0), (3.4, 1.5, 0.2)]]), + _mol_with_conformers("CCCO", [[(0.0, 0.0, 0.0), (1.5, 0.2, 0.0), (2.0, 1.4, float("inf")), (3.4, 1.5, 0.2)]]), + _embed("CCCO", 1, 47), + ] + values = Calc3DProperties(mols, Property3D.GETAWAY, precision=precision)[Property3D.GETAWAY].numpy() + assert np.isnan(values[:2, :2]).all() # ITH and ISH + _assert_rounded_matches_rdkit(values[2:], _rdkit_getaway_rows(mols[2:]), precision) + + +@pytest.mark.parametrize("precision_digits", [0, 2.5]) +def test_getaway_precision_ignored_when_not_requested(precision_digits): + options = Property3DOptions(getaway=GetawayOptions(precision=precision_digits)) + result = Calc3DProperties([_embed("CCO", 1, 5)], Property3D.PMI1, options=options) + assert result[Property3D.PMI1].torch().shape == (1,) + + @pytest.mark.parametrize("precision", PRECISIONS) def test_autocorr3d_accepts_batches_without_bonds(precision): mols = [ diff --git a/src/descriptors3d.h b/src/descriptors3d.h index c5ee7197..0c426c67 100644 --- a/src/descriptors3d.h +++ b/src/descriptors3d.h @@ -199,9 +199,13 @@ struct Property3DDeviceInputs { const int32_t* defaultConformerRows = nullptr; //! PBF, GETAWAY: per-conformer RDKit is3D flags (one per coordinate row); null treats every row as 3D. const int8_t* conformerIs3D = nullptr; - //! WHIM: largest molecule atom count in the batch (host value); sizes the per-conformer symmetry-search - //! scratch. Rows with more atoms produce NaN. + //! WHIM, AUTOCORR3D, GETAWAY: largest molecule atom count in the batch (host value). Sizes WHIM's + //! per-conformer symmetry-search scratch (rows with more atoms produce NaN) and, with 0 or above 3072 atoms, + //! leaves AUTOCORR3D and GETAWAY without the shared bond-distance table. int32_t maxMoleculeAtoms = 0; + //! AUTOCORR3D, GETAWAY: atom pairs n (n - 1) / 2 summed over the batch's molecules (host value); sizes the + //! shared bond-distance table. 0 sizes it as if every molecule had @ref maxMoleculeAtoms atoms. + int64_t moleculeAtomPairs = 0; }; //! One row-major device vector of length numConformers * property3DWidth(property) per property. diff --git a/src/descriptors3d_getaway.cuh b/src/descriptors3d_getaway.cuh index 8cd6dd9b..77462b5c 100644 --- a/src/descriptors3d_getaway.cuh +++ b/src/descriptors3d_getaway.cuh @@ -92,8 +92,8 @@ __device__ __forceinline__ double roundToSignificantDigits(const double value, int& exponent) { mantissa = 0; exponent = 0; - if (!(value > 0.0)) { - return value; // 0 stays 0; NaN stays NaN. + if (!(value > 0.0) || !isfinite(value)) { + return value; // 0 stays 0; NaN and infinity stay as they are. } exponent = static_cast(floor(log10(value))); const double lowest = exactPowerOfTen(digits - 1); @@ -304,7 +304,8 @@ __device__ __forceinline__ bool rdkitIsClose(ClusterKey a, ClusterKey b) { * * The lanes round, key and rank the heavy atoms (equal values print identically, so their relative order * does not matter); lane 0 then runs the greedy scan on the sorted keys and alone receives @p ith and @p ish. - * @p heavyValues and @p unsortedKeys hold one entry and @p sortedKeys one entry per heavy atom. + * Both are NaN when any heavy-atom leverage is not finite. @p heavyValues and @p unsortedKeys hold one entry and @p + * sortedKeys one entry per heavy atom. */ __device__ __forceinline__ void computeInformationIndices(const double* leverages, const uint8_t* heavyAtomFlags, @@ -316,12 +317,14 @@ __device__ __forceinline__ void computeInformationIndices(const double* leverag int64_t* sortedKeys, double& ith, double& ish) { - const auto warp = laneTile(); - int numHeavy = 0; + const auto warp = laneTile(); + int numHeavy = 0; + bool nonFinite = false; for (int base = 0; base < numAtoms; base += kWarpSize) { const int atomIdx = base + lane; const bool isHeavy = atomIdx < numAtoms && heavyAtomFlags[atomIdx] != 0; const unsigned heavy = warp.ballot(isHeavy); + nonFinite = nonFinite || (isHeavy && !isfinite(leverages[atomIdx])); if (isHeavy) { const int slot = numHeavy + __popc(heavy & ((1u << lane) - 1u)); double rounded; @@ -331,6 +334,12 @@ __device__ __forceinline__ void computeInformationIndices(const double* leverag } numHeavy += __popc(heavy); } + // Non-finite leverages (from non-finite coordinates) have no rank or printed digits to cluster. + if (warp.any(nonFinite)) { + ith = nan(""); + ish = nan(""); + return; + } warp.sync(); for (int slot = lane; slot < numHeavy; slot += kWarpSize) { const double value = heavyValues[slot]; diff --git a/src/descriptors3d_mol.cpp b/src/descriptors3d_mol.cpp index e5233e1f..2a34e885 100644 --- a/src/descriptors3d_mol.cpp +++ b/src/descriptors3d_mol.cpp @@ -355,7 +355,9 @@ Property3DBatchResult calc3DProperties(const std::vector(mol->getNumAtoms())); + const int64_t numAtoms = mol->getNumAtoms(); + inputs.maxMoleculeAtoms = std::max(inputs.maxMoleculeAtoms, static_cast(numAtoms)); + inputs.moleculeAtomPairs += numAtoms * (numAtoms - 1) / 2; } Property3DBatchResult result; diff --git a/src/descriptors3d_topology.cuh b/src/descriptors3d_topology.cuh index 064ec895..a11b4ea2 100644 --- a/src/descriptors3d_topology.cuh +++ b/src/descriptors3d_topology.cuh @@ -110,8 +110,9 @@ struct BondDistanceTableStorage { /** * @brief Builds the bond-distance table for the batch's molecules on @p stream, without host - * synchronization. Capacity is the smaller of kBondTableBudgetBytes and the worst case for - * `inputs.maxMoleculeAtoms`; batches whose largest molecule exceeds kBondTableMaxAtoms (or whose + * synchronization. Capacity is the smaller of kBondTableBudgetBytes and `inputs.moleculeAtomPairs` + * (or, when unset, the worst case for `inputs.maxMoleculeAtoms`); batches whose largest molecule + * exceeds kBondTableMaxAtoms (or whose * maxMoleculeAtoms is unset) get an empty table, so every conformer searches per row. */ inline void buildBondDistanceTable(const DeviceCoordView& coordinates, @@ -124,7 +125,8 @@ inline void buildBondDistanceTable(const DeviceCoordView& coordinates, return; } const int64_t maxPairs = static_cast(maxAtoms) * (maxAtoms - 1) / 2; - const int64_t capacity = std::min(kBondTableBudgetBytes, maxPairs * nMols); + const int64_t pairs = inputs.moleculeAtomPairs > 0 ? inputs.moleculeAtomPairs : maxPairs * nMols; + const int64_t capacity = std::min(kBondTableBudgetBytes, pairs); storage.starts = AsyncDeviceVector(static_cast(nMols) + 1, stream); moleculePairCountsKernel<<<(nMols + kBlockSize - 1) / kBlockSize, kBlockSize, 0, stream>>>(inputs.moleculeAtomStarts, diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index aff15ee3..11cd743a 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -408,6 +408,10 @@ add_executable(test_descriptors3d_whim_arithmetic target_link_libraries(test_descriptors3d_whim_arithmetic PRIVATE descriptors3d cuda_error_check) +add_executable(test_descriptors3d_topology test_descriptors3d_topology.cu) +target_link_libraries(test_descriptors3d_topology PRIVATE descriptors3d + cuda_error_check) + add_executable(test_tfd_cpu test_tfd_cpu.cpp) target_link_libraries(test_tfd_cpu PRIVATE ${RDKit_LIBS} tfd_cpu test_utils nvmolkit_versions) @@ -428,6 +432,7 @@ target_link_libraries(test_tfd PRIVATE ${RDKit_LIBS} tfd_gpu tfd_cpu set(TEST_LIST test_descriptors3d test_descriptors3d_whim_arithmetic + test_descriptors3d_topology test_boolean_tree test_butina test_diversity_pickers diff --git a/tests/test_descriptors3d_topology.cu b/tests/test_descriptors3d_topology.cu new file mode 100644 index 00000000..cc628276 --- /dev/null +++ b/tests/test_descriptors3d_topology.cu @@ -0,0 +1,97 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +#include + +#include +#include +#include + +#include "src/descriptors3d_topology.cuh" +#include "src/utils/device_vector.h" + +using namespace nvMolKit::descriptors3d_detail; + +namespace { + +//! Linear chains with the given atom counts: CSR neighbor lists (molecule-local indices) and atom starts. +struct ChainBatch { + std::vector atomStarts{0}; + std::vector neighborStarts{0}; + std::vector neighbors; + + explicit ChainBatch(const std::vector& sizes) { + for (const int size : sizes) { + for (int atom = 0; atom < size; ++atom) { + if (atom > 0) { + neighbors.push_back(atom - 1); + } + if (atom + 1 < size) { + neighbors.push_back(atom + 1); + } + neighborStarts.push_back(static_cast(neighbors.size())); + } + atomStarts.push_back(atomStarts.back() + size); + } + } +}; + +std::vector buildTable(const std::vector& sizes, const int64_t moleculeAtomPairs) { + const ChainBatch batch(sizes); + nvMolKit::AsyncDeviceVector atomStarts(batch.atomStarts.size()); + nvMolKit::AsyncDeviceVector neighborStarts(batch.neighborStarts.size()); + nvMolKit::AsyncDeviceVector neighbors(batch.neighbors.size()); + atomStarts.copyFromHost(batch.atomStarts); + neighborStarts.copyFromHost(batch.neighborStarts); + neighbors.copyFromHost(batch.neighbors); + + nvMolKit::DeviceCoordView coordinates; + coordinates.nMols = static_cast(sizes.size()); + nvMolKit::Property3DDeviceInputs inputs; + inputs.moleculeAtomStarts = atomStarts.data(); + inputs.bondNeighborStarts = neighborStarts.data(); + inputs.bondNeighbors = neighbors.data(); + for (const int size : sizes) { + inputs.maxMoleculeAtoms = std::max(inputs.maxMoleculeAtoms, size); + } + inputs.moleculeAtomPairs = moleculeAtomPairs; + + BondDistanceTableStorage storage; + buildBondDistanceTable(coordinates, inputs, storage, nullptr); + std::vector table(storage.distances.size()); + storage.distances.copyToHost(table); + EXPECT_EQ(cudaStreamSynchronize(nullptr), cudaSuccess); + return table; +} + +//! Pairs (j < k) of each chain in row-major order: bond count k - j up to kBondTableDepth, else unreached. +std::vector expectedChainDistances(const std::vector& sizes) { + std::vector expected; + for (const int size : sizes) { + for (int j = 0; j < size; ++j) { + for (int k = j + 1; k < size; ++k) { + expected.push_back(k - j <= kBondTableDepth ? static_cast(k - j) : kUnreachedDepth); + } + } + } + return expected; +} + +} // namespace + +TEST(BondDistanceTable, SizedByMoleculeAtomPairs) { + // One large molecule among small ones: the table holds exactly the batch's pairs, not + // nMols * maxPairs. + const std::vector sizes = {3, 40, 1, 12, 2}; + const auto expected = expectedChainDistances(sizes); + EXPECT_EQ(buildTable(sizes, static_cast(expected.size())), expected); +} + +TEST(BondDistanceTable, WithoutPairCountSizesForLargestMolecule) { + const std::vector sizes = {3, 40, 1, 12, 2}; + const auto expected = expectedChainDistances(sizes); + auto table = buildTable(sizes, 0); + ASSERT_EQ(table.size(), sizes.size() * (40 * 39 / 2)); + table.resize(expected.size()); // Entries past the batch's pairs are unused. + EXPECT_EQ(table, expected); +}