#ifndef OPENMM_HIPKERNELS_H_
#define OPENMM_HIPKERNELS_H_
/* -------------------------------------------------------------------------- *
* OpenMM *
* -------------------------------------------------------------------------- *
* This is part of the OpenMM molecular simulation toolkit originating from *
* Simbios, the NIH National Center for Physics-Based Simulation of *
* Biological Structures at Stanford, funded under the NIH Roadmap for *
* Medical Research, grant U54 GM072970. See https://simtk.org. *
* *
* Portions copyright (c) 2008-2024 Stanford University and the Authors. *
* Portions copyright (c) 2020-2022 Advanced Micro Devices, Inc. *
* Authors: Peter Eastman, Nicholas Curtis *
* Contributors: *
* *
* This program is free software: you can redistribute it and/or modify *
* it under the terms of the GNU Lesser General Public License as published *
* by the Free Software Foundation, either version 3 of the License, or *
* (at your option) any later version. *
* *
* This program is distributed in the hope that it will be useful, *
* but WITHOUT ANY WARRANTY; without even the implied warranty of *
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the *
* GNU Lesser General Public License for more details. *
* *
* You should have received a copy of the GNU Lesser General Public License *
* along with this program. If not, see . *
* -------------------------------------------------------------------------- */
#include "HipPlatform.h"
#include "HipArray.h"
#include "HipContext.h"
#include "HipFFT3D.h"
#include "HipSort.h"
#include "openmm/kernels.h"
#include "openmm/System.h"
#include "openmm/common/CommonKernels.h"
namespace OpenMM {
/**
* This kernel is invoked at the beginning and end of force and energy computations. It gives the
* Platform a chance to clear buffers and do other initialization at the beginning, and to do any
* necessary work at the end to determine the final results.
*/
class HipCalcForcesAndEnergyKernel : public CalcForcesAndEnergyKernel {
public:
HipCalcForcesAndEnergyKernel(std::string name, const Platform& platform, HipContext& cu) : CalcForcesAndEnergyKernel(name, platform), cu(cu) {
}
/**
* Initialize the kernel.
*
* @param system the System this kernel will be applied to
*/
void initialize(const System& system);
/**
* This is called at the beginning of each force/energy computation, before calcForcesAndEnergy() has been called on
* any ForceImpl.
*
* @param context the context in which to execute this kernel
* @param includeForce true if forces should be computed
* @param includeEnergy true if potential energy should be computed
* @param groups a set of bit flags for which force groups to include
*/
void beginComputation(ContextImpl& context, bool includeForce, bool includeEnergy, int groups);
/**
* This is called at the end of each force/energy computation, after calcForcesAndEnergy() has been called on
* every ForceImpl.
*
* @param context the context in which to execute this kernel
* @param includeForce true if forces should be computed
* @param includeEnergy true if potential energy should be computed
* @param groups a set of bit flags for which force groups to include
* @param valid the method may set this to false to indicate the results are invalid and the force/energy
* calculation should be repeated
* @return the potential energy of the system. This value is added to all values returned by ForceImpls'
* calcForcesAndEnergy() methods. That is, each force kernel may either return its contribution to the
* energy directly, or add it to an internal buffer so that it will be included here.
*/
double finishComputation(ContextImpl& context, bool includeForce, bool includeEnergy, int groups, bool& valid);
private:
HipContext& cu;
};
/**
* This kernel is invoked by NonbondedForce to calculate the forces acting on the system.
*/
class HipCalcNonbondedForceKernel : public CalcNonbondedForceKernel {
public:
HipCalcNonbondedForceKernel(std::string name, const Platform& platform, HipContext& cu, const System& system) : CalcNonbondedForceKernel(name, platform),
cu(cu), hasInitializedFFT(false), sort(NULL), dispersionFft(NULL), fft(NULL), pmeio(NULL), usePmeStream(false) {
}
~HipCalcNonbondedForceKernel();
/**
* Initialize the kernel.
*
* @param system the System this kernel will be applied to
* @param force the NonbondedForce this kernel will be used for
*/
void initialize(const System& system, const NonbondedForce& force);
/**
* Execute the kernel to calculate the forces and/or energy.
*
* @param context the context in which to execute this kernel
* @param includeForces true if forces should be calculated
* @param includeEnergy true if the energy should be calculated
* @param includeDirect true if direct space interactions should be included
* @param includeReciprocal true if reciprocal space interactions should be included
* @return the potential energy due to the force
*/
double execute(ContextImpl& context, bool includeForces, bool includeEnergy, bool includeDirect, bool includeReciprocal);
/**
* Copy changed parameters over to a context.
*
* @param context the context to copy parameters to
* @param force the NonbondedForce to copy the parameters from
*/
void copyParametersToContext(ContextImpl& context, const NonbondedForce& force);
/**
* Get the parameters being used for PME.
*
* @param alpha the separation parameter
* @param nx the number of grid points along the X axis
* @param ny the number of grid points along the Y axis
* @param nz the number of grid points along the Z axis
*/
void getPMEParameters(double& alpha, int& nx, int& ny, int& nz) const;
/**
* Get the dispersion parameters being used for the dispersion term in LJPME.
*
* @param alpha the separation parameter
* @param nx the number of grid points along the X axis
* @param ny the number of grid points along the Y axis
* @param nz the number of grid points along the Z axis
*/
void getLJPMEParameters(double& alpha, int& nx, int& ny, int& nz) const;
private:
class SortTrait : public HipSort::SortTrait {
int getDataSize() const {return 8;}
int getKeySize() const {return 4;}
const char* getDataType() const {return "int2";}
const char* getKeyType() const {return "int";}
const char* getMinKey() const {return "(-2147483647-1)";}
const char* getMaxKey() const {return "2147483647";}
const char* getMaxValue() const {return "make_int2(2147483647, 2147483647)";}
const char* getSortKey() const {return "value.y";}
};
class ForceInfo;
class PmeIO;
class PmePreComputation;
class PmePostComputation;
class SyncStreamPreComputation;
class SyncStreamPostComputation;
HipContext& cu;
ForceInfo* info;
bool hasInitializedFFT;
HipArray charges;
HipArray sigmaEpsilon;
HipArray exceptionParams;
HipArray exclusionAtoms;
HipArray exclusionParams;
HipArray baseParticleParams;
HipArray baseExceptionParams;
HipArray particleParamOffsets;
HipArray exceptionParamOffsets;
HipArray particleOffsetIndices;
HipArray exceptionOffsetIndices;
HipArray globalParams;
HipArray cosSinSums;
HipArray pmeGrid1;
HipArray pmeGrid2;
HipArray pmeBsplineModuliX;
HipArray pmeBsplineModuliY;
HipArray pmeBsplineModuliZ;
HipArray pmeDispersionBsplineModuliX;
HipArray pmeDispersionBsplineModuliY;
HipArray pmeDispersionBsplineModuliZ;
HipArray pmeAtomGridIndex;
HipArray pmeEnergyBuffer;
HipSort* sort;
Kernel cpuPme;
PmeIO* pmeio;
hipStream_t pmeStream;
hipEvent_t pmeSyncEvent, paramsSyncEvent;
HipFFT3D* fft;
HipFFT3D* dispersionFft;
hipFunction_t computeParamsKernel, computeExclusionParamsKernel;
hipFunction_t ewaldSumsKernel;
hipFunction_t ewaldForcesKernel;
hipFunction_t pmeGridIndexKernel;
hipFunction_t pmeDispersionGridIndexKernel;
hipFunction_t pmeSpreadChargeKernel;
hipFunction_t pmeDispersionSpreadChargeKernel;
hipFunction_t pmeFinishSpreadChargeKernel;
hipFunction_t pmeDispersionFinishSpreadChargeKernel;
hipFunction_t pmeEvalEnergyKernel;
hipFunction_t pmeEvalDispersionEnergyKernel;
hipFunction_t pmeConvolutionKernel;
hipFunction_t pmeDispersionConvolutionKernel;
hipFunction_t pmeInterpolateForceKernel;
hipFunction_t pmeInterpolateDispersionForceKernel;
std::vector > exceptionAtoms;
std::vector paramNames;
std::vector paramValues;
double ewaldSelfEnergy, dispersionCoefficient, alpha, dispersionAlpha;
int interpolateForceThreads;
int gridSizeX, gridSizeY, gridSizeZ;
int dispersionGridSizeX, dispersionGridSizeY, dispersionGridSizeZ;
bool hasCoulomb, hasLJ, usePmeStream, doLJPME, usePosqCharges, recomputeParams, hasOffsets;
NonbondedMethod nonbondedMethod;
static const int PmeOrder = 5;
};
/**
* This kernel is invoked by CustomCVForce to calculate the forces acting on the system and the energy of the system.
*/
class HipCalcCustomCVForceKernel : public CommonCalcCustomCVForceKernel {
public:
HipCalcCustomCVForceKernel(std::string name, const Platform& platform, ComputeContext& cc) : CommonCalcCustomCVForceKernel(name, platform, cc) {
}
ComputeContext& getInnerComputeContext(ContextImpl& innerContext) {
return *reinterpret_cast(innerContext.getPlatformData())->contexts[0];
}
};
/**
* This kernel is invoked by ATMForce to calculate the forces acting on the system and the energy of the system.
*/
class HipCalcATMForceKernel : public CommonCalcATMForceKernel {
public:
HipCalcATMForceKernel(std::string name, const Platform& platform, ComputeContext& cc) : CommonCalcATMForceKernel(name, platform, cc) {
}
ComputeContext& getInnerComputeContext(ContextImpl& innerContext) {
return *reinterpret_cast(innerContext.getPlatformData())->contexts[0];
}
};
} // namespace OpenMM
#endif /*OPENMM_HIPKERNELS_H_*/