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();