CudaGlobalMasterClient.C

Go to the documentation of this file.

#include "CudaGlobalMasterClient.h"
#include "CudaGlobalMasterClientKernel.h"
#include "CudaGlobalMasterServer.h"
#include "CudaUtils.h"
#include "InfoStream.h"
#include "Lattice.h"
#include "Molecule.h"
#include "ReductionMgr.h"
#include "Node.h"
#include "NamdState.h"
#include "DataExchanger.h"
#include "SimParameters.h"
#include "ScriptTcl.h"

namespace CudaGlobalMaster {

#if (defined(NAMD_CUDA) || defined(NAMD_HIP)) && defined(NODEGROUP_FORCE_REGISTER)

void CudaGlobalMasterClient::initialize(
  const std::vector<std::string>& arguments, int deviceID, cudaStream_t stream) {
  m_name = arguments[1];
  m_device_id = deviceID;
  const std::string msg = "Creating CudaGlobalMasterClient " + m_name +
                          " on PE " + std::to_string(CkMyPe()) + "\n";
  iout << iINFO << msg << endi;
  int savedDevice;
  cudaCheck(cudaGetDevice(&savedDevice));
  cudaCheck(cudaSetDevice(m_device_id));
  allocate_device(&m_d_ExtVirial, 1);
  allocate_device(&m_d_ExtForce, 1);
  allocate_host(&m_h_virial, 1);
  allocate_host(&m_h_ExtForce, 1);
  cudaCheck(cudaSetDevice(savedDevice));
}

CudaGlobalMasterClient::~CudaGlobalMasterClient() {
  if (m_d_ExtVirial)
    deallocate_device(&m_d_ExtVirial);
  if (m_d_ExtForce)
    deallocate_device(&m_d_ExtForce);
  if (m_h_virial)
    deallocate_host(&m_h_virial);
  if (m_h_ExtForce)
    deallocate_host(&m_h_ExtForce);
}

void CudaGlobalMasterClient::unsubscribe() {
  auto master = m_master.lock();
  if (master) {
    master->removeClient(shared_from_this());
    master.reset();
  } else {
    iout << iWARN << "Trying to remove client \"" << name()
         << "\" which is not in the master";
  }
}

void CudaGlobalMasterClient::subscribe(
    std::shared_ptr<CudaGlobalMasterServer> master) {
  auto tmp = m_master.lock();
  if (tmp) {
    return;
  } else {
    m_master = master;
    master->addClient(shared_from_this());
  }
}

int CudaGlobalMasterClient::updateFromTCLCommand(const std::vector<std::string>& arguments) {
  iout << iINFO << "Updating client " << name() << " with arguments:\n" << endi;
  for (size_t i = 0; i < arguments.size(); ++i) {
    iout << iINFO << "  arguments[" << i << "]: " << arguments[i] << '\n' << endi;
  }
#ifdef NAMD_TCL
  return TCL_OK;
#else
  return 0;
#endif // NAMD_TCL
}

std::string CudaGlobalMasterClient::getTCLUpdateResult() {
  return "undefined";
}

bool CudaGlobalMasterClient::tclAvailable() const {
#ifdef NAMD_TCL
  return true;
#else
  return false;
#endif // NAMD_TCL
}

void CudaGlobalMasterClient::finishReductions(bool doEnergy, bool doVirial,
                                              SubmitReduction *reduction) {
  if (doEnergy && hasEnergy()) {
    reduction->item(REDUCTION_MISC_ENERGY) += getEnergy();
  }
  if (useDefaultExtForceAndVirial()) {
    if (getRequestedAtoms().size() == getRequestedForcedAtoms().size()) {
      if (doVirial) {
        int savedDevice;
        cudaCheck(cudaGetDevice(&savedDevice));
        cudaCheck(cudaSetDevice(m_device_id));
        cudaCheck(cudaMemsetAsync(m_d_ExtVirial, 0, sizeof(cudaTensor), getStream()));
        cudaCheck(cudaMemsetAsync(m_d_ExtForce, 0, sizeof(Vector), getStream()));
        clientVirialAndExtForce(this->getPositions(), this->getAppliedForces(),
                                this->getRequestedAtoms().size(), m_d_ExtVirial,
                                m_d_ExtForce, getStream());
        copy_DtoH(m_d_ExtVirial, m_h_virial, 1, getStream());
        copy_DtoH(m_d_ExtForce, m_h_ExtForce, 1, getStream());
        cudaCheck(cudaStreamSynchronize(getStream()));
        ADD_TENSOR_OBJECT(reduction, REDUCTION_VIRIAL_NORMAL, (*m_h_virial));
        ADD_VECTOR_OBJECT(reduction, REDUCTION_EXT_FORCE_NORMAL, (*m_h_ExtForce));
        cudaCheck(cudaSetDevice(savedDevice));
      }
    } else {
      if (doVirial) {
        iout << iWARN << "Virial is not available for CudaGlobalMasterClient \""
            << name() << "\" since the number of atoms requested ("
            << std::to_string(getRequestedAtoms().size())
            << ") does not match the number of applied forces ("
            << std::to_string(getRequestedForcedAtoms().size()) << ")\n";
      }
    }
  } else {
    if (hasVirial()) {
      ADD_TENSOR_OBJECT(reduction, REDUCTION_VIRIAL_NORMAL, getVirial());
    }
    if (hasExtForce()) {
      ADD_VECTOR_OBJECT(reduction, REDUCTION_EXT_FORCE_NORMAL, getExtForce());
    }
  }
}

SimParameters* CudaGlobalMasterClient::getSimParameters() const {
  return Node::Object()->simParameters;
}

Molecule* CudaGlobalMasterClient::getMolecule() const {

  return Node::Object()->molecule;
}

const Controller* CudaGlobalMasterClient::getController() const {
  return &(Node::Object()->state->getController());
}

ScriptTcl* CudaGlobalMasterClient::getScript() {
#ifdef NAMD_TCL
  return Node::Object()->getScript();
#else
  return nullptr;
#endif // NAMD_TCL
}

bool CudaGlobalMasterClient::replica_enabled() const {
#if CMK_HAS_PARTITION
  return true;
#else
  return false;
#endif
}

int CudaGlobalMasterClient::replica_index() const {
  return CmiMyPartition();
}

int CudaGlobalMasterClient::num_replicas() const {
  return CmiNumPartitions();
}

void replica_comm_barrier() {
  replica_barrier();
}

int replica_comm_recv(char* msg_data, int buf_len, int src_rep) {
  DataMessage *recvMsg = NULL;
  replica_recv(&recvMsg, src_rep, CkMyPe());
  CmiAssert(recvMsg != NULL);
  int retval = recvMsg->size;
  if (buf_len >= retval) {
    memcpy(msg_data,recvMsg->data,retval);
  } else {
    retval = 0;
  }
  CmiFree(recvMsg);
  return retval;
}

int replica_comm_send(char* msg_data, int msg_len, int dest_rep) {
  replica_send(msg_data, msg_len, dest_rep, CkMyPe());
  return msg_len;
}

#else

CudaGlobalMasterClient::CudaGlobalMasterClient() {
  NAMD_die("CudaGlobalMasterClient requires NAMD to be built with CUDA");
}

#endif // (defined(NAMD_CUDA) || defined(NAMD_HIP)) && defined(NODEGROUP_FORCE_REGISTER)
}