Skip to content

GitLab

Unverified Commit 5d8d5874 authored by Peter Eastman's avatar Peter Eastman Committed by GitHub
Browse files

Avoid multiple forces running on the worker thread at once (#5243)

parent 1c528ca8
Hide whitespace changes
Inline Side-by-side
Showing with 47 additions and 37 deletions
olla/include/openmm/kernels.h
View file @ 5d8d5874
......@@ -1994,10 +1994,10 @@ public:
/**
* Initialize the kernel.
*
* @param system the System this kernel will be applied to
* @param context the ContextImpl this kernel will be applied to
* @param force the CustomCPPForceImpl this kernel will be used for
*/
virtual void initialize(const System& system, CustomCPPForceImpl& force) = 0;
virtual void initialize(const ContextImpl& context, CustomCPPForceImpl& force) = 0;
/**
* Execute the kernel to calculate the forces and/or energy.
*
......@@ -2022,10 +2022,10 @@ public:
/**
* Initialize the kernel.
*
* @param system the System this kernel will be applied to
* @param context the ContextImpl this kernel will be applied to
* @param force the PythonForce this kernel will be used for
*/
virtual void initialize(const System& system, const PythonForce& force) = 0;
virtual void initialize(const ContextImpl& context, const PythonForce& force) = 0;
/**
* Execute the kernel to calculate the forces and/or energy.
*
......
openmmapi/src/CustomCPPForceImpl.cpp
View file @ 5d8d5874
......@@ -4,7 +4,7 @@
* This is part of the OpenMM molecular simulation toolkit. *
* See https://openmm.org/development. *
* *
* Portions copyright (c) 2008-2021 Stanford University and the Authors. *
* Portions copyright (c) 2008-2026 Stanford University and the Authors. *
* Authors: Peter Eastman *
* Contributors: *
* *
......@@ -40,7 +40,7 @@ CustomCPPForceImpl::CustomCPPForceImpl(const Force& owner) {
void CustomCPPForceImpl::initialize(ContextImpl& context) {
kernel = context.getPlatform().createKernel(CalcCustomCPPForceKernel::Name(), context);
kernel.getAs<CalcCustomCPPForceKernel>().initialize(context.getSystem(), *this);
kernel.getAs<CalcCustomCPPForceKernel>().initialize(context, *this);
}
double CustomCPPForceImpl::calcForcesAndEnergy(ContextImpl& context, bool includeForces, bool includeEnergy, int groups) {
......
openmmapi/src/PythonForceImpl.cpp
View file @ 5d8d5874
......@@ -4,7 +4,7 @@
* This is part of the OpenMM molecular simulation toolkit. *
* See https://openmm.org/development. *
* *
* Portions copyright (c) 2025 Stanford University and the Authors. *
* Portions copyright (c) 2025-2026 Stanford University and the Authors. *
* Authors: Peter Eastman *
* Contributors: *
* *
......@@ -47,7 +47,7 @@ PythonForceImpl::~PythonForceImpl() {
void PythonForceImpl::initialize(ContextImpl& context) {
kernel = context.getPlatform().createKernel(CalcPythonForceKernel::Name(), context);
kernel.getAs<CalcPythonForceKernel>().initialize(context.getSystem(), owner);
kernel.getAs<CalcPythonForceKernel>().initialize(context, owner);
}
double PythonForceImpl::calcForcesAndEnergy(ContextImpl& context, bool includeForces, bool includeEnergy, int groups) {
......
platforms/common/include/openmm/common/CommonKernels.h
View file @ 5d8d5874
......@@ -7,7 +7,7 @@
* This is part of the OpenMM molecular simulation toolkit. *
* See https://openmm.org/development. *
* *
* Portions copyright (c) 2008-2025 Stanford University and the Authors. *
* Portions copyright (c) 2008-2026 Stanford University and the Authors. *
* Authors: Peter Eastman *
* Contributors: *
* *
......@@ -1469,10 +1469,10 @@ public:
/**
* Initialize the kernel.
*
* @param system the System this kernel will be applied to
* @param context the ContextImpl this kernel will be applied to
* @param force the CustomCPPForceImpl this kernel will be used for
*/
void initialize(const System& system, CustomCPPForceImpl& force);
void initialize(const ContextImpl& context, CustomCPPForceImpl& force);
/**
* Execute the kernel to calculate the forces and/or energy.
*
......@@ -1507,6 +1507,7 @@ private:
std::vector<float> floatForces;
int forceGroupFlag;
double energy;
bool useWorkerThread;
};
/**
......@@ -1520,10 +1521,10 @@ public:
/**
* Initialize the kernel.
*
* @param system the System this kernel will be applied to
* @param context the ContextImpl this kernel will be applied to
* @param force the PythonForce this kernel will be used for
*/
void initialize(const System& system, const PythonForce& force);
void initialize(const ContextImpl& context, const PythonForce& force);
/**
* Execute the kernel to calculate the forces and/or energy.
*
......@@ -1558,7 +1559,7 @@ private:
std::vector<double> forcesVec;
int forceGroupFlag;
double energy;
bool usePeriodic;
bool usePeriodic, useWorkerThread;
};
} // namespace OpenMM
......
platforms/common/src/CommonKernels.cpp
View file @ 5d8d5874
......@@ -4,7 +4,7 @@
* This is part of the OpenMM molecular simulation toolkit. *
* See https://openmm.org/development. *
* *
* Portions copyright (c) 2008-2025 Stanford University and the Authors. *
* Portions copyright (c) 2008-2026 Stanford University and the Authors. *
* Authors: Peter Eastman *
* Contributors: *
* *
......@@ -34,6 +34,7 @@
#include "openmm/internal/CustomCompoundBondForceImpl.h"
#include "openmm/internal/DPDIntegratorUtilities.h"
#include "openmm/internal/OSRngSeed.h"
#include "openmm/internal/PythonForceImpl.h"
#include "openmm/internal/ThreadPool.h"
#include "openmm/internal/timer.h"
#include "CommonKernelSources.h"
......@@ -4688,10 +4689,10 @@ public:
CommonCalcCustomCPPForceKernel& owner;
};
void CommonCalcCustomCPPForceKernel::initialize(const System& system, CustomCPPForceImpl& force) {
void CommonCalcCustomCPPForceKernel::initialize(const ContextImpl& context, CustomCPPForceImpl& force) {
ContextSelector selector(cc);
this->force = &force;
int numParticles = system.getNumParticles();
int numParticles = context.getSystem().getNumParticles();
forcesVec.resize(numParticles);
positionsVec.resize(numParticles);
floatForces.resize(3*numParticles);
......@@ -4706,14 +4707,18 @@ void CommonCalcCustomCPPForceKernel::initialize(const System& system, CustomCPPF
addForcesKernel->addArg(cc.getLongForceBuffer());
addForcesKernel->addArg(cc.getAtomIndexArray());
forceGroupFlag = (1<<force.getOwner().getForceGroup());
if (cc.getNumContexts() == 1) {
useWorkerThread = (cc.getNumContexts() == 1);
for (const ForceImpl* impl : context.getForceImpls())
if (dynamic_cast<const CustomCPPForceImpl*>(impl) != NULL || dynamic_cast<const PythonForceImpl*>(impl) != NULL)
useWorkerThread = false;
if (useWorkerThread) {
cc.addPreComputation(new StartCalculationPreComputation(*this));
cc.addPostComputation(new AddForcesPostComputation(*this));
}
}
double CommonCalcCustomCPPForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
if (cc.getNumContexts() == 1) {
if (useWorkerThread) {
// This method does nothing. The actual calculation is started by the pre-computation, continued on
// the worker thread, and finished by the post-computation.
......@@ -4765,7 +4770,7 @@ double CommonCalcCustomCPPForceKernel::addForces(bool includeForces, bool includ
// Wait until executeOnWorkerThread() is finished.
if (cc.getNumContexts() == 1)
if (useWorkerThread)
cc.getWorkThread().flush();
// Add in the forces.
......@@ -4811,11 +4816,11 @@ public:
CommonCalcPythonForceKernel& owner;
};
void CommonCalcPythonForceKernel::initialize(const System& system, const PythonForce& force) {
void CommonCalcPythonForceKernel::initialize(const ContextImpl& context, const PythonForce& force) {
ContextSelector selector(cc);
computation = &force.getComputation();
usePeriodic = force.usesPeriodicBoundaryConditions();
int numParticles = system.getNumParticles();
int numParticles = context.getSystem().getNumParticles();
positionsVec.resize(numParticles);
forcesVec.resize(3*numParticles);
int elementSize = (cc.getUseDoublePrecision() ? sizeof(double) : sizeof(float));
......@@ -4829,14 +4834,18 @@ void CommonCalcPythonForceKernel::initialize(const System& system, const PythonF
addForcesKernel->addArg(cc.getLongForceBuffer());
addForcesKernel->addArg(cc.getAtomIndexArray());
forceGroupFlag = (1<<force.getForceGroup());
if (cc.getNumContexts() == 1) {
useWorkerThread = (cc.getNumContexts() == 1);
for (const ForceImpl* impl : context.getForceImpls())
if (dynamic_cast<const CustomCPPForceImpl*>(impl) != NULL || dynamic_cast<const PythonForceImpl*>(impl) != NULL)
useWorkerThread = false;
if (useWorkerThread) {
cc.addPreComputation(new StartCalculationPreComputation(*this));
cc.addPostComputation(new AddForcesPostComputation(*this));
}
}
double CommonCalcPythonForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
if (cc.getNumContexts() == 1) {
if (useWorkerThread) {
// This method does nothing. The actual calculation is started by the pre-computation, continued on
// the worker thread, and finished by the post-computation.
......@@ -4857,9 +4866,9 @@ double CommonCalcPythonForceKernel::execute(ContextImpl& context, bool includeFo
void CommonCalcPythonForceKernel::beginComputation(bool includeForces, bool includeEnergy, int groups) {
if ((groups&forceGroupFlag) == 0)
return;
// The actual force computation will be done on a different thread.
cc.getWorkThread().addTask(new ExecuteTask(*this, includeForces));
}
......@@ -4886,8 +4895,8 @@ double CommonCalcPythonForceKernel::addForces(bool includeForces, bool includeEn
return 0;
// Wait until executeOnWorkerThread() is finished.
if (cc.getNumContexts() == 1)
if (useWorkerThread)
cc.getWorkThread().flush();
// Add in the forces.
......@@ -4896,7 +4905,7 @@ double CommonCalcPythonForceKernel::addForces(bool includeForces, bool includeEn
ContextSelector selector(cc);
addForcesKernel->execute(cc.getNumAtoms());
}
// Return the energy.
return energy;
......
platforms/reference/include/ReferenceKernels.h
View file @ 5d8d5874
......@@ -1992,10 +1992,10 @@ public:
/**
* Initialize the kernel.
*
* @param system the System this kernel will be applied to
* @param context the ContextImpl this kernel will be applied to
* @param force the CustomCPPForceImpl this kernel will be used for
*/
void initialize(const System& system, CustomCPPForceImpl& force);
void initialize(const ContextImpl& context, CustomCPPForceImpl& force);
/**
* Execute the kernel to calculate the forces and/or energy.
*
......@@ -2020,10 +2020,10 @@ public:
/**
* Initialize the kernel.
*
* @param system the System this kernel will be applied to
* @param context the ContextImpl this kernel will be applied to
* @param force the PythonForce this kernel will be used for
*/
void initialize(const System& system, const PythonForce& force);
void initialize(const ContextImpl& context, const PythonForce& force);
/**
* Execute the kernel to calculate the forces and/or energy.
*
......
platforms/reference/src/ReferenceKernels.cpp
View file @ 5d8d5874
......@@ -3535,9 +3535,9 @@ void ReferenceCalcATMForceKernel::copyParametersToContext(ContextImpl& context,
loadParams(numParticles, force);
}
void ReferenceCalcCustomCPPForceKernel::initialize(const System& system, CustomCPPForceImpl& force) {
void ReferenceCalcCustomCPPForceKernel::initialize(const ContextImpl& context, CustomCPPForceImpl& force) {
this->force = &force;
forces.resize(system.getNumParticles());
forces.resize(context.getSystem().getNumParticles());
}
double ReferenceCalcCustomCPPForceKernel::execute(ContextImpl& context, bool includeForces, bool includeEnergy) {
......@@ -3550,9 +3550,9 @@ double ReferenceCalcCustomCPPForceKernel::execute(ContextImpl& context, bool inc
return energy;
}
void ReferenceCalcPythonForceKernel::initialize(const System& system, const PythonForce& force) {
void ReferenceCalcPythonForceKernel::initialize(const ContextImpl& context, const PythonForce& force) {
computation = &force.getComputation();
forces.resize(system.getNumParticles());
forces.resize(context.getSystem().getNumParticles());
usePeriodic = force.usesPeriodicBoundaryConditions();
}
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment