Skip to content

GitLab

Commit c26187aa authored by one's avatar one
Browse files

Revert "Add wave64 LDS spreading in HIP LJ-PME" 

This reverts commit 4e7070c2.

The split-pme-stream optimization to be implemented will put
LJ-PME reciprocal kernels on a non-critical path.
parent 05ad28f2
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 153 deletions
platforms/common/include/openmm/common/CommonCalcNonbondedForce.h
View file @ c26187aa
......@@ -46,9 +46,7 @@ class CommonCalcNonbondedForceKernel : public CalcNonbondedForceKernel {
public:
CommonCalcNonbondedForceKernel(std::string name, const Platform& platform, ComputeContext& cc, const System& system) : CalcNonbondedForceKernel(name, platform),
hasInitializedKernel(false), cc(cc), pmeio(NULL), stepsToSort(0), dispersionStepsToSort(0),
usePmeDispersionWave64LdsSpread(false), pmeDispersionSpreadWaveSize(0), pmeDispersionSpreadBlockSize(0),
pmeDispersionAtomsPerWave(0), pmeDispersionAtomsPerBlock(0), pmeGridIndexBlockSize(-1),
pmeSpreadChargeBlockSize(-1), pmeFinishSpreadChargeBlockSize(-1) {
pmeGridIndexBlockSize(-1), pmeSpreadChargeBlockSize(-1), pmeFinishSpreadChargeBlockSize(-1) {
}
~CommonCalcNonbondedForceKernel();
/**
......@@ -173,8 +171,6 @@ private:
int stepsToSort;
int dispersionStepsToSort;
bool usePmeQueue, deviceIsCpu, useFixedPointChargeSpreading, useCpuPme;
bool usePmeDispersionWave64LdsSpread;
int pmeDispersionSpreadWaveSize, pmeDispersionSpreadBlockSize, pmeDispersionAtomsPerWave, pmeDispersionAtomsPerBlock;
int pmeGridIndexBlockSize, pmeSpreadChargeBlockSize, pmeFinishSpreadChargeBlockSize;
bool hasCoulomb, hasLJ, doLJPME, usePosqCharges, recomputeParams, hasOffsets;
NonbondedMethod nonbondedMethod;
......
platforms/common/src/CommonCalcNonbondedForce.cpp
View file @ c26187aa
......@@ -290,15 +290,6 @@ void CommonCalcNonbondedForceKernel::commonInitialize(const System& system, cons
pmeGridIndexBlockSize = useLargeHipPmeBlocks ? 128 : -1;
pmeSpreadChargeBlockSize = useLargeHipPmeBlocks ? 128 : -1;
pmeFinishSpreadChargeBlockSize = isHip ? 128 : -1;
pmeDispersionSpreadWaveSize = 64;
pmeDispersionSpreadBlockSize = 256;
pmeDispersionAtomsPerWave = pmeDispersionSpreadWaveSize/PmeOrder;
pmeDispersionAtomsPerBlock = (pmeDispersionSpreadBlockSize/pmeDispersionSpreadWaveSize)*pmeDispersionAtomsPerWave;
// The LDS spread path assumes wave64 execution and is only used for HIP LJ-PME fixed point spreading.
usePmeDispersionWave64LdsSpread = (isHip &&
doLJPME && useFixedPointChargeSpreading && PmeOrder == 5 &&
cc.getSIMDWidth() == pmeDispersionSpreadWaveSize &&
cc.getMaxThreadBlockSize() >= pmeDispersionSpreadBlockSize);
map<string, string> defines;
defines["HAS_COULOMB"] = (hasCoulomb ? "1" : "0");
defines["HAS_LENNARD_JONES"] = (hasLJ ? "1" : "0");
......@@ -845,16 +836,9 @@ double CommonCalcNonbondedForceKernel::execute(ContextImpl& context, bool includ
pmeDefines["RECIP_EXP_FACTOR"] = cc.doubleToString(M_PI*M_PI/(dispersionAlpha*dispersionAlpha));
pmeDefines["USE_LJPME"] = "1";
pmeDefines["CHARGE_FROM_SIGEPS"] = "1";
if (usePmeDispersionWave64LdsSpread) {
pmeDefines["PME_USE_WAVE64_LDS_SPREAD"] = "1";
pmeDefines["PME_SPREAD_WAVE_SIZE"] = cc.intToString(pmeDispersionSpreadWaveSize);
pmeDefines["PME_SPREAD_BLOCK_SIZE"] = cc.intToString(pmeDispersionSpreadBlockSize);
pmeDefines["PME_SPREAD_ATOMS_PER_WAVE"] = cc.intToString(pmeDispersionAtomsPerWave);
pmeDefines["PME_SPREAD_ATOMS_PER_BLOCK"] = cc.intToString(pmeDispersionAtomsPerBlock);
}
program = cc.compileProgram(CommonKernelSources::pme, pmeDefines);
pmeDispersionGridIndexKernel = program->createKernel("findAtomGridIndex");
pmeDispersionSpreadChargeKernel = program->createKernel(usePmeDispersionWave64LdsSpread ? "gridSpreadChargeWave64Lds" : "gridSpreadCharge");
pmeDispersionSpreadChargeKernel = program->createKernel("gridSpreadCharge");
pmeDispersionConvolutionKernel = program->createKernel("reciprocalConvolution");
pmeDispersionEvalEnergyKernel = program->createKernel("gridEvaluateEnergy");
pmeDispersionInterpolateForceKernel = program->createKernel("gridInterpolateForce");
......@@ -1089,11 +1073,6 @@ double CommonCalcNonbondedForceKernel::execute(ContextImpl& context, bool includ
pmeDispersionSpreadChargeKernel->setArg(8, recipBoxVectorsFloat[1]);
pmeDispersionSpreadChargeKernel->setArg(9, recipBoxVectorsFloat[2]);
}
if (usePmeDispersionWave64LdsSpread) {
const int workSize = ((cc.getNumAtoms()+pmeDispersionAtomsPerBlock-1)/pmeDispersionAtomsPerBlock)*pmeDispersionSpreadBlockSize;
pmeDispersionSpreadChargeKernel->execute(workSize, pmeDispersionSpreadBlockSize);
}
else
pmeDispersionSpreadChargeKernel->execute(cc.getNumAtoms(), pmeSpreadChargeBlockSize);
if (useFixedPointChargeSpreading)
pmeDispersionFinishSpreadChargeKernel->execute(dispersionGridSizeX*dispersionGridSizeY*dispersionGridSizeZ, pmeFinishSpreadChargeBlockSize);
......
platforms/common/src/kernels/pme.cc
View file @ c26187aa
......@@ -120,131 +120,6 @@ KERNEL void gridSpreadCharge(GLOBAL const real4* RESTRICT posq,
}
}
#if defined(USE_HIP) && defined(CHARGE_FROM_SIGEPS) && PME_ORDER == 5 && defined(PME_USE_WAVE64_LDS_SPREAD)
// HIP LJ-PME dispersion spread variant for wave64 hardware. It keeps the generic
// PME spread algorithm unchanged, but maps one atom across PME_ORDER wavefront
// lanes so the five z-slices can be accumulated in parallel after one lane
// computes the atom's charge, grid index, and B-spline coefficients.
KERNEL void gridSpreadChargeWave64Lds(GLOBAL const real4* RESTRICT posq,
#ifdef USE_FIXED_POINT_CHARGE_SPREADING
GLOBAL mm_ulong* RESTRICT pmeGrid,
#else
GLOBAL real* RESTRICT pmeGrid,
#endif
real4 periodicBoxSize, real4 invPeriodicBoxSize, real4 periodicBoxVecX, real4 periodicBoxVecY, real4 periodicBoxVecZ,
real4 recipBoxVecX, real4 recipBoxVecY, real4 recipBoxVecZ, GLOBAL const int2* RESTRICT pmeAtomGridIndex,
GLOBAL const float2* RESTRICT sigmaEpsilon
) {
real3 data[PME_ORDER];
const real scale = RECIP((real) (PME_ORDER-1));
// B-spline coefficients are stored as [localAtom][order] so all wavefront
// lanes for the same atom read contiguous LDS entries during the spread phase.
LOCAL real sharedDataX[PME_SPREAD_BLOCK_SIZE];
LOCAL real sharedDataY[PME_SPREAD_BLOCK_SIZE];
LOCAL real sharedDataZ[PME_SPREAD_BLOCK_SIZE];
LOCAL real sharedCharge[PME_SPREAD_ATOMS_PER_BLOCK];
LOCAL int sharedGridX[PME_SPREAD_ATOMS_PER_BLOCK];
LOCAL int sharedGridY[PME_SPREAD_ATOMS_PER_BLOCK];
LOCAL int sharedGridZ[PME_SPREAD_ATOMS_PER_BLOCK];
// Each wavefront handles floor(PME_SPREAD_WAVE_SIZE/PME_ORDER) atoms. For
// PME_ORDER=5, lanes 0..4 spread one atom, lanes 5..9 spread the next, and
// the remaining lanes in the wavefront are inactive.
const int waveLane = LOCAL_ID & (PME_SPREAD_WAVE_SIZE-1);
const int wave = LOCAL_ID/PME_SPREAD_WAVE_SIZE;
int atomLane = waveLane-(waveLane/PME_ORDER)*PME_ORDER;
int atomInWave = waveLane/PME_ORDER;
int localAtom = wave*PME_SPREAD_ATOMS_PER_WAVE+atomInWave;
const int atomBlockStride = (GLOBAL_SIZE/PME_SPREAD_BLOCK_SIZE)*PME_SPREAD_ATOMS_PER_BLOCK;
for (int atomBlockStart = GROUP_ID*PME_SPREAD_ATOMS_PER_BLOCK; atomBlockStart < NUM_ATOMS; atomBlockStart += atomBlockStride) {
int atomIndex = atomBlockStart+localAtom;
bool isActive = (waveLane < PME_SPREAD_ATOMS_PER_WAVE*PME_ORDER && atomIndex < NUM_ATOMS);
if (isActive && atomLane == 0) {
int atom = pmeAtomGridIndex[atomIndex].x;
real4 pos = posq[atom];
const float2 sigEps = sigmaEpsilon[atom];
const real charge = 8*sigEps.x*sigEps.x*sigEps.x*sigEps.y;
APPLY_PERIODIC_TO_POS(pos)
real3 t = make_real3(pos.x*recipBoxVecX.x+pos.y*recipBoxVecY.x+pos.z*recipBoxVecZ.x,
pos.y*recipBoxVecY.y+pos.z*recipBoxVecZ.y,
pos.z*recipBoxVecZ.z);
t.x = (t.x-floor(t.x))*GRID_SIZE_X;
t.y = (t.y-floor(t.y))*GRID_SIZE_Y;
t.z = (t.z-floor(t.z))*GRID_SIZE_Z;
int3 gridIndex = make_int3(((int) t.x) % GRID_SIZE_X,
((int) t.y) % GRID_SIZE_Y,
((int) t.z) % GRID_SIZE_Z);
sharedCharge[localAtom] = charge;
sharedGridX[localAtom] = gridIndex.x;
sharedGridY[localAtom] = gridIndex.y;
sharedGridZ[localAtom] = gridIndex.z;
real3 dr = make_real3(t.x-(int) t.x, t.y-(int) t.y, t.z-(int) t.z);
data[PME_ORDER-1] = make_real3(0);
data[1] = dr;
data[0] = make_real3(1)-dr;
for (int j = 3; j < PME_ORDER; j++) {
real div = RECIP((real) (j-1));
data[j-1] = div*dr*data[j-2];
for (int k = 1; k < (j-1); k++)
data[j-k-1] = div*((dr+make_real3(k))*data[j-k-2] + (make_real3(j-k)-dr)*data[j-k-1]);
data[0] = div*(make_real3(1)-dr)*data[0];
}
data[PME_ORDER-1] = scale*dr*data[PME_ORDER-2];
for (int j = 1; j < (PME_ORDER-1); j++)
data[PME_ORDER-j-1] = scale*((dr+make_real3(j))*data[PME_ORDER-j-2] + (make_real3(PME_ORDER-j)-dr)*data[PME_ORDER-j-1]);
data[0] = scale*(make_real3(1)-dr)*data[0];
for (int j = 0; j < PME_ORDER; j++) {
int offset = localAtom*PME_ORDER+j;
sharedDataX[offset] = data[j].x;
sharedDataY[offset] = data[j].y;
sharedDataZ[offset] = data[j].z;
}
}
// All consumers of an atom's LDS data are in the same wavefront as the
// producing lane, so a wave-level barrier is sufficient here.
SYNC_WARPS
if (isActive) {
real charge = sharedCharge[localAtom];
if (charge != 0) {
int gridX = sharedGridX[localAtom];
int gridY = sharedGridY[localAtom];
int gridZ = sharedGridZ[localAtom];
int zindex = gridZ+atomLane;
zindex -= (zindex >= GRID_SIZE_Z ? GRID_SIZE_Z : 0);
real dz = sharedDataZ[localAtom*PME_ORDER+atomLane]*charge;
for (int ix = 0; ix < PME_ORDER; ix++) {
int xbase = gridX+ix;
xbase -= (xbase >= GRID_SIZE_X ? GRID_SIZE_X : 0);
xbase = xbase*GRID_SIZE_Y*GRID_SIZE_Z;
real dzdx = dz*sharedDataX[localAtom*PME_ORDER+ix];
for (int iy = 0; iy < PME_ORDER; iy++) {
int ybase = gridY+iy;
ybase -= (ybase >= GRID_SIZE_Y ? GRID_SIZE_Y : 0);
ybase = ybase*GRID_SIZE_Z;
int index = xbase + ybase + zindex;
real add = dzdx*sharedDataY[localAtom*PME_ORDER+iy];
if (fabs(add) > 2.3e-10f) { // Smallest value representable in 64 bit fixed point
#ifdef USE_FIXED_POINT_CHARGE_SPREADING
ATOMIC_ADD(&pmeGrid[index], (mm_ulong) realToFixedPoint(add));
#if defined(__GFX12__) && defined(USE_HIP)
// Workaround for rare cases when few values of pmeGrid are very large and
// incorrect. The cause is unknown. Why this workaround or other irrelevant
// changes like printf help is also unknown.
asm volatile("s_wait_storecnt 0x0");
#endif
#else
ATOMIC_ADD(&pmeGrid[index], add);
#endif
}
}
}
}
}
SYNC_WARPS
}
}
#endif
#ifdef USE_FIXED_POINT_CHARGE_SPREADING
KERNEL void finishSpreadCharge(
......
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