Skip to content

GitLab

Unverified Commit 01e99e77 authored by Peter Eastman's avatar Peter Eastman Committed by GitHub
Browse files

Uniform interface for FFTs (#4911) 

* Unified interface for FFTs

* AMOEBA uses unified interface for FFTs

* HIP implementation of common FFT interface
parent a3909c8e
Hide whitespace changes
Inline Side-by-side
Showing with 15 additions and 241 deletions
plugins/amoeba/platforms/common/src/AmoebaCommonKernels.h
View file @ 01e99e77
......@@ -9,7 +9,7 @@
* Biological Structures at Stanford, funded under the NIH Roadmap for *
* Medical Research, grant U54 GM072970. See https://simtk.org. *
* *
* Portions copyright (c) 2008-2021 Stanford University and the Authors. *
* Portions copyright (c) 2008-2025 Stanford University and the Authors. *
* Authors: Mark Friedrichs, Peter Eastman *
* Contributors: *
* *
......@@ -32,6 +32,7 @@
#include "openmm/System.h"
#include "openmm/common/ComputeContext.h"
#include "openmm/common/ComputeArray.h"
#include "openmm/common/FFT3D.h"
#include "openmm/common/NonbondedUtilities.h"
namespace OpenMM {
......@@ -152,10 +153,6 @@ public:
* @param nz the number of grid points along the Z axis
*/
void getPMEParameters(double& alpha, int& nx, int& ny, int& nz) const;
/**
* Compute the FFT.
*/
virtual void computeFFT(bool forward) = 0;
/**
* Get whether charge spreading should be done in fixed point.
*/
......@@ -238,6 +235,7 @@ protected:
ComputeKernel pmeFixedPotentialKernel, pmeInducedPotentialKernel, pmeFixedForceKernel, pmeInducedForceKernel, pmeRecordInducedFieldDipolesKernel;
ComputeKernel pmeTransformMultipolesKernel, pmeTransformPotentialKernel;
ComputeEvent syncEvent;
FFT3D* fft;
CommonCalcAmoebaGeneralizedKirkwoodForceKernel* gkKernel;
static const int PmeOrder = 5;
static const int MaxPrevDIISDipoles = 20;
......@@ -424,6 +422,7 @@ private:
class CommonCalcHippoNonbondedForceKernel : public CalcHippoNonbondedForceKernel {
public:
CommonCalcHippoNonbondedForceKernel(const std::string& name, const Platform& platform, ComputeContext& cc, const System& system);
virtual ~CommonCalcHippoNonbondedForceKernel();
/**
* Initialize the kernel.
*
......@@ -431,10 +430,6 @@ public:
* @param force the HippoNonbondedForce this kernel will be used for
*/
void initialize(const System& system, const HippoNonbondedForce& force);
/**
* Compute the FFT.
*/
virtual void computeFFT(bool forward, bool dispersion) = 0;
/**
* Get whether charge spreading should be done in fixed point.
*/
......@@ -534,6 +529,8 @@ protected:
ComputeArray lastPositions;
ComputeArray exceptionScales[6];
ComputeArray exceptionAtoms;
FFT3D* fft;
FFT3D* dfft;
ComputeKernel computeMomentsKernel, recordInducedDipolesKernel, mapTorqueKernel;
ComputeKernel fixedFieldKernel, fixedFieldExceptionKernel, mutualFieldKernel, mutualFieldExceptionKernel, computeExceptionsKernel;
ComputeKernel pmeSpreadFixedMultipolesKernel, pmeSpreadInducedDipolesKernel, pmeFinishSpreadChargeKernel, pmeConvolutionKernel;
......
plugins/amoeba/platforms/cuda/src/AmoebaCudaKernels.cpp
View file @ 01e99e77
......@@ -6,7 +6,7 @@
* Biological Structures at Stanford, funded under the NIH Roadmap for *
* Medical Research, grant U54 GM072970. See https://simtk.org. *
* *
* Portions copyright (c) 2008-2021 Stanford University and the Authors. *
* Portions copyright (c) 2008-2025 Stanford University and the Authors. *
* Authors: Peter Eastman, Mark Friedrichs *
* Contributors: *
* *
......@@ -38,7 +38,6 @@
#include "openmm/internal/AmoebaVdwForceImpl.h"
#include "openmm/internal/NonbondedForceImpl.h"
#include "CudaBondedUtilities.h"
#include "CudaFFT3D.h"
#include "CudaForceInfo.h"
#include "CudaKernelSources.h"
#include "SimTKOpenMMRealType.h"
......@@ -79,44 +78,6 @@ static void setPeriodicBoxArgs(ComputeContext& cc, ComputeKernel kernel, int ind
}
}
/* -------------------------------------------------------------------------- *
* AmoebaMultipole *
* -------------------------------------------------------------------------- */
CudaCalcAmoebaMultipoleForceKernel::~CudaCalcAmoebaMultipoleForceKernel() {
ContextSelector selector(cc);
if (hasInitializedFFT)
cufftDestroy(fft);
}
void CudaCalcAmoebaMultipoleForceKernel::initialize(const System& system, const AmoebaMultipoleForce& force) {
CommonCalcAmoebaMultipoleForceKernel::initialize(system, force);
if (usePME) {
ContextSelector selector(cc);
cufftResult result = cufftPlan3d(&fft, gridSizeX, gridSizeY, gridSizeZ, cc.getUseDoublePrecision() ? CUFFT_Z2Z : CUFFT_C2C);
if (result != CUFFT_SUCCESS)
throw OpenMMException("Error initializing FFT: "+cc.intToString(result));
hasInitializedFFT = true;
}
}
void CudaCalcAmoebaMultipoleForceKernel::computeFFT(bool forward) {
CudaArray& grid1 = dynamic_cast<CudaContext&>(cc).unwrap(pmeGrid1);
CudaArray& grid2 = dynamic_cast<CudaContext&>(cc).unwrap(pmeGrid2);
if (forward) {
if (cc.getUseDoublePrecision())
cufftExecZ2Z(fft, (double2*) grid1.getDevicePointer(), (double2*) grid2.getDevicePointer(), CUFFT_FORWARD);
else
cufftExecC2C(fft, (float2*) grid1.getDevicePointer(), (float2*) grid2.getDevicePointer(), CUFFT_FORWARD);
}
else {
if (cc.getUseDoublePrecision())
cufftExecZ2Z(fft, (double2*) grid2.getDevicePointer(), (double2*) grid1.getDevicePointer(), CUFFT_INVERSE);
else
cufftExecC2C(fft, (float2*) grid2.getDevicePointer(), (float2*) grid1.getDevicePointer(), CUFFT_INVERSE);
}
}
/* -------------------------------------------------------------------------- *
* HippoNonbondedForce *
* -------------------------------------------------------------------------- */
......@@ -125,12 +86,6 @@ CudaCalcHippoNonbondedForceKernel::~CudaCalcHippoNonbondedForceKernel() {
ContextSelector selector(cc);
if (sort != NULL)
delete sort;
if (hasInitializedFFT) {
cufftDestroy(fftForward);
cufftDestroy(fftBackward);
cufftDestroy(dfftForward);
cufftDestroy(dfftBackward);
}
}
void CudaCalcHippoNonbondedForceKernel::initialize(const System& system, const HippoNonbondedForce& force) {
......@@ -139,38 +94,6 @@ void CudaCalcHippoNonbondedForceKernel::initialize(const System& system, const H
ContextSelector selector(cc);
CudaContext& cu = dynamic_cast<CudaContext&>(cc);
sort = new CudaSort(cu, new SortTrait(), cc.getNumAtoms());
cufftResult result = cufftPlan3d(&fftForward, gridSizeX, gridSizeY, gridSizeZ, cc.getUseDoublePrecision() ? CUFFT_D2Z : CUFFT_R2C);
if (result != CUFFT_SUCCESS)
throw OpenMMException("Error initializing FFT: "+cc.intToString(result));
result = cufftPlan3d(&fftBackward, gridSizeX, gridSizeY, gridSizeZ, cc.getUseDoublePrecision() ? CUFFT_Z2D : CUFFT_C2R);
if (result != CUFFT_SUCCESS)
throw OpenMMException("Error initializing FFT: "+cc.intToString(result));
result = cufftPlan3d(&dfftForward, dispersionGridSizeX, dispersionGridSizeY, dispersionGridSizeZ, cc.getUseDoublePrecision() ? CUFFT_D2Z : CUFFT_R2C);
if (result != CUFFT_SUCCESS)
throw OpenMMException("Error initializing FFT: "+cc.intToString(result));
result = cufftPlan3d(&dfftBackward, dispersionGridSizeX, dispersionGridSizeY, dispersionGridSizeZ, cc.getUseDoublePrecision() ? CUFFT_Z2D : CUFFT_C2R);
if (result != CUFFT_SUCCESS)
throw OpenMMException("Error initializing FFT: "+cc.intToString(result));
hasInitializedFFT = true;
}
}
void CudaCalcHippoNonbondedForceKernel::computeFFT(bool forward, bool dispersion) {
CudaArray& grid1 = dynamic_cast<CudaContext&>(cc).unwrap(pmeGrid1);
CudaArray& grid2 = dynamic_cast<CudaContext&>(cc).unwrap(pmeGrid2);
if (forward) {
cufftHandle fft = dispersion ? dfftForward : fftForward;
if (cc.getUseDoublePrecision())
cufftExecD2Z(fft, (double*) grid1.getDevicePointer(), (double2*) grid2.getDevicePointer());
else
cufftExecR2C(fft, (float*) grid1.getDevicePointer(), (float2*) grid2.getDevicePointer());
}
else {
cufftHandle fft = dispersion ? dfftBackward : fftBackward;
if (cc.getUseDoublePrecision())
cufftExecZ2D(fft, (double2*) grid2.getDevicePointer(), (double*) grid1.getDevicePointer());
else
cufftExecC2R(fft, (float2*) grid2.getDevicePointer(), (float*) grid1.getDevicePointer());
}
}
......
plugins/amoeba/platforms/cuda/src/AmoebaCudaKernels.h
View file @ 01e99e77
......@@ -9,7 +9,7 @@
* Biological Structures at Stanford, funded under the NIH Roadmap for *
* Medical Research, grant U54 GM072970. See https://simtk.org. *
* *
* Portions copyright (c) 2008-2021 Stanford University and the Authors. *
* Portions copyright (c) 2008-2025 Stanford University and the Authors. *
* Authors: Mark Friedrichs, Peter Eastman *
* Contributors: *
* *
......@@ -34,7 +34,6 @@
#include "CudaNonbondedUtilities.h"
#include "CudaSort.h"
#include "AmoebaCommonKernels.h"
#include <cufft.h>
namespace OpenMM {
......@@ -44,29 +43,14 @@ namespace OpenMM {
class CudaCalcAmoebaMultipoleForceKernel : public CommonCalcAmoebaMultipoleForceKernel {
public:
CudaCalcAmoebaMultipoleForceKernel(const std::string& name, const Platform& platform, CudaContext& cu, const System& system) :
CommonCalcAmoebaMultipoleForceKernel(name, platform, cu, system), hasInitializedFFT(false) {
CommonCalcAmoebaMultipoleForceKernel(name, platform, cu, system) {
}
~CudaCalcAmoebaMultipoleForceKernel();
/**
* Initialize the kernel.
*
* @param system the System this kernel will be applied to
* @param force the AmoebaMultipoleForce this kernel will be used for
*/
void initialize(const System& system, const AmoebaMultipoleForce& force);
/**
* Compute the FFT.
*/
void computeFFT(bool forward);
/**
* Get whether charge spreading should be done in fixed point.
*/
bool useFixedPointChargeSpreading() const {
return cc.getUseDoublePrecision();
}
private:
bool hasInitializedFFT;
cufftHandle fft;
};
/**
......@@ -75,7 +59,7 @@ private:
class CudaCalcHippoNonbondedForceKernel : public CommonCalcHippoNonbondedForceKernel {
public:
CudaCalcHippoNonbondedForceKernel(const std::string& name, const Platform& platform, CudaContext& cu, const System& system) :
CommonCalcHippoNonbondedForceKernel(name, platform, cu, system), sort(NULL), hasInitializedFFT(false) {
CommonCalcHippoNonbondedForceKernel(name, platform, cu, system), sort(NULL) {
}
~CudaCalcHippoNonbondedForceKernel();
/**
......@@ -85,10 +69,6 @@ public:
* @param force the HippoNonbondedForce this kernel will be used for
*/
void initialize(const System& system, const HippoNonbondedForce& force);
/**
* Compute the FFT.
*/
void computeFFT(bool forward, bool dispersion);
/**
* Get whether charge spreading should be done in fixed point.
*/
......@@ -110,9 +90,7 @@ private:
const char* getMaxValue() const {return "make_int2(2147483647, 2147483647)";}
const char* getSortKey() const {return "value.y";}
};
bool hasInitializedFFT;
CudaSort* sort;
cufftHandle fftForward, fftBackward, dfftForward, dfftBackward;
};
} // namespace OpenMM
......
plugins/amoeba/platforms/hip/src/AmoebaHipKernels.cpp
View file @ 01e99e77
......@@ -38,7 +38,6 @@
#include "openmm/internal/AmoebaVdwForceImpl.h"
#include "openmm/internal/NonbondedForceImpl.h"
#include "HipBondedUtilities.h"
#include "HipFFT3D.h"
#include "HipForceInfo.h"
#include "HipKernelSources.h"
#include "SimTKOpenMMRealType.h"
......@@ -53,30 +52,6 @@
using namespace OpenMM;
using namespace std;
/* -------------------------------------------------------------------------- *
* AmoebaMultipole *
* -------------------------------------------------------------------------- */
HipCalcAmoebaMultipoleForceKernel::~HipCalcAmoebaMultipoleForceKernel() {
ContextSelector selector(cc);
if (fft != NULL)
delete fft;
}
void HipCalcAmoebaMultipoleForceKernel::initialize(const System& system, const AmoebaMultipoleForce& force) {
CommonCalcAmoebaMultipoleForceKernel::initialize(system, force);
if (usePME) {
ContextSelector selector(cc);
HipArray& grid1 = cu.unwrap(pmeGrid1);
HipArray& grid2 = cu.unwrap(pmeGrid2);
fft = new HipFFT3D(cu, gridSizeX, gridSizeY, gridSizeZ, false, cu.getCurrentStream(), grid1, grid2);
}
}
void HipCalcAmoebaMultipoleForceKernel::computeFFT(bool forward) {
fft->execFFT(forward);
}
/* -------------------------------------------------------------------------- *
* HippoNonbondedForce *
* -------------------------------------------------------------------------- */
......@@ -85,10 +60,6 @@ HipCalcHippoNonbondedForceKernel::~HipCalcHippoNonbondedForceKernel() {
ContextSelector selector(cc);
if (sort != NULL)
delete sort;
if (fft != NULL)
delete fft;
if (dfft != NULL)
delete dfft;
}
void HipCalcHippoNonbondedForceKernel::initialize(const System& system, const HippoNonbondedForce& force) {
......@@ -96,19 +67,6 @@ void HipCalcHippoNonbondedForceKernel::initialize(const System& system, const Hi
if (usePME) {
ContextSelector selector(cc);
sort = new HipSort(cu, new SortTrait(), cc.getNumAtoms());
HipArray& grid1 = cu.unwrap(pmeGrid1);
HipArray& grid2 = cu.unwrap(pmeGrid2);
fft = new HipFFT3D(cu, gridSizeX, gridSizeY, gridSizeZ, true, cu.getCurrentStream(), grid1, grid2);
dfft = new HipFFT3D(cu, dispersionGridSizeX, dispersionGridSizeY, dispersionGridSizeZ, true, cu.getCurrentStream(), grid1, grid2);
}
}
void HipCalcHippoNonbondedForceKernel::computeFFT(bool forward, bool dispersion) {
if (dispersion) {
dfft->execFFT(forward);
}
else {
fft->execFFT(forward);
}
}
......
plugins/amoeba/platforms/hip/src/AmoebaHipKernels.h
View file @ 01e99e77
......@@ -34,7 +34,6 @@
#include "HipContext.h"
#include "HipNonbondedUtilities.h"
#include "HipSort.h"
#include "HipFFT3D.h"
#include "AmoebaCommonKernels.h"
namespace OpenMM {
......@@ -45,20 +44,8 @@ namespace OpenMM {
class HipCalcAmoebaMultipoleForceKernel : public CommonCalcAmoebaMultipoleForceKernel {
public:
HipCalcAmoebaMultipoleForceKernel(const std::string& name, const Platform& platform, HipContext& cu, const System& system) :
CommonCalcAmoebaMultipoleForceKernel(name, platform, cu, system), cu(cu), fft(NULL) {
CommonCalcAmoebaMultipoleForceKernel(name, platform, cu, system), cu(cu) {
}
~HipCalcAmoebaMultipoleForceKernel();
/**
* Initialize the kernel.
*
* @param system the System this kernel will be applied to
* @param force the AmoebaMultipoleForce this kernel will be used for
*/
void initialize(const System& system, const AmoebaMultipoleForce& force);
/**
* Compute the FFT.
*/
void computeFFT(bool forward);
/**
* Get whether charge spreading should be done in fixed point.
*/
......@@ -67,7 +54,6 @@ public:
}
private:
HipContext& cu;
HipFFT3D* fft;
};
/**
......@@ -76,7 +62,7 @@ private:
class HipCalcHippoNonbondedForceKernel : public CommonCalcHippoNonbondedForceKernel {
public:
HipCalcHippoNonbondedForceKernel(const std::string& name, const Platform& platform, HipContext& cu, const System& system) :
CommonCalcHippoNonbondedForceKernel(name, platform, cu, system), cu(cu), sort(NULL), fft(NULL), dfft(NULL) {
CommonCalcHippoNonbondedForceKernel(name, platform, cu, system), cu(cu), sort(NULL) {
}
~HipCalcHippoNonbondedForceKernel();
/**
......@@ -86,10 +72,6 @@ public:
* @param force the HippoNonbondedForce this kernel will be used for
*/
void initialize(const System& system, const HippoNonbondedForce& force);
/**
* Compute the FFT.
*/
void computeFFT(bool forward, bool dispersion);
/**
* Get whether charge spreading should be done in fixed point.
*/
......@@ -113,8 +95,6 @@ private:
};
HipContext& cu;
HipSort* sort;
HipFFT3D* fft;
HipFFT3D* dfft;
};
} // namespace OpenMM
......
plugins/amoeba/platforms/opencl/src/AmoebaOpenCLKernels.cpp
View file @ 01e99e77
......@@ -29,30 +29,6 @@
using namespace OpenMM;
using namespace std;
/* -------------------------------------------------------------------------- *
* AmoebaMultipole *
* -------------------------------------------------------------------------- */
OpenCLCalcAmoebaMultipoleForceKernel::~OpenCLCalcAmoebaMultipoleForceKernel() {
if (fft != NULL)
delete fft;
}
void OpenCLCalcAmoebaMultipoleForceKernel::initialize(const System& system, const AmoebaMultipoleForce& force) {
CommonCalcAmoebaMultipoleForceKernel::initialize(system, force);
if (usePME)
fft = new OpenCLFFT3D(dynamic_cast<OpenCLContext&>(cc), gridSizeX, gridSizeY, gridSizeZ, false);
}
void OpenCLCalcAmoebaMultipoleForceKernel::computeFFT(bool forward) {
OpenCLArray& grid1 = dynamic_cast<OpenCLContext&>(cc).unwrap(pmeGrid1);
OpenCLArray& grid2 = dynamic_cast<OpenCLContext&>(cc).unwrap(pmeGrid2);
if (forward)
fft->execFFT(grid1, grid2, true);
else
fft->execFFT(grid2, grid1, false);
}
/* -------------------------------------------------------------------------- *
* HippoNonbondedForce *
* -------------------------------------------------------------------------- */
......@@ -60,10 +36,6 @@ void OpenCLCalcAmoebaMultipoleForceKernel::computeFFT(bool forward) {
OpenCLCalcHippoNonbondedForceKernel::~OpenCLCalcHippoNonbondedForceKernel() {
if (sort != NULL)
delete sort;
if (hasInitializedFFT) {
delete fftForward;
delete dfftForward;
}
}
void OpenCLCalcHippoNonbondedForceKernel::initialize(const System& system, const HippoNonbondedForce& force) {
......@@ -71,22 +43,9 @@ void OpenCLCalcHippoNonbondedForceKernel::initialize(const System& system, const
if (usePME) {
OpenCLContext& cl = dynamic_cast<OpenCLContext&>(cc);
sort = new OpenCLSort(cl, new SortTrait(), cc.getNumAtoms());
fftForward = new OpenCLFFT3D(cl, gridSizeX, gridSizeY, gridSizeZ, true);
dfftForward = new OpenCLFFT3D(cl, dispersionGridSizeX, dispersionGridSizeY, dispersionGridSizeZ, true);
hasInitializedFFT = true;
}
}
void OpenCLCalcHippoNonbondedForceKernel::computeFFT(bool forward, bool dispersion) {
OpenCLArray& grid1 = dynamic_cast<OpenCLContext&>(cc).unwrap(pmeGrid1);
OpenCLArray& grid2 = dynamic_cast<OpenCLContext&>(cc).unwrap(pmeGrid2);
OpenCLFFT3D* fft = (dispersion ? dfftForward : fftForward);
if (forward)
fft->execFFT(grid1, grid2, true);
else
fft->execFFT(grid2, grid1, false);
}
void OpenCLCalcHippoNonbondedForceKernel::sortGridIndex() {
sort->sort(dynamic_cast<OpenCLContext&>(cc).unwrap(pmeAtomGridIndex));
}
plugins/amoeba/platforms/opencl/src/AmoebaOpenCLKernels.h
View file @ 01e99e77
......@@ -9,7 +9,7 @@
* Biological Structures at Stanford, funded under the NIH Roadmap for *
* Medical Research, grant U54 GM072970. See https://simtk.org. *
* *
* Portions copyright (c) 2008-2021 Stanford University and the Authors. *
* Portions copyright (c) 2008-2025 Stanford University and the Authors. *
* Authors: Mark Friedrichs, Peter Eastman *
* Contributors: *
* *
......@@ -32,7 +32,6 @@
#include "openmm/System.h"
#include "AmoebaCommonKernels.h"
#include "OpenCLContext.h"
#include "OpenCLFFT3D.h"
#include "OpenCLSort.h"
namespace OpenMM {
......@@ -43,28 +42,14 @@ namespace OpenMM {
class OpenCLCalcAmoebaMultipoleForceKernel : public CommonCalcAmoebaMultipoleForceKernel {
public:
OpenCLCalcAmoebaMultipoleForceKernel(const std::string& name, const Platform& platform, OpenCLContext& cl, const System& system) :
CommonCalcAmoebaMultipoleForceKernel(name, platform, cl, system), fft(NULL) {
CommonCalcAmoebaMultipoleForceKernel(name, platform, cl, system) {
}
~OpenCLCalcAmoebaMultipoleForceKernel();
/**
* Initialize the kernel.
*
* @param system the System this kernel will be applied to
* @param force the AmoebaMultipoleForce this kernel will be used for
*/
void initialize(const System& system, const AmoebaMultipoleForce& force);
/**
* Compute the FFT.
*/
void computeFFT(bool forward);
/**
* Get whether charge spreading should be done in fixed point.
*/
bool useFixedPointChargeSpreading() const {
return true;
}
private:
OpenCLFFT3D* fft;
};
......@@ -74,7 +59,7 @@ private:
class OpenCLCalcHippoNonbondedForceKernel : public CommonCalcHippoNonbondedForceKernel {
public:
OpenCLCalcHippoNonbondedForceKernel(const std::string& name, const Platform& platform, OpenCLContext& cl, const System& system) :
CommonCalcHippoNonbondedForceKernel(name, platform, cl, system), sort(NULL), hasInitializedFFT(false) {
CommonCalcHippoNonbondedForceKernel(name, platform, cl, system), sort(NULL) {
}
~OpenCLCalcHippoNonbondedForceKernel();
/**
......@@ -84,10 +69,6 @@ public:
* @param force the HippoNonbondedForce this kernel will be used for
*/
void initialize(const System& system, const HippoNonbondedForce& force);
/**
* Compute the FFT.
*/
void computeFFT(bool forward, bool dispersion);
/**
* Get whether charge spreading should be done in fixed point.
*/
......@@ -109,9 +90,7 @@ private:
const char* getMaxValue() const {return "make_int2(2147483647, 2147483647)";}
const char* getSortKey() const {return "value.y";}
};
bool hasInitializedFFT;
OpenCLSort* sort;
OpenCLFFT3D *fftForward, *fftBackward, *dfftForward, *dfftBackward;
};
} // namespace OpenMM
......
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