From 1eb7055f1cc26eed186728dcb5c6f6c61ebc518e Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Fri, 26 Jun 2026 21:26:29 -0400 Subject: [PATCH 1/2] Add FIRE per-molecule backend state --- src/minimizer/bfgs_types.h | 15 +++++++ src/minimizer/fire_minimizer.cu | 76 ++++++++++++++++++++++++++++++--- src/minimizer/fire_minimizer.h | 24 +++++++++-- tests/test_fire_minimizer.cu | 28 ++++++++++++ 4 files changed, 133 insertions(+), 10 deletions(-) diff --git a/src/minimizer/bfgs_types.h b/src/minimizer/bfgs_types.h index 25e9d210..ccd6b895 100644 --- a/src/minimizer/bfgs_types.h +++ b/src/minimizer/bfgs_types.h @@ -42,6 +42,21 @@ enum class BfgsBackend { //! Atom count threshold for HYBRID backend selection (use PER_MOLECULE if max atoms <= this value) constexpr int kHybridBackendAtomThreshold = 64; +//! \brief Backend implementation type for FIRE minimization. +//! +//! Mirrors ::BfgsBackend: ::BATCHED runs streaming kernels over the batch with +//! a host-driven loop; ::PER_MOLECULE runs the full FIRE iteration loop in one +//! CUDA kernel per molecule; ::HYBRID auto-selects based on the largest molecule +//! in the batch. +enum class FireBackend { + BATCHED = 0, + PER_MOLECULE = 1, + HYBRID = 2, +}; + +//! Atom count threshold for HYBRID FIRE backend selection (use PER_MOLECULE if max atoms <= this value). +constexpr int kHybridFireBackendAtomThreshold = 64; + } // namespace nvMolKit #endif // NVMOLKIT_BFGS_TYPES_H diff --git a/src/minimizer/fire_minimizer.cu b/src/minimizer/fire_minimizer.cu index 0c60682c..782071b1 100644 --- a/src/minimizer/fire_minimizer.cu +++ b/src/minimizer/fire_minimizer.cu @@ -385,11 +385,13 @@ __global__ void fireStuckCheckKernel(cuda::std::span activeSystemI FireBatchMinimizer::FireBatchMinimizer(const int dataDim, const FireOptions& options, cudaStream_t stream, - const bool debugMode) + const bool debugMode, + const FireBackend backend) : dataDim_(dataDim), fireOptions_(options), stream_(stream), - debugMode_(debugMode) { + debugMode_(debugMode), + backend_(backend) { velocities_.setStream(stream_); statuses_.setStream(stream_); dt_.setStream(stream_); @@ -405,10 +407,23 @@ FireBatchMinimizer::FireBatchMinimizer(const int dataDim, energyMaxStreak_.setStream(stream_); stuckStreak_.setStream(stream_); convergeReason_.setStream(stream_); + activeMolIdsDevice_.setStream(stream_); loopStatusHost_.resize(1); loopStatusHost_[0] = 0; } +FireBackend FireBatchMinimizer::resolveBackend(const std::vector& atomStartsHost) const { + if (backend_ != FireBackend::HYBRID) { + return backend_; + } + for (size_t i = 0; i + 1 < atomStartsHost.size(); ++i) { + if (atomStartsHost[i + 1] - atomStartsHost[i] > kHybridFireBackendAtomThreshold) { + return FireBackend::BATCHED; + } + } + return FireBackend::PER_MOLECULE; +} + void FireBatchMinimizer::setMasses(const std::vector& masses) { hostMasses_ = masses; } @@ -430,14 +445,17 @@ void FireBatchMinimizer::setConvergencePollInterval(const int interval) { void FireBatchMinimizer::initialize(const std::vector& atomStartsHost, const double* masses, - const uint8_t* activeThisStage) { + const uint8_t* activeThisStage, + const FireBackend effectiveBackend) { step_ = 0; const int totalAtoms = atomStartsHost.back(); const int numSystems = static_cast(atomStartsHost.size()) - 1; - const bool isContinuation = hasInitializedBatch_ && cachedNumSystems_ == numSystems && - cachedTotalAtoms_ == totalAtoms && cachedActiveThisStage_ == activeThisStage && - cachedMasses_ == masses; + // Continuation cache only applies to the BATCHED backend; the per-mol path + // resets per-system state on every call. + const bool isContinuation = effectiveBackend == FireBackend::BATCHED && hasInitializedBatch_ && + cachedNumSystems_ == numSystems && cachedTotalAtoms_ == totalAtoms && + cachedActiveThisStage_ == activeThisStage && cachedMasses_ == masses; numSystems_ = numSystems; @@ -473,6 +491,50 @@ void FireBatchMinimizer::initialize(const std::vector& atomStartsHost, dt_.resize(numSystems); setAll(dt_, fireOptions_.dtInit); + if (effectiveBackend == FireBackend::PER_MOLECULE) { + energyMinStreak_.resize(0); + energyMaxStreak_.resize(0); + stuckStreak_.resize(0); + convergeReason_.resize(0); + debugPowers_.resize(0); + debugOutputs_.clear(); + activeSystemIndices_.resize(0); + allSystemIndices_.resize(0); + pollsSinceLastEnergyEval_ = 0; + + activeHost_.resize(numSystems); + convergenceHost_.resize(numSystems); + std::fill_n(activeHost_.begin(), numSystems, 1); + if (activeThisStage != nullptr) { + cudaCheckError(cudaMemcpyAsync(activeHost_.data(), + activeThisStage, + numSystems * sizeof(uint8_t), + cudaMemcpyDeviceToHost, + stream_)); + cudaCheckError(cudaStreamSynchronize(stream_)); + } + activeMolIds_.clear(); + maxAtomsInBatch_ = 0; + for (int sysIdx = 0; sysIdx < numSystems; ++sysIdx) { + if (activeHost_[sysIdx] == 0) { + continue; + } + activeMolIds_.push_back(sysIdx); + const int numAtoms = atomStartsHost[sysIdx + 1] - atomStartsHost[sysIdx]; + if (numAtoms > maxAtomsInBatch_) { + maxAtomsInBatch_ = numAtoms; + } + } + if (!activeMolIds_.empty()) { + activeMolIdsDevice_.resize(activeMolIds_.size()); + activeMolIdsDevice_.setFromVector(activeMolIds_); + } + + resetContinuationCache(); + lastKnownNumUnfinished_ = static_cast(activeMolIds_.size()); + return; + } + activeSystemIndices_.resize(numSystems); allSystemIndices_.resize(numSystems); std::vector indicesHost(numSystems); @@ -696,7 +758,7 @@ bool FireBatchMinimizer::minimize(const int nu const GradFunctor gFunc, const uint8_t* activeThisStage) { const ScopedNvtxRange minimizeRange("FireBatchMinimizer::minimize (batched)"); - initialize(atomStartsHost, nullptr, activeThisStage); + initialize(atomStartsHost, nullptr, activeThisStage, FireBackend::BATCHED); for (int iter = 0; iter < numIters; ++iter) { if (debugMode_) { diff --git a/src/minimizer/fire_minimizer.h b/src/minimizer/fire_minimizer.h index a8c8df15..2d149f4a 100644 --- a/src/minimizer/fire_minimizer.h +++ b/src/minimizer/fire_minimizer.h @@ -18,6 +18,7 @@ #include +#include "src/minimizer/bfgs_types.h" #include "src/minimizer/fire_options.h" #include "src/minimizer/minimizer_api.h" #include "src/utils/device_vector.h" @@ -58,16 +59,23 @@ class FireBatchMinimizer final : public BatchMinimizer { explicit FireBatchMinimizer(int dataDim = 3, const FireOptions& options = FireOptions(), cudaStream_t stream = nullptr, - bool debugMode = false); + bool debugMode = false, + FireBackend backend = FireBackend::BATCHED); ~FireBatchMinimizer() override = default; + //! \brief Resolve the effective backend for the provided batch under HYBRID selection. + FireBackend resolveBackend(const std::vector& atomStartsHost) const; + //! \brief Initialize internal buffers for a new batch. //! \param atomStartsHost Host offsets for the first atom of each system. //! \param masses Optional pointer to per-atom masses; nullptr means use any masses set via setMasses(). //! \param activeThisStage Optional uint8_t mask (1 = active). When nullptr all systems start active. + //! \param effectiveBackend Selects which backend's auxiliary buffers to materialize. + //! Pass the value returned by ::resolveBackend so HYBRID is collapsed first. void initialize(const std::vector& atomStartsHost, - const double* masses = nullptr, - const uint8_t* activeThisStage = nullptr); + const double* masses = nullptr, + const uint8_t* activeThisStage = nullptr, + FireBackend effectiveBackend = FireBackend::BATCHED); //! \brief Provide per-atom masses to be used on the next initialization //! when explicit masses are not supplied. Passing an empty vector clears @@ -107,6 +115,8 @@ class FireBatchMinimizer final : public BatchMinimizer { //! \brief Cadence (in iterations) at which the minimize() loop reads the //! still-running system count back to the host. Default 8. + //! \note Only the BATCHED backend uses this; per-molecule kernels iterate + //! entirely device-side and ignore the poll interval. void setConvergencePollInterval(int interval); //! \brief Read back internal per-system state for testing. @@ -144,6 +154,7 @@ class FireBatchMinimizer final : public BatchMinimizer { int numSystems_ = 0; int convergencePollInterval_ = 8; int lastKnownNumUnfinished_ = 0; + FireBackend backend_ = FireBackend::BATCHED; AsyncDeviceVector velocities_; AsyncDeviceVector masses_; @@ -184,6 +195,13 @@ class FireBatchMinimizer final : public BatchMinimizer { //! Per-system convergence reason for diagnostics: 0=active, 1=grad-tol, 2=stuck-plateau. AsyncDeviceVector convergeReason_; + + // Per-molecule kernel data (used when backend_ == PER_MOLECULE / HYBRID resolves to it). + int maxAtomsInBatch_ = 0; //!< Largest molecule in batch (for kernel dispatch). + std::vector activeMolIds_; //!< Active molecule IDs (host). + AsyncDeviceVector activeMolIdsDevice_; //!< Device copy of @c activeMolIds_. + PinnedHostVector activeHost_; //!< Pinned scratch for caller-supplied active mask. + PinnedHostVector convergenceHost_; //!< Pinned scratch for status readback. }; } // namespace nvMolKit diff --git a/tests/test_fire_minimizer.cu b/tests/test_fire_minimizer.cu index 6bc31c1f..803bab91 100644 --- a/tests/test_fire_minimizer.cu +++ b/tests/test_fire_minimizer.cu @@ -965,3 +965,31 @@ TEST(FireMinimizer, StaggeredConvergenceCount) { uniqueConvIters.erase(std::unique(uniqueConvIters.begin(), uniqueConvIters.end()), uniqueConvIters.end()); EXPECT_GE(uniqueConvIters.size(), 2u) << "Test should produce at least two distinct convergence iterations"; } + +TEST(FireMinimizer, HybridBackendSelectionAndPerMolInitialization) { + nvMolKit::FireOptions options; + nvMolKit::FireBatchMinimizer minimizer(kDim, + options, + /*stream=*/nullptr, + /*debugMode=*/false, + nvMolKit::FireBackend::HYBRID); + + EXPECT_EQ(minimizer.resolveBackend({0, 5, 10, 30}), nvMolKit::FireBackend::PER_MOLECULE); + EXPECT_EQ(minimizer.resolveBackend({0, 5, 10, 200}), nvMolKit::FireBackend::BATCHED); + + const std::vector atomStarts = {0, 4, 10, 18}; + minimizer.initialize(atomStarts, nullptr, nullptr, nvMolKit::FireBackend::PER_MOLECULE); + + EXPECT_EQ(minimizer.numActiveSystemsHost(), 3); + const auto state = minimizer.snapshotInternalState(); + ASSERT_EQ(state.statuses.size(), 3u); + ASSERT_EQ(state.dt.size(), 3u); + ASSERT_EQ(state.alpha.size(), 3u); + ASSERT_EQ(state.nStepsPositive.size(), 3u); + for (size_t i = 0; i < state.statuses.size(); ++i) { + EXPECT_EQ(state.statuses[i], 1); + EXPECT_EQ(state.nStepsPositive[i], 0); + EXPECT_NEAR(state.dt[i], options.dtInit, 0.0); + EXPECT_NEAR(state.alpha[i], options.alphaInit, 0.0); + } +} From b24eb98d274a7f8cb20f600370319fb0fccfd23b Mon Sep 17 00:00:00 2001 From: Kevin Boyd Date: Mon, 29 Jun 2026 10:00:27 -0400 Subject: [PATCH 2/2] Add MMFF per-molecule FIRE kernel --- src/minimizer/CMakeLists.txt | 6 +- src/minimizer/fire_minimize_permol_kernels.cu | 341 ++++++++++++++++++ src/minimizer/fire_minimize_permol_kernels.h | 67 ++++ tests/CMakeLists.txt | 17 +- tests/test_fire_minimizer_permol.cu | 212 +++++++++++ 5 files changed, 640 insertions(+), 3 deletions(-) create mode 100644 src/minimizer/fire_minimize_permol_kernels.cu create mode 100644 src/minimizer/fire_minimize_permol_kernels.h create mode 100644 tests/test_fire_minimizer_permol.cu diff --git a/src/minimizer/CMakeLists.txt b/src/minimizer/CMakeLists.txt index 889c6876..40c2ea73 100644 --- a/src/minimizer/CMakeLists.txt +++ b/src/minimizer/CMakeLists.txt @@ -31,11 +31,13 @@ target_link_libraries( PUBLIC host_vector device conformer_types PRIVATE ${RDKit_LIBS} OpenMP::OpenMP_CXX) -add_library(fire_minimizer fire_minimizer.cu fire_minimizer.h) +add_library( + fire_minimizer fire_minimizer.cu fire_minimizer.h + fire_minimize_permol_kernels.cu fire_minimize_permol_kernels.h) target_link_libraries( fire_minimizer PUBLIC device_vector batched_forcefield - PRIVATE cub_helpers cccl_interface) + PRIVATE mmff cub_helpers cccl_interface) add_library(bfgs_mmff bfgs_mmff.cpp) target_link_libraries( diff --git a/src/minimizer/fire_minimize_permol_kernels.cu b/src/minimizer/fire_minimize_permol_kernels.cu new file mode 100644 index 00000000..df1569f2 --- /dev/null +++ b/src/minimizer/fire_minimize_permol_kernels.cu @@ -0,0 +1,341 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include "src/forcefields/mmff_kernels.h" +#include "src/forcefields/mmff_kernels_device.cuh" +#include "src/minimizer/fire_minimize_permol_kernels.h" +#include "src/utils/device_vector.h" + +namespace nvMolKit { + +namespace { + +constexpr int kFirePerMolBlockSize = 128; +constexpr int kDataDim = 3; + +//! Acceleration conversion factor: 1 kcal/mol/Å applied to 1 amu produces 4.184 * 100 Å/ps^2. +//! Must remain identical to the batched FIRE conversion factor. +constexpr double kForceKcalMolPerAng_PerAmu_to_AngPerPs2 = 4.184 * 100.0; + +struct FirePerMolKernelParams { + double dtIncrementFactor; + double dtDecrementFactor; + double minDt; + double maxDt; + double dMax; + double alphaStart; + double alphaDecrementFactor; + double gradTol; + int nMinForIncrease; +}; + +FirePerMolKernelParams buildKernelParams(const FireOptions& opts, const double gradTol) { + FirePerMolKernelParams params{}; + params.dtIncrementFactor = opts.timeStepIncrement; + params.dtDecrementFactor = opts.timeStepDecrement; + params.minDt = opts.dtInit * opts.dtMinFactor; + params.maxDt = opts.dtInit * opts.dtMaxFactor; + params.dMax = opts.dMax; + params.alphaStart = opts.alphaInit; + params.alphaDecrementFactor = opts.alphaDecrement; + params.gradTol = gradTol; + params.nMinForIncrease = opts.nMinForIncrease; + return params; +} + +__launch_bounds__(kFirePerMolBlockSize) + __global__ void firePerMolMmffKernel(const int numIters, + const FirePerMolKernelParams params, + const bool takeHalfStepBack, + const bool useAbc, + const bool useMass, + const MMFF::EnergyForceContribsDevicePtr* terms, + const MMFF::BatchedIndicesDevicePtr* systemIndices, + const int* molIdList, + const int* atomStarts, + double* positions, + double* grad, + double* velocities, + double* alphas, + double* dts, + int* nStepsPositive, + const double* masses, + double* energyOuts, + uint8_t* statuses) { + const int molIdx = molIdList[blockIdx.x]; + const int tid = threadIdx.x; + + if (statuses[molIdx] == 0) { + return; + } + + const int atomStart = atomStarts[molIdx]; + const int atomEnd = atomStarts[molIdx + 1]; + const int numTerms = (atomEnd - atomStart) * kDataDim; + + double* const molCoords = positions + atomStart * kDataDim; + double* const molGrad = grad + atomStart * kDataDim; + double* const molVel = velocities + atomStart * kDataDim; + const double* const massSys = useMass ? (masses + atomStart) : nullptr; + + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage tempStorage; + + __shared__ double sharedDt; + __shared__ double sharedAlpha; + __shared__ int sharedNsteps; + __shared__ double sharedScalar0; + __shared__ double sharedScalar1; + __shared__ bool sharedConverged; + + if (tid == 0) { + sharedDt = dts[molIdx]; + sharedAlpha = alphas[molIdx]; + sharedNsteps = nStepsPositive[molIdx]; + sharedConverged = false; + } + __syncthreads(); + + for (int iter = 0; iter < numIters; ++iter) { + const bool isFirstStep = (iter == 0); + + for (int i = tid; i < numTerms; i += kFirePerMolBlockSize) { + molGrad[i] = 0.0; + } + __syncthreads(); + MMFF::molGrad(*terms, *systemIndices, molCoords, molGrad, molIdx, tid); + __syncthreads(); + + double power = 0.0; + double gradSq = 0.0; + for (int i = tid; i < numTerms; i += kFirePerMolBlockSize) { + const double fi = molGrad[i]; + if (!isFirstStep) { + power += molVel[i] * -fi; + } + gradSq += fi * fi; + } + double powerSum = 0.0; + if (!isFirstStep) { + powerSum = BlockReduce(tempStorage).Sum(power); + __syncthreads(); + } + const double gradSqSum = BlockReduce(tempStorage).Sum(gradSq); + if (tid == 0) { + sharedScalar0 = powerSum; + sharedScalar1 = gradSqSum; + } + __syncthreads(); + const double powerShared = sharedScalar0; + const double gradSqShared = sharedScalar1; + + if (tid == 0 && sqrt(gradSqShared) <= params.gradTol) { + sharedConverged = true; + statuses[molIdx] = 0; + } + __syncthreads(); + if (sharedConverged) { + break; + } + + if (tid == 0 && !isFirstStep) { + double newDt = sharedDt; + double newAlpha = sharedAlpha; + int newNsteps = sharedNsteps; + if (powerShared >= 0.0) { + newNsteps = sharedNsteps + 1; + if (newNsteps > params.nMinForIncrease) { + newDt = fmin(sharedDt * params.dtIncrementFactor, params.maxDt); + newAlpha = sharedAlpha * params.alphaDecrementFactor; + } + } else { + newNsteps = 0; + newAlpha = params.alphaStart; + newDt = fmax(sharedDt * params.dtDecrementFactor, params.minDt); + } + sharedDt = newDt; + sharedAlpha = newAlpha; + sharedNsteps = newNsteps; + } + __syncthreads(); + + const bool negative = !isFirstStep && (powerShared < 0.0); + if (negative) { + const double dtNow = sharedDt; + for (int i = tid; i < numTerms; i += kFirePerMolBlockSize) { + if (takeHalfStepBack) { + molCoords[i] -= 0.5 * dtNow * molVel[i]; + } + molVel[i] = 0.0; + molGrad[i] = 0.0; + } + __syncthreads(); + MMFF::molGrad(*terms, *systemIndices, molCoords, molGrad, molIdx, tid); + __syncthreads(); + } + + const double dt = sharedDt; + const double alpha = sharedAlpha; + const int nsteps = sharedNsteps; + + double vSqAccum = 0.0; + double gradSqAccum = 0.0; + for (int i = tid; i < numTerms; i += kFirePerMolBlockSize) { + double accel; + if (useMass) { + const double accelMag = -molGrad[i] * kForceKcalMolPerAng_PerAmu_to_AngPerPs2; + const int coordIdx = i / kDataDim; + accel = accelMag / massSys[coordIdx]; + } else { + accel = -molGrad[i]; + } + const double newV = molVel[i] + dt * accel; + molVel[i] = newV; + vSqAccum += newV * newV; + gradSqAccum += molGrad[i] * molGrad[i]; + } + const double vSqReduced = BlockReduce(tempStorage).Sum(vSqAccum); + if (tid == 0) { + sharedScalar0 = vSqReduced; + } + __syncthreads(); + const double vSqSum = sharedScalar0; + const double gradSqReduced = BlockReduce(tempStorage).Sum(gradSqAccum); + if (tid == 0) { + sharedScalar0 = gradSqReduced; + } + __syncthreads(); + const double gradSqSum2 = sharedScalar0; + + const double mixCoef1 = 1.0 - alpha; + const double mixCoef2 = (gradSqSum2 > 1e-30) ? (alpha * sqrt(vSqSum) / sqrt(gradSqSum2)) : 0.0; + double abcMult = 1.0; + if (useAbc) { + const double oneMinusA = 1.0 - fmax(alpha, 1e-10); + const double powTerm = pow(oneMinusA, static_cast(nsteps + 1)); + const double denom = 1.0 - powTerm; + abcMult = (denom > 1e-30) ? (1.0 / denom) : 1.0; + } + + for (int i = tid; i < numTerms; i += kFirePerMolBlockSize) { + const double vMix = mixCoef1 * molVel[i] + mixCoef2 * (-molGrad[i]); + molVel[i] = abcMult * vMix; + } + __syncthreads(); + + double drScale = 1.0; + if (useAbc) { + if (params.dMax > 0.0) { + const double maxV = params.dMax / dt; + for (int i = tid; i < numTerms; i += kFirePerMolBlockSize) { + molVel[i] = fmax(-maxV, fmin(maxV, molVel[i])); + } + __syncthreads(); + } + } else if (params.dMax > 0.0) { + double drSqAccum = 0.0; + for (int i = tid; i < numTerms; i += kFirePerMolBlockSize) { + const double dri = dt * molVel[i]; + drSqAccum += dri * dri; + } + const double drSqReduced = BlockReduce(tempStorage).Sum(drSqAccum); + if (tid == 0) { + sharedScalar0 = drSqReduced; + } + __syncthreads(); + const double drNorm = sqrt(sharedScalar0); + if (drNorm > params.dMax) { + drScale = params.dMax / drNorm; + } + } + + for (int i = tid; i < numTerms; i += kFirePerMolBlockSize) { + molCoords[i] += drScale * dt * molVel[i]; + } + __syncthreads(); + } + + if (tid == 0) { + dts[molIdx] = sharedDt; + alphas[molIdx] = sharedAlpha; + nStepsPositive[molIdx] = sharedNsteps; + } + __syncthreads(); + + const double finalThreadEnergy = + MMFF::molEnergy(*terms, *systemIndices, molCoords, molIdx, tid); + const double finalEnergy = BlockReduce(tempStorage).Sum(finalThreadEnergy); + if (tid == 0) { + energyOuts[molIdx] = finalEnergy; + } +} + +} // namespace + +cudaError_t launchFirePerMolKernel(const int numMols, + const int* molIds, + [[maybe_unused]] const int maxAtoms, + const int* atomStarts, + const FireOptions& fireOptions, + const int numIters, + const double gradTol, + const MMFF::EnergyForceContribsDevicePtr& terms, + const MMFF::BatchedIndicesDevicePtr& systemIndices, + const bool hasConstraints, + double* positions, + double* grad, + double* velocities, + double* alphas, + double* dts, + int* nStepsPositive, + const double* masses, + double* energyOuts, + uint8_t* statuses, + const cudaStream_t stream) { + if (numMols == 0) { + return cudaSuccess; + } + if (hasConstraints) { + return cudaErrorNotSupported; + } + + const AsyncDevicePtr devTerms(terms, stream); + const AsyncDevicePtr devSysIdx(systemIndices, stream); + const FirePerMolKernelParams params = buildKernelParams(fireOptions, gradTol); + const bool useMass = fireOptions.useMass && masses != nullptr; + firePerMolMmffKernel<<>>(numIters, + params, + fireOptions.takeHalfStepBack, + fireOptions.abcCorrection, + useMass, + devTerms.data(), + devSysIdx.data(), + molIds, + atomStarts, + positions, + grad, + velocities, + alphas, + dts, + nStepsPositive, + masses, + energyOuts, + statuses); + return cudaGetLastError(); +} + +} // namespace nvMolKit diff --git a/src/minimizer/fire_minimize_permol_kernels.h b/src/minimizer/fire_minimize_permol_kernels.h new file mode 100644 index 00000000..7b1da1c6 --- /dev/null +++ b/src/minimizer/fire_minimize_permol_kernels.h @@ -0,0 +1,67 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef NVMOLKIT_FIRE_MINIMIZE_PERMOL_KERNELS_H +#define NVMOLKIT_FIRE_MINIMIZE_PERMOL_KERNELS_H + +#include + +#include + +#include "src/forcefields/mmff_kernels.h" +#include "src/minimizer/fire_minimizer.h" + +namespace nvMolKit { + +//! \brief Per-block launch configuration for the per-molecule FIRE kernels. +//! +//! The caller supplies per-system state buffers and force-field term descriptors. The +//! kernel iterates the full FIRE 2.0 loop internally and writes a per-system +//! status (0 = converged, 1 = active) into @p statuses. +struct FirePerMolLaunchParams { + int numIters = 0; //!< Maximum FIRE iterations to run inside the kernel. + double gradTol = 0.0; //!< sqrt(sum(grad^2)) per-system convergence tolerance. + bool takeHalfStepBack = true; //!< When true and power<0, take a half step back and zero v. + bool useAbc = false; //!< Apply ABC-FIRE mixer correction. + bool useMass = false; //!< Mass-weight the force kick (requires non-null masses). +}; + +//! \brief Launch per-molecule FIRE 2.0 minimization - MMFF specialization. +//! \note ::FireOptions stuck-detection fields are not supported on the per-mol path. +//! \note Returns cudaErrorNotSupported when @p hasConstraints is true. +cudaError_t launchFirePerMolKernel(int numMols, + const int* molIds, + int maxAtoms, + const int* atomStarts, + const FireOptions& fireOptions, + int numIters, + double gradTol, + const MMFF::EnergyForceContribsDevicePtr& terms, + const MMFF::BatchedIndicesDevicePtr& systemIndices, + bool hasConstraints, + double* positions, + double* grad, + double* velocities, + double* alphas, + double* dts, + int* nStepsPositive, + const double* masses, + double* energyOuts, + uint8_t* statuses, + cudaStream_t stream = nullptr); + +} // namespace nvMolKit + +#endif // NVMOLKIT_FIRE_MINIMIZE_PERMOL_KERNELS_H diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index cfd03058..a73b750a 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -236,6 +236,20 @@ add_executable(test_fire_minimizer test_fire_minimizer.cu) target_link_libraries(test_fire_minimizer PRIVATE fire_minimizer device_vector cuda_error_check) +add_executable(test_fire_minimizer_permol test_fire_minimizer_permol.cu) +target_link_libraries( + test_fire_minimizer_permol + PRIVATE fire_minimizer + batched_forcefield + mmff_batched_forcefield + rdkit_mmff_flattened + device_vector + cuda_error_check + ${RDKit_LIBS} + test_utils + mmff + device) + add_executable(test_etkdg_etk_minimize test_etkdg_etk_minimize.cu) target_link_libraries( test_etkdg_etk_minimize @@ -387,7 +401,8 @@ set(TEST_LIST test_similarity test_thread_safe_queue test_work_splitting - test_fire_minimizer) + test_fire_minimizer + test_fire_minimizer_permol) # Auto populate the NVMOLKIT_TESTDATA environment variable for ctest. foreach(arg ${TEST_LIST}) diff --git a/tests/test_fire_minimizer_permol.cu b/tests/test_fire_minimizer_permol.cu new file mode 100644 index 00000000..9c51b64e --- /dev/null +++ b/tests/test_fire_minimizer_permol.cu @@ -0,0 +1,212 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include +// clang-format off +// Bug in RDKit, includes need to be ordered. +#include +#include +// clang-format on +#include + +#include +#include +#include +#include +#include + +#include "rdkit_extensions/mmff_flattened_builder.h" +#include "src/forcefields/mmff.h" +#include "src/forcefields/mmff_batched_forcefield.h" +#include "src/minimizer/fire_minimize_permol_kernels.h" +#include "src/utils/cuda_error_check.h" +#include "src/utils/device_vector.h" +#include "tests/test_utils.h" + +using nvMolKit::checkReturnCode; +using ::nvMolKit::MMFF::BatchedMolecularDeviceBuffers; +using ::nvMolKit::MMFF::BatchedMolecularSystemHost; + +namespace { + +void perturbConformer(RDKit::Conformer& conf, const float delta = 0.1, const int seed = 0) { + std::mt19937 gen(seed); + std::uniform_real_distribution dist(-delta, delta); + for (unsigned int i = 0; i < conf.getNumAtoms(); ++i) { + RDGeom::Point3D pos = conf.getAtomPos(i); + pos.x += delta * dist(gen); + pos.y += delta * dist(gen); + pos.z += delta * dist(gen); + conf.setAtomPos(i, pos); + } +} + +struct PerMolFireFixture { + std::vector> mols; + BatchedMolecularSystemHost systemHost; + BatchedMolecularDeviceBuffers systemDevice; + + void setup(int numMols) { + getMols(getTestDataFolderPath() + "/MMFF94_dative.sdf", mols, numMols); + int runningSeed = 0; + for (const auto& mol : mols) { + perturbConformer(mol->getConformer(), 0.3, runningSeed++); + std::vector positions(3 * mol->getNumAtoms()); + for (unsigned int i = 0; i < mol->getNumAtoms(); ++i) { + const RDGeom::Point3D pos = mol->getConformer().getAtomPos(i); + positions[3 * i] = pos.x; + positions[3 * i + 1] = pos.y; + positions[3 * i + 2] = pos.z; + } + const auto ffParams = nvMolKit::MMFF::constructForcefieldContribs(*mol); + nvMolKit::MMFF::addMoleculeToBatch(ffParams, positions, systemHost); + } + nvMolKit::MMFF::sendContribsAndIndicesToDevice(systemHost, systemDevice); + nvMolKit::MMFF::allocateIntermediateBuffers(systemHost, systemDevice); + systemDevice.energyOuts.zero(); + systemDevice.positions.setFromVector(systemHost.positions); + systemDevice.grad.resize(systemDevice.positions.size()); + systemDevice.grad.zero(); + } +}; + +std::vector computeReferenceEnergies(const std::vector>& mols) { + std::vector energies; + energies.reserve(mols.size()); + for (const auto& mol : mols) { + const auto molProps = std::make_unique(*mol); + const std::unique_ptr molFF(RDKit::MMFF::constructForceField(*mol, molProps.get())); + molFF->initialize(); + molFF->minimize(500, 1e-4); + energies.push_back(molFF->calcEnergy()); + } + return energies; +} + +int maxAtomsInBatch(const std::vector& atomStarts) { + int maxAtoms = 0; + for (size_t i = 0; i + 1 < atomStarts.size(); ++i) { + maxAtoms = std::max(maxAtoms, atomStarts[i + 1] - atomStarts[i]); + } + return maxAtoms; +} + +bool allConverged(const nvMolKit::AsyncDeviceVector& statuses) { + std::vector statusHost(statuses.size()); + statuses.copyToHost(statusHost); + cudaCheckError(cudaDeviceSynchronize()); + return std::all_of(statusHost.begin(), statusHost.end(), [](const uint8_t status) { return status == 0; }); +} + +} // namespace + +TEST(FireMinimizerPerMolMMFF, DirectLauncherConvergesNearReference) { + PerMolFireFixture fixture; + fixture.setup(/*numMols=*/4); + const std::vector refEnergies = computeReferenceEnergies(fixture.mols); + const int numMols = static_cast(fixture.systemHost.indices.atomStarts.size()) - 1; + + nvMolKit::FireOptions options{}; + options.useMass = false; + options.stuckDetectionEnabled = false; + options.gradTol = 1e-3; + options.dtInit = 0.05; + options.dMax = 0.2; + + std::vector molIdsHost(numMols); + std::iota(molIdsHost.begin(), molIdsHost.end(), 0); + + nvMolKit::AsyncDeviceVector molIds; + molIds.setFromVector(molIdsHost); + + nvMolKit::AsyncDeviceVector velocities; + velocities.resize(fixture.systemDevice.positions.size()); + velocities.zero(); + + nvMolKit::AsyncDeviceVector alphas; + alphas.setFromVector(std::vector(numMols, options.alphaInit)); + + nvMolKit::AsyncDeviceVector dts; + dts.setFromVector(std::vector(numMols, options.dtInit)); + + nvMolKit::AsyncDeviceVector nStepsPositive; + nStepsPositive.resize(numMols); + nStepsPositive.zero(); + + nvMolKit::AsyncDeviceVector statuses; + statuses.setFromVector(std::vector(numMols, 1)); + + auto terms = nvMolKit::MMFF::toEnergyForceContribsDevicePtr(fixture.systemDevice); + auto systemIndices = nvMolKit::MMFF::toBatchedIndicesDevicePtr(fixture.systemDevice); + const bool hasConstraints = nvMolKit::MMFF::batchHasConstraints(fixture.systemDevice.contribs); + + bool converged = false; + for (int extension = 0; extension < 5 && !converged; ++extension) { + const cudaError_t err = nvMolKit::launchFirePerMolKernel(numMols, + molIds.data(), + maxAtomsInBatch(fixture.systemHost.indices.atomStarts), + fixture.systemDevice.indices.atomStarts.data(), + options, + /*numIters=*/2000, + options.gradTol, + terms, + systemIndices, + hasConstraints, + fixture.systemDevice.positions.data(), + fixture.systemDevice.grad.data(), + velocities.data(), + alphas.data(), + dts.data(), + nStepsPositive.data(), + /*masses=*/nullptr, + fixture.systemDevice.energyOuts.data(), + statuses.data()); + cudaCheckError(err); + converged = allConverged(statuses); + } + + std::vector energiesHost(fixture.systemDevice.energyOuts.size()); + fixture.systemDevice.energyOuts.copyToHost(energiesHost); + cudaCheckError(cudaDeviceSynchronize()); + + for (size_t i = 0; i < energiesHost.size(); ++i) { + const double tolerance = 1.0 + 0.05 * std::abs(refEnergies[i]); + EXPECT_NEAR(energiesHost[i], refEnergies[i], tolerance) << "system " << i; + } +} + +TEST(FireMinimizerPerMolMMFF, DirectLauncherNoopsForEmptyBatch) { + nvMolKit::FireOptions options{}; + const cudaError_t err = nvMolKit::launchFirePerMolKernel(/*numMols=*/0, + /*molIds=*/nullptr, + /*maxAtoms=*/0, + /*atomStarts=*/nullptr, + options, + /*numIters=*/10, + /*gradTol=*/options.gradTol, + {}, + {}, + /*hasConstraints=*/false, + /*positions=*/nullptr, + /*grad=*/nullptr, + /*velocities=*/nullptr, + /*alphas=*/nullptr, + /*dts=*/nullptr, + /*nStepsPositive=*/nullptr, + /*masses=*/nullptr, + /*energyOuts=*/nullptr, + /*statuses=*/nullptr); + EXPECT_EQ(err, cudaSuccess); +}