From aacfa73ff093f3e09256044c2b3e6f429d0044d8 Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Mon, 27 Jul 2026 12:31:54 -0400 Subject: [PATCH] Fix data race on shared molecule in batched MMFF minimizers Batch threads built each molecule's MMFF force field inside the OpenMP loop, and RDKit's MMFFMolProperties constructor mutates the molecule it types. Conformers of one molecule can span several batches, so threads raced on the same ROMol, producing wrong atom types and heap corruption. The trigger is thread count, not GPU count. Build each molecule's MMFFMolProperties once under std::call_once and share it; contribs still build per batch through the non-mutating overload, so preprocessing stays interleaved with GPU work. --- rdkit_extensions/mmff_flattened_builder.cpp | 32 +++-- rdkit_extensions/mmff_flattened_builder.h | 10 ++ src/minimizer/mmff_minimize.cpp | 125 ++++++++++++++------ tests/test_fmcs_primitives.cu | 6 +- tests/test_mmff.cu | 58 +++++++++ 5 files changed, 176 insertions(+), 55 deletions(-) diff --git a/rdkit_extensions/mmff_flattened_builder.cpp b/rdkit_extensions/mmff_flattened_builder.cpp index 93f59c86..720812a0 100644 --- a/rdkit_extensions/mmff_flattened_builder.cpp +++ b/rdkit_extensions/mmff_flattened_builder.cpp @@ -549,23 +549,29 @@ MMFF::EnergyForceContribsHost constructForcefieldContribs(RDKit::ROMol& mol, return constructForcefieldContribs(mol, &mmffMolProperties, nonBondedThresh, confId, ignoreInterfragInteractions); } +std::shared_ptr makeMMFFMolProperties(RDKit::ROMol& mol, + const nvMolKit::MMFFProperties& props) { + auto mmffMolProperties = std::make_shared(mol, props.variant); + PRECONDITION(mmffMolProperties->isValid(), "missing atom types - invalid force-field"); + mmffMolProperties->setMMFFVariant(props.variant); + mmffMolProperties->setMMFFDielectricConstant(props.dielectricConstant); + mmffMolProperties->setMMFFDielectricModel(props.dielectricModel); + mmffMolProperties->setMMFFBondTerm(props.bondTerm); + mmffMolProperties->setMMFFAngleTerm(props.angleTerm); + mmffMolProperties->setMMFFStretchBendTerm(props.stretchBendTerm); + mmffMolProperties->setMMFFOopTerm(props.oopTerm); + mmffMolProperties->setMMFFTorsionTerm(props.torsionTerm); + mmffMolProperties->setMMFFVdWTerm(props.vdwTerm); + mmffMolProperties->setMMFFEleTerm(props.eleTerm); + return mmffMolProperties; +} + MMFF::EnergyForceContribsHost constructForcefieldContribs(RDKit::ROMol& mol, const nvMolKit::MMFFProperties& props, int confId) { - RDKit::MMFF::MMFFMolProperties mmffMolProperties(mol, props.variant); - PRECONDITION(mmffMolProperties.isValid(), "missing atom types - invalid force-field"); - mmffMolProperties.setMMFFVariant(props.variant); - mmffMolProperties.setMMFFDielectricConstant(props.dielectricConstant); - mmffMolProperties.setMMFFDielectricModel(props.dielectricModel); - mmffMolProperties.setMMFFBondTerm(props.bondTerm); - mmffMolProperties.setMMFFAngleTerm(props.angleTerm); - mmffMolProperties.setMMFFStretchBendTerm(props.stretchBendTerm); - mmffMolProperties.setMMFFOopTerm(props.oopTerm); - mmffMolProperties.setMMFFTorsionTerm(props.torsionTerm); - mmffMolProperties.setMMFFVdWTerm(props.vdwTerm); - mmffMolProperties.setMMFFEleTerm(props.eleTerm); + const auto mmffMolProperties = makeMMFFMolProperties(mol, props); return constructForcefieldContribs( - mol, &mmffMolProperties, props.nonBondedThreshold, confId, props.ignoreInterfragInteractions); + mol, mmffMolProperties.get(), props.nonBondedThreshold, confId, props.ignoreInterfragInteractions); } } // namespace MMFF diff --git a/rdkit_extensions/mmff_flattened_builder.h b/rdkit_extensions/mmff_flattened_builder.h index 8dbdf56c..f474555a 100644 --- a/rdkit_extensions/mmff_flattened_builder.h +++ b/rdkit_extensions/mmff_flattened_builder.h @@ -16,6 +16,7 @@ #ifndef NVMOLKIT_MMFF_FLATTENED_BUILDER_H #define NVMOLKIT_MMFF_FLATTENED_BUILDER_H +#include #include #include "src/forcefields/mmff.h" @@ -61,6 +62,15 @@ EnergyForceContribsHost constructForcefieldContribs(RDKit::ROMol& const nvMolKit::MMFFProperties& props, int confId = -1); +//! Builds the RDKit MMFF atom typing/charges for \c mol. +//! NOTE: RDKit's MMFFMolProperties constructor MUTATES the molecule (Kekulize, setMMFFAromaticity, +//! and a _MMFFSanitized property write). Callers that share a molecule across threads must ensure +//! this runs exactly once per molecule -- see the std::call_once use in the batch minimizers. +//! Returns a shared_ptr so callers need only the forward declaration above: the deleter is bound +//! here, where MMFFMolProperties is complete (its RDKit header is not self-contained). +std::shared_ptr makeMMFFMolProperties(RDKit::ROMol& mol, + const nvMolKit::MMFFProperties& props); + } // namespace nvMolKit::MMFF #endif // NVMOLKIT_MMFF_FLATTENED_BUILDER_H diff --git a/src/minimizer/mmff_minimize.cpp b/src/minimizer/mmff_minimize.cpp index 93e9f753..c53f0512 100644 --- a/src/minimizer/mmff_minimize.cpp +++ b/src/minimizer/mmff_minimize.cpp @@ -19,6 +19,8 @@ #include #include +#include +#include #include #include @@ -38,6 +40,31 @@ struct CachedMoleculeData { EnergyForceContribsHost ffParams; }; +//! Per-molecule MMFF atom typing/charges, shared across batch threads. +//! +//! RDKit's MMFFMolProperties constructor mutates the molecule it types (Kekulize, +//! setMMFFAromaticity, and a _MMFFSanitized property write), and setMMFFAromaticity runs on every +//! construction rather than only the first. Since a molecule's conformers can be spread over +//! several batches, building properties inside the batch loop had concurrent threads mutating the +//! same ROMol -- yielding wrong atom types (and so wrong energies) or a corrupted property dict. +//! +//! Build each molecule's properties exactly once instead. Only the first batch to touch a molecule +//! pays for it, and batches for other molecules never block, so preprocessing stays interleaved +//! with GPU work. The comparatively expensive contribs build stays per-batch and is read-only. +class SharedMolProperties { + public: + explicit SharedMolProperties(size_t numMols) : initFlags_(numMols), properties_(numMols) {} + + RDKit::MMFF::MMFFMolProperties* get(size_t molIdx, RDKit::ROMol& mol, const MMFFProperties& props) { + std::call_once(initFlags_[molIdx], [&]() { properties_[molIdx] = makeMMFFMolProperties(mol, props); }); + return properties_[molIdx].get(); + } + + private: + std::vector initFlags_; + std::vector> properties_; +}; + std::vector> MMFFOptimizeMoleculesConfsBfgs(std::vector& mols, const int maxIters, const MMFFProperties& properties, @@ -135,24 +162,26 @@ MMFFMinimizeResult MMFFMinimizeMoleculesConfs(std::vector& collector.converged.setStream(collector.stream); } } + SharedMolProperties sharedMolProperties(mols.size()); detail::OpenMPExceptionRegistry exceptionHandler; -#pragma omp parallel for num_threads(ctx.numThreads) schedule(dynamic) default(none) shared(allConformers, \ - moleculeEnergies, \ - moleculeConverged, \ - totalConformers, \ - effectiveBatchSize, \ - maxIters, \ - gradTol, \ - properties, \ - constraints, \ - ctx, \ - threadBuffers, \ - deviceCollectors, \ - deviceOutput, \ - useDeviceInput, \ - deviceInput, \ - deviceInputIndex, \ - backend, \ +#pragma omp parallel for num_threads(ctx.numThreads) schedule(dynamic) default(none) shared(allConformers, \ + moleculeEnergies, \ + moleculeConverged, \ + totalConformers, \ + effectiveBatchSize, \ + maxIters, \ + gradTol, \ + properties, \ + constraints, \ + ctx, \ + threadBuffers, \ + deviceCollectors, \ + deviceOutput, \ + useDeviceInput, \ + deviceInput, \ + deviceInputIndex, \ + backend, \ + sharedMolProperties, \ exceptionHandler) for (size_t batchStart = 0; batchStart < totalConformers; batchStart += effectiveBatchSize) { try { @@ -194,10 +223,18 @@ MMFFMinimizeResult MMFFMinimizeMoleculesConfs(std::vector& auto it = moleculeCache.find(mol); if (it == moleculeCache.end()) { - ScopedNvtxRange computeCacheRange("Preprocess single molecule"); - CachedMoleculeData cached; - cached.ffParams = constructForcefieldContribs(*mol, properties[confInfo.molIdx]); - it = moleculeCache.insert({mol, std::move(cached)}).first; + ScopedNvtxRange computeCacheRange("Preprocess single molecule"); + CachedMoleculeData cached; + const MMFFProperties& molProps = properties[confInfo.molIdx]; + // Atom typing is shared and built once per molecule; the contribs build below is + // read-only and stays per-batch. + auto* mmffMolProperties = sharedMolProperties.get(confInfo.molIdx, *mol, molProps); + cached.ffParams = constructForcefieldContribs(*mol, + mmffMolProperties, + molProps.nonBondedThreshold, + /*confId=*/-1, + molProps.ignoreInterfragInteractions); + it = moleculeCache.insert({mol, std::move(cached)}).first; } ScopedNvtxRange addToBatchRange("Add conformer to batch data"); @@ -394,21 +431,23 @@ MMFFMinimizeResult MMFFMinimizeMoleculesConfsFire( collector.converged.setStream(collector.stream); } } + SharedMolProperties sharedMolProperties(mols.size()); detail::OpenMPExceptionRegistry exceptionHandler; -#pragma omp parallel for num_threads(ctx.numThreads) schedule(dynamic) default(none) shared(allConformers, \ - moleculeEnergies, \ - moleculeConverged, \ - totalConformers, \ - effectiveBatchSize, \ - maxIters, \ - fireOptions, \ - properties, \ - constraints, \ - ctx, \ - threadBuffers, \ - deviceCollectors, \ - deviceOutput, \ - backend, \ +#pragma omp parallel for num_threads(ctx.numThreads) schedule(dynamic) default(none) shared(allConformers, \ + moleculeEnergies, \ + moleculeConverged, \ + totalConformers, \ + effectiveBatchSize, \ + maxIters, \ + fireOptions, \ + properties, \ + constraints, \ + ctx, \ + threadBuffers, \ + deviceCollectors, \ + deviceOutput, \ + backend, \ + sharedMolProperties, \ exceptionHandler) for (size_t batchStart = 0; batchStart < totalConformers; batchStart += effectiveBatchSize) { try { @@ -435,10 +474,18 @@ MMFFMinimizeResult MMFFMinimizeMoleculesConfsFire( auto it = moleculeCache.find(mol); if (it == moleculeCache.end()) { - ScopedNvtxRange computeCacheRange("Preprocess single molecule"); - CachedMoleculeData cached; - cached.ffParams = constructForcefieldContribs(*mol, properties[confInfo.molIdx]); - it = moleculeCache.insert({mol, std::move(cached)}).first; + ScopedNvtxRange computeCacheRange("Preprocess single molecule"); + CachedMoleculeData cached; + const MMFFProperties& molProps = properties[confInfo.molIdx]; + // Atom typing is shared and built once per molecule; the contribs build below is + // read-only and stays per-batch. + auto* mmffMolProperties = sharedMolProperties.get(confInfo.molIdx, *mol, molProps); + cached.ffParams = constructForcefieldContribs(*mol, + mmffMolProperties, + molProps.nonBondedThreshold, + /*confId=*/-1, + molProps.ignoreInterfragInteractions); + it = moleculeCache.insert({mol, std::move(cached)}).first; } ScopedNvtxRange addToBatchRange("Add conformer to batch data"); diff --git a/tests/test_fmcs_primitives.cu b/tests/test_fmcs_primitives.cu index 324371c8..bb4338f3 100644 --- a/tests/test_fmcs_primitives.cu +++ b/tests/test_fmcs_primitives.cu @@ -13,9 +13,9 @@ #include #include -#include "fmcs_cuda/fmcs_seed.cuh" -#include "fmcs_cuda/fmcs_seed_queue.cuh" -#include "mcs_common/mcs_cooperative_copy.cuh" +#include "src/mcs/fmcs_cuda/fmcs_seed.cuh" +#include "src/mcs/fmcs_cuda/fmcs_seed_queue.cuh" +#include "src/mcs/mcs_common/mcs_cooperative_copy.cuh" #include "src/utils/device_vector.h" namespace { diff --git a/tests/test_mmff.cu b/tests/test_mmff.cu index 92f58f7c..fe86233e 100644 --- a/tests/test_mmff.cu +++ b/tests/test_mmff.cu @@ -1959,6 +1959,64 @@ TEST(MMFFMultiGPU, NonZeroGPUID) { } } +// Regression test: several batch threads minimizing conformers of the same molecule used to build +// that molecule's MMFF properties concurrently, and RDKit's MMFFMolProperties constructor mutates +// the molecule it types. That produced wrong atom types (energies far from the RDKit reference) or +// a corrupted property dict (heap corruption). Only one GPU is needed -- the trigger is the thread +// count, not the device count, so batchesPerGpu is what matters here. +TEST(MMFFThreadSafety, SharedMoleculeAcrossBatchThreads) { + nvMolKit::BatchHardwareOptions options; + options.batchSize = 2; + options.batchesPerGpu = 4; // 4 threads on one GPU + options.gpuIds.push_back(0); + + constexpr int numMols = 3; + constexpr int numConfsPerMol = 6; + + std::vector> mols; + getMols(getTestDataFolderPath() + "/MMFF94_dative.sdf", mols, numMols); + + std::vector molPtrs; + for (auto& mol : mols) { + for (int i = 1; i < numConfsPerMol; i++) { + auto conf = new RDKit::Conformer(mol->getConformer()); + perturbConformer(*conf, 0.5f, i + 101); + mol->addConformer(conf, true); + } + molPtrs.push_back(mol.get()); + } + + std::vector> molCopies; + for (const auto& mol : mols) { + molCopies.push_back(std::make_unique(*mol)); + std::vector> res(molCopies.back()->getNumConformers(), {-1, -1}); + RDKit::MMFF::MMFFOptimizeMoleculeConfs(*molCopies.back(), res); + } + + std::vector> gotEnergies = + nvMolKit::MMFF::MMFFOptimizeMoleculesConfsBfgs(molPtrs, 200, nvMolKit::MMFFProperties{}, options); + + ASSERT_EQ(gotEnergies.size(), mols.size()); + for (size_t molIdx = 0; molIdx < mols.size(); ++molIdx) { + auto& molRef = *molCopies[molIdx]; + auto molProps = std::make_unique(molRef); + std::unique_ptr refFF(RDKit::MMFF::constructForceField(molRef, molProps.get())); + + const auto& energiesForMol = gotEnergies[molIdx]; + ASSERT_EQ(energiesForMol.size(), molRef.getNumConformers()); + + int confIdx = 0; + for (auto confIter = molRef.beginConformers(); confIter != molRef.endConformers(); ++confIter) { + std::vector posRef; + nvMolKit::confPosToVect(**confIter, posRef); + const double refEnergy = refFF->calcEnergy(posRef.data()); + ASSERT_NEAR(energiesForMol[confIdx], refEnergy, 1e-4) + << "Energy mismatch vs RDKit reference for molecule " << molIdx << ", conformer " << confIdx; + confIdx++; + } + } +} + TEST(MMFFMultiGPU, MultiGPUSpecificIds) { // Requires multiple GPUs const int numDevices = nvMolKit::countCudaDevices();