ComputeGridForceCUDA.C

Go to the documentation of this file.

#include "ComputeGridForceCUDA.h"
#include "GridforceGridCUDAKernel.h"
#include "ComputeGridForceCUDAKernel.h"
#include "Molecule.h"
#include "Node.h"
#include "HomePatch.h"
#include "SimParameters.h"


#define GF_OVERLAPCHECK_FREQ 1000
#define MIN_DEBUG_LEVEL 4
//#define DEBUGM
#include "Debug.h"
#ifdef NODEGROUP_FORCE_REGISTER

ComputeGridForceCUDA::ComputeGridForceCUDA(
  std::vector<HomePatch*> &a_patchList,
  std::vector<AtomMap*> &atomMapsList,
  cudaStream_t stream)
{
  DebugM(4, "ComputeGridForceCUDA\n" << endi);
  this->stream=stream;
  this->patchList=a_patchList;
  Molecule *molecule       = Node::Object()->molecule;
  SimParameters *simParams = Node::Object()->simParameters;
  allocate_device<unsigned int>(&d_tbcatomic, 1);

  // create the list of atoms that need grid forces
  // allocate the grid structures we need on the device
  // extract the data we need from molecule
  // sync grid data from the host to the device

  numGriddedAtoms=createGriddedLists();

}

/*
 * update griddedAtomsSOAIndex for atoms that are local
 */
int ComputeGridForceCUDA::updateGriddedAtoms(
      std::vector<AtomMap*> &atomMapsList,
      std::vector<CudaLocalRecord> &localRecords,
      std::vector<HomePatch*> &patches,
      const int *h_globalToLocalID,
      bool mGpuOn)
{

  Molecule *molecule       = Node::Object()->molecule;
  DebugM(4, "ComputeGridForceCUDA::updateGriddedLists "<< numGriddedAtomsLocal <<"\n"<< endi);

  if(mGpuOn)
    { // we can't use the atomMapsList from SequencerCUDA in mGpuOn
      // go HomePatch by HomePatch with the vector from DeviceData
      for(int i = 0; i <patches.size(); ++i){
        HomePatch *p = patches[i];
        for(int j = 0; j < p->getNumAtoms(); j++){
          for (int gridnum = 0; gridnum < molecule->numGridforceGrids; gridnum++){
            int gid = p->getAtomList()[j].id; // Gets global ID of corresponding Ful}
            if (molecule->is_atom_gridforced(gid, gridnum)){
              LocalID lid;
              // the local id is a tuple of patch Id and the patch local index
              lid.pid=p->getPatchID();
              lid.index=j;
              int soaPid = h_globalToLocalID[lid.pid]; // Converts global patch ID to its local position in our SOA data structures
              std::pair<int,int> mapkey(gridnum,gid);
              int soaIndex = localRecords[soaPid].bufferOffset + lid.index;
              DebugM(1, "["<< CkMyPe() <<"] ComputeGridForceCUDA::updateGriddedLists gid "<< gid  << " patch "  << lid.pid << " atom " << j << " soaIndex[" << griddedAtomsGlobalIndexMap[mapkey] <<"] = "<< griddedAtomsSOAIndex[griddedAtomsGlobalIndexMap[mapkey]] <<"->" << soaIndex <<"\n"<< endi);
              griddedAtomsSOAIndex[griddedAtomsGlobalIndexMap[mapkey]] = soaIndex;
            }
          }
        }
      }
    }
  else
    { // use the atomMaplist on single gpu to save time
      for(int i = 0 ; i < numGriddedAtoms; i++){
        int gid = griddedAtomsGlobalIndex[i];
        LocalID lid;
        // Search for a valid localID in all atoms
        for(int j = 0 ; j < atomMapsList.size(); j++)
          {
            lid = atomMapsList[j]->localID(gid);
            if( lid.pid != -1) break;
          }
        //JM NOTE: Fields of lid need to be != -1, bc the atom needs to be somewhere
        //          otherwise we have a bug
        if(lid.pid == -1)
          {
            NAMD_bug(" LocalAtomID not found in patchMap");
          }
        int soaPid = h_globalToLocalID[lid.pid]; // Converts global patch ID to its local position in our SOA data structures
        int soaIndex = localRecords[soaPid].bufferOffset + lid.index;
        DebugM(1, "ComputeGridForceCUDA::updateGriddedLists soa[" << i <<"] <- "<< griddedAtomsSOAIndex[i] <<"->" << soaIndex <<"\n"<< endi);
        griddedAtomsSOAIndex[i] = soaIndex;
      }
    }
  // now that we've fixed the SOAIndex for gridded atoms, we need to
  // trim our local atom structures to the gridded atoms in the local
  // patch list.  If we knew exactly which ones left and which ones
  // arrived, this could be done more surgically, but this works.
  
  gridOffsetsLocal.clear();
  griddedAtomsLocalIndex.clear();
  gridOffsetsLocal.push_back(0); // we start at offset 0
  DebugM(3, "ComputeGridForceCUDA::updateGriddedLists grids "<< molecule->numGridforceGrids <<"\n" << endi);
  for (int gridnum = 0; gridnum < molecule->numGridforceGrids; gridnum++)
    {
      DebugM(2, "ComputeGridForceCUDA::updateGriddedLists patches "<< patches.size() <<"\n" << endi);
      for(int i = 0; i <patches.size(); ++i){
        HomePatch *p = patches[i];
        DebugM(1, "ComputeGridForceCUDA::updateGriddedLists patch " << p->getPatchID() <<" atoms "<< p->getNumAtoms() <<"\n" << endi);
        for(int j = 0; j < p->getNumAtoms(); j++){
          int gid = p->getAtomList()[j].id; // Gets global ID of corresponding Ful}
          if (molecule->is_atom_gridforced(gid, gridnum)){
            std::pair<int,int> mapkey(gridnum,gid);
            int mapoffset= griddedAtomsGlobalIndexMap[mapkey];
            if( mapoffset <0)
              {
                NAMD_bug("how is this not in our map?");
              }
            griddedAtomsLocalIndex.push_back(mapoffset);
            DebugM(1, "ComputeGridForceCUDA::updateGriddedLists patch " << p->getPatchID() <<" atom "<< j << " gid "<<gid<< " gridnum " << gridnum << " mapoffset "<< mapoffset << " localIndex.size() " << griddedAtomsLocalIndex.size()<< "\n" << endi);
          }
        }
        DebugM(2, "ComputeGridForceCUDA::updateGriddedLists patch " << p->getPatchID() <<" gridded atoms "<< griddedAtomsLocalIndex.size() <<"\n" << endi);
      }
      int size = griddedAtomsLocalIndex.size();
      gridOffsetsLocal.push_back(size);
    }
  
  numGriddedAtomsLocal = griddedAtomsLocalIndex.size();
  DebugM(4, "ComputeGridForceCUDA::updateGriddedLists numGriddedAtomsLocal now "<< numGriddedAtomsLocal <<"\n" << endi);
  // copy indices to device
  if(numGriddedAtomsLocal>0)
    {
      int* soaPtr=griddedAtomsSOAIndex.data();
      copy_HtoD_sync<int>(soaPtr, d_griddedAtomsSOAIndex, griddedAtomsSOAIndex.size());
      int* idxPtr=griddedAtomsLocalIndex.data();
      copy_HtoD_sync<int>(idxPtr, d_griddedAtomsLocalIndex, griddedAtomsLocalIndex.size());     
    }
  return numGriddedAtomsLocal;
}


/*
 * We make host side arrays that we can sync over for the atomIdx,
 * gridFrcParams.k, gridFrcParams.q for all the atoms that have
 * is_atom_gridforced == TRUE.  That way we don't skip around
 * non-contiguous references into the molecule for the handful of
 * atoms.

 * The SOA indexand LocalIndex get updated by migrations, and that is
 * what we'll iterate over.

 */


int ComputeGridForceCUDA::createGriddedLists()
{
  // get the total count across all grids so we can just stream
  // through the whole thing
  DebugM(4, "ComputeGridForceCUDA::createGriddedLists\n" << endi);
  Molecule *molecule       = Node::Object()->molecule;
  SimParameters *simParams = Node::Object()->simParameters;
  Atom *atoms = molecule->getAtoms();
  int gridTotalSize=0;
  numGriddedAtoms=0;
  numGriddedAtomsLocal=0;
  gridOffsets.clear();
  gridOffsets.push_back(0); // 0th grid has 0th offset
  gridOffsetsLocal.push_back(0); // 0th grid has 0th offset
  gridCoordsOffsets.push_back(0);
  griddedAtomsGlobalIndex.clear();
  griddedAtomsSOAIndex.clear();
  for (int gridnum = 0; gridnum < molecule->numGridforceGrids; gridnum++) {
    // atoms can be in multiple grids, so we go grid by grid
    int numAtoms = molecule->numAtoms;
    for(int j = 0; j < numAtoms; j++){
      if (molecule->is_atom_gridforced(j, gridnum)){
        griddedAtomsGlobalIndex.push_back(j);
        // we're going to have to look up the reverse of GID to index
        // a lot, so make a map
        std::pair<int,int> mapkey(gridnum,j);
        griddedAtomsGlobalIndexMap[mapkey]= griddedAtomsGlobalIndex.size();
      }
    }
    gridOffsets.push_back(griddedAtomsGlobalIndex.size());
    auto *mol_grid=molecule->get_gridfrc_grid(gridnum);
    GridforceFullBaseGrid* h_grid;
    if( mol_grid->get_grid_type() == 1)
      h_grid = (GridforceFullMainGrid *) mol_grid;
    else
      NAMD_bug("GridforceGridCUDA::createGriddedLists called with unsupported grid type!");
    int size =h_grid->size;
    gridCoordsOffsets.push_back(size);
    gridTotalSize += size;
    DebugM(2, "ComputeGridForceCUDA::createGriddedLists grid "<<gridnum <<" has " << size << " points" << " and " << gridOffsets[gridnum+1] - gridOffsets[gridnum] <<" atoms\n"<<endi);
  }

  for (int gridnum = 0; gridnum < molecule->numGridforceGrids; gridnum++) {
    // atoms can be in multiple grids, so we go grid by grid
    for(int i = 0; i < patchList.size(); i++){
      HomePatch *p = patchList[i];
      for(int j = 0; j < p->getNumAtoms(); j++){
        int gid = p->getAtomList()[j].id; // Gets global ID of corresponding FullAtom
        if (molecule->is_atom_gridforced(gid, gridnum)){
          std::pair<int,int> mapkey(gridnum,gid);
          int globalIndex=griddedAtomsGlobalIndexMap[mapkey];
          griddedAtomsLocalIndex.push_back(globalIndex);
        }
      }
    }
    gridOffsetsLocal.push_back(griddedAtomsLocalIndex.size());
  }
  
  numGriddedAtoms=griddedAtomsGlobalIndex.size();
  DebugM(3, "ComputeGridForceCUDA::createGriddedLists numGriddedAtoms " << numGriddedAtoms << " numAtoms "<< molecule->numAtoms << "grids " << molecule->numGridforceGrids << "\n" << endi);
  // size to fit all the gridded atoms
  griddedAtomsSOAIndex.resize(numGriddedAtoms);
  numGriddedAtomsLocal=0;
  if(numGriddedAtoms > 0)
    {
      allocate_host<float>(&h_gridded_charge, numGriddedAtoms);
      allocate_host<float>(&h_gridded_scale, numGriddedAtoms);
      allocate_device<float>(&d_gridded_charge, numGriddedAtoms);
      allocate_device<float>(&d_gridded_scale, numGriddedAtoms);
      allocate_device<GridforceGridCUDA>(&d_grids, molecule->numGridforceGrids);
      allocate_device<int>(&d_griddedAtomsSOAIndex, numGriddedAtoms);
      allocate_device<int>(&d_griddedAtomsLocalIndex, numGriddedAtoms);
      /*
       * we flatten it all into one array of gridded atoms for all the atoms
       */
      allocate_device<float>(&d_gridCoords, gridTotalSize);
      // for the multiple grid case, we need grid specific buffers for
      // these, because their kernel launches may run in parallel
      allocate_device<double>(&d_extEnergy_G, molecule->numGridforceGrids);
      allocate_device<double3>(&d_netForce_G, molecule->numGridforceGrids*3);
      allocate_device<cudaTensor>(&d_extVirial_G, molecule->numGridforceGrids);
      allocate_host<double>(&h_extEnergy_G, molecule->numGridforceGrids);
      allocate_host<double3>(&h_netForce_G, molecule->numGridforceGrids*3);
      allocate_host<cudaTensor>(&h_extVirial_G, molecule->numGridforceGrids);

      for (int gridnum = 0; gridnum < molecule->numGridforceGrids; gridnum++) {
        int offset = gridOffsets[gridnum];
        auto *mol_grid=molecule->get_gridfrc_grid(gridnum);
        GridforceFullBaseGrid* h_grid;
        if( mol_grid->get_grid_type() == 1)
          h_grid = (GridforceFullMainGrid *) mol_grid;
        else
          NAMD_bug("GridforceGridCUDA::createGriddedLists called with unsupported grid type!");
        int grid_size = h_grid->size;
        for(int gridIdx = offset; gridIdx < gridOffsets[gridnum+1];  gridIdx++)
          {
            int gid = griddedAtomsGlobalIndex[gridIdx];
            if (molecule->is_atom_gridforced(gid, gridnum)){
              molecule->get_gridfrc_params(h_gridded_scale[gridIdx], h_gridded_charge[gridIdx], gid, gridnum);

            }
          }

        h_grids.push_back(GridforceGridCUDA(h_grid->get_k0(),
                                            h_grid->get_k1(),
                                            h_grid->get_k2(),
                                            h_grid->size,
                                            h_grid->dk[0],
                                            h_grid->dk[1],
                                            h_grid->dk[2],
                                            h_grid->factor,
                                            h_grid->get_origin(),
                                            h_grid->get_center(),
                                            h_grid->cont[0],
                                            h_grid->cont[1],
                                            h_grid->cont[2],
                                            h_grid->gapinv[0],
                                            h_grid->gapinv[1],
                                            h_grid->gapinv[2],
                                            h_grid->gap[0],
                                            h_grid->gap[1],
                                            h_grid->gap[2],
                                            h_grid->inv,
                                            h_grid->offset[0],
                                            h_grid->offset[1],
                                            h_grid->offset[2],
                                            h_grid->get_scale(),
                                            &d_gridCoords[gridCoordsOffsets[gridnum]]));
        copy_HtoD_sync<float>(h_grid->grid, &d_gridCoords[gridCoordsOffsets[gridnum]], h_grid->size );

      }
      GridforceGridCUDA* gPtr=h_grids.data();
      copy_HtoD_sync<GridforceGridCUDA>(gPtr, d_grids, molecule->numGridforceGrids );
      copy_HtoD_sync<float>(h_gridded_charge, d_gridded_charge, numGriddedAtoms);
      copy_HtoD_sync<float>(h_gridded_scale, d_gridded_scale, numGriddedAtoms);
      cudaCheck(cudaMemset(d_tbcatomic, 0, sizeof(unsigned int))); // sets the scala
    }
  return numGriddedAtoms;
}



void ComputeGridForceCUDA::doForce(
                                   const int doEnergy,
                                   const int doVirial,
                                   const Lattice lat,
                                   const int timeStep,
                                   const double* d_pos_x,
                                   const double* d_pos_y,
                                   const double* d_pos_z,
                                   const char3* d_transform,
                                   double* f_normal_x,
                                   double* f_normal_y,
                                   double* f_normal_z,
                                   cudaStream_t stream)
{
  // iterate through the grids like nonCUDA
  DebugM(4, "ComputeGridForceCUDA::doForce\n" << endi);
  SimParameters *simParams = Node::Object()->simParameters;
  Molecule *mol = Node::Object()->molecule;
  if(numGriddedAtomsLocal>0)
    {
      HomePatch *homePatch = patchList[0];
      // EJB- CPU side code has this check, so blindly reimplemented here
      if (homePatch->flags.step % GF_OVERLAPCHECK_FREQ == 0) {
        for (int gridnuml = 0; gridnuml < mol->numGridforceGrids; gridnuml++) {
          // only check on node 0 and every GF_OVERLAPCHECK_FREQ steps
          if (simParams->langevinPistonOn || simParams->berendsenPressureOn) {
            // check for grid overlap if pressure control is on not needed
            // without pressure control, since the check is also performed
            // on startup
            GridforceGrid *grid = mol->get_gridfrc_grid(gridnuml);
            if (!grid->fits_lattice(homePatch->lattice)) {
              char errmsg[512];
              if (grid->get_checksize()) {
                sprintf(errmsg, "Warning: Periodic cell basis too small for Gridforce grid %d.  Set gridforcechecksize off in configuration file to ignore.\n", gridnuml);
                NAMD_die(errmsg);
              }
            }
          }
        }
      }
      // The atoms that need gridding, along with their static data
      // (charge, scaling, etc.) are already collated.  Launch computeGridForce
      // kernel with the grid, the collated arrays of charge and
      // scale, the atom position arrays, to compute accumulated
      // forces into the device side force arrays along with the
      // energy and virial contribution from each gridForced atom
      for (int gridnum = 0; gridnum < mol->numGridforceGrids; gridnum++) {
        DebugM(3, "ComputeGridForceCUDA::doForce grid "<< gridnum << " from local offset " << gridOffsetsLocal[gridnum] << " to " << gridOffsetsLocal[gridnum+1]<< "\n" << endi);
        h_extEnergy_G[gridnum] = 0.;
        h_netForce_G[gridnum].x = 0;
        h_netForce_G[gridnum].y = 0;
        h_netForce_G[gridnum].z = 0;
        h_extVirial_G[gridnum].xx = 0;
        h_extVirial_G[gridnum].xy = 0;
        h_extVirial_G[gridnum].xz = 0;
        h_extVirial_G[gridnum].yx = 0;
        h_extVirial_G[gridnum].yy = 0;
        h_extVirial_G[gridnum].yz = 0;
        h_extVirial_G[gridnum].zx = 0;
        h_extVirial_G[gridnum].zy = 0;
        h_extVirial_G[gridnum].zz = 0;
        computeGridForce(doEnergy, doVirial, d_grids[gridnum], lat,  d_pos_x, d_pos_y, d_pos_z, d_transform, d_griddedAtomsSOAIndex, &d_griddedAtomsLocalIndex[gridOffsetsLocal[gridnum]], d_gridded_charge, d_gridded_scale, f_normal_x, f_normal_y, f_normal_z, &h_netForce_G[gridnum], &d_netForce_G[gridnum], &h_extEnergy_G[gridnum], &d_extEnergy_G[gridnum], &h_extVirial_G[gridnum], &d_extVirial_G[gridnum], gridOffsetsLocal[gridnum+1]-gridOffsetsLocal[gridnum], d_tbcatomic, stream);
      }
    }
  else
    {
      //      CkPrintf("[%d] computeGridForce NO griddedAtoms %d\n",CkMyPe(),numGriddedAtoms);
    }
}
void ComputeGridForceCUDA::zeroOutEnergyVirialForcesAcrossGrids(
                                                                double* h_extEnergy, double3* h_netForce, cudaTensor* h_virial)
{
  Molecule *mol = Node::Object()->molecule;
  h_extEnergy[0] = 0;
  h_netForce->x = 0;
  h_netForce->y = 0;
  h_netForce->z = 0;
  h_virial->xx = 0;
  h_virial->xy = 0;
  h_virial->xz = 0;
  h_virial->yx = 0;
  h_virial->yy = 0;
  h_virial->yz = 0;
  h_virial->zx = 0;
  h_virial->zy = 0;
  h_virial->zz = 0;
  for (int gridnum = 0; gridnum < mol->numGridforceGrids; gridnum++) {
    h_extEnergy_G[gridnum] = 0.;
    h_netForce_G[gridnum].x = 0;
    h_netForce_G[gridnum].y = 0;
    h_netForce_G[gridnum].z = 0;
    h_extVirial_G->xx = 0;
    h_extVirial_G->xy = 0;
    h_extVirial_G->xz = 0;
    h_extVirial_G->yx = 0;
    h_extVirial_G->yy = 0;
    h_extVirial_G->yz = 0;
    h_extVirial_G->zx = 0;
    h_extVirial_G->zy = 0;
    h_extVirial_G->zz = 0;
  }
  
}

void ComputeGridForceCUDA::sumEnergyVirialForcesAcrossGrids(double* h_extEnergy,
                                         double3* h_netForce,
                                         cudaTensor* h_virial)
{
  Molecule *mol = Node::Object()->molecule;
  for (int gridnum = 0; gridnum < mol->numGridforceGrids; gridnum++) {
    DebugM(2, "ComputeGridForceCUDA::sumEnergyVirialForcesAcrossGrids in grid "<< gridnum <<" energy "<< h_extEnergy_G[gridnum] <<"\n" << endi);
    h_extEnergy[0] += h_extEnergy_G[gridnum];


    h_netForce->x += h_netForce_G[gridnum].x;
    h_netForce->y += h_netForce_G[gridnum].y;
    h_netForce->z += h_netForce_G[gridnum].z;
    h_virial[0] += h_extVirial_G[gridnum];
    DebugM(2, "ComputeGridForceCUDA::sumEnergyVirialForcesAcrossGrids out grid "<< gridnum <<" energy "<< h_extEnergy_G[gridnum] <<"\n" << endi);
  }
}
void ComputeGridForceCUDA::destroyGriddedLists()
{
  DebugM(4, "ComputeGridForceCUDA::destroyGriddedLists\n" << endi);
  if(numGriddedAtoms>0){
    deallocate_host<float>(&h_gridded_charge);
    deallocate_host<float>(&h_gridded_scale);
    deallocate_device<float>(&d_gridded_charge);
    deallocate_device<float>(&d_gridded_scale);
    deallocate_device<float>(&d_gridCoords);
    deallocate_device<int>(&d_griddedAtomsSOAIndex);
    deallocate_device<int>(&d_griddedAtomsLocalIndex);
    deallocate_device<GridforceGridCUDA>(&d_grids);
    deallocate_device<float>(&d_gridCoords);
    deallocate_device<double>(&d_extEnergy_G);
    deallocate_device<double3>(&d_netForce_G);
    deallocate_device<cudaTensor>(&d_extVirial_G);
    deallocate_host<double>(&h_extEnergy_G);
    deallocate_host<double3>(&h_netForce_G);
    deallocate_host<cudaTensor>(&h_extVirial_G);

  }
}

ComputeGridForceCUDA::~ComputeGridForceCUDA()
{
  DebugM(4, "ComputeGridForceCUDA::~ComputeGridForceCUDA\n" << endi);
  destroyGriddedLists();
  deallocate_device<unsigned int>(&d_tbcatomic);
}



#endif //NODEGROUP_FORCE_REGISTER