CollectiveDeviceBuffer.C

Go to the documentation of this file.

#include "GlobalGPUMgr.h"
#include "SynchronousCollectives.h"
#include "CollectiveDeviceBuffer.h"
#include "NamdEventsProfiling.h"
#include "TupleTypesCUDA.h"

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

template<typename T>
void CollectiveDeviceBuffer<T>::allocate(CollectiveBufferType type_in, const size_t numElemsIn, 
  SynchronousCollectiveScope scope) {
  SynchronousCollectives* syncColl = SynchronousCollectives::Object();
  GlobalGPUMgr* globalGPUMgr = GlobalGPUMgr::Object();
  const bool isMasterPe = globalGPUMgr->getIsMasterPe();

  size_t numElemsTemp;
  if (isMasterPe) { 
    numElemsTemp = numElemsIn;
  } else {
    numElemsTemp = 0;
  }

  // Compute size of buffer
  std::vector<size_t> numElemsVec = {numElemsTemp};
  const size_t maxNumElems = syncColl->allReduce<size_t>(numElemsVec, CkReduction::max_ulong_long,
    scope)[0];

  // Call the no check version of allocate
  allocate_no_check(type_in, maxNumElems);
}

template<typename T>
void CollectiveDeviceBuffer<T>::allocate_no_check(CollectiveBufferType type_in, const size_t numElemsIn) {
  SynchronousCollectives* syncColl = SynchronousCollectives::Object();
  GlobalGPUMgr* globalGPUMgr = GlobalGPUMgr::Object();
  const bool isMasterPe = globalGPUMgr->getIsMasterPe();
  const int numDevices = globalGPUMgr->getNumDevices();
  const int deviceIndex = globalGPUMgr->getDeviceIndex();

  type = type_in;

  numElemsAlloc = numElemsIn;
  
  if (CollectiveBufferType::SingleProcess == type) {
    if (isMasterPe) {
      // Allocate Buffer locally
      allocate_device<T>(&(buffer), numElemsAlloc);

      // Communicate buffer to peers
      auto temp_peerBuffers = syncColl->allGather<unsigned long long>(
          (unsigned long long) buffer, SynchronousCollectiveScope::master);

      // Copy peer buffers to vector
      h_peerBuffers.resize(numDevices);
      for (int i = 0; i < numDevices; i++) {
        h_peerBuffers[i] = (T*) temp_peerBuffers[i];
      }

      // Copy P2P buffer to device
      allocate_device<T*>(&(d_peerBuffers), numDevices);
      copy_HtoD<T*>(h_peerBuffers.data(), d_peerBuffers, numDevices, nullptr);
      cudaStreamSynchronize(nullptr);
    }
#if ! defined(NAMD_HIP)
  } else if (CollectiveBufferType::IPC == type) {
    if (isMasterPe) {
      cudaIpcMemHandle_t handle;
      // Allocate Buffer locally
      allocate_device<T>(&(buffer), numElemsAlloc);
      cudaCheck(cudaIpcGetMemHandle(&handle, (void*) buffer));

      // Communicate handles to peers
      auto temp_peerHandles = syncColl->allGather<cudaIpcMemHandle_t>(handle, 
        SynchronousCollectiveScope::master);

      // Open CUDA IPC mem handles and store in vector
      h_peerBuffers.resize(numDevices);
      for (int i = 0; i < numDevices; i++) {
        if (i != deviceIndex) {
          cudaCheck(cudaIpcOpenMemHandle((void**) &(h_peerBuffers[i]), 
            temp_peerHandles[i],
            cudaIpcMemLazyEnablePeerAccess));
        } else {
          h_peerBuffers[i] = buffer;
        }
      }

      // Copy P2P buffer to device
      allocate_device<T*>(&(d_peerBuffers), numDevices);
      copy_HtoD<T*>(h_peerBuffers.data(), d_peerBuffers, numDevices, nullptr);
      cudaCheck(cudaStreamSynchronize(nullptr));
    }
#endif
  } else {
    NAMD_die("CollectiveBufferType not currently implemented");
  }
}

template<typename T>
void CollectiveDeviceBuffer<T>::reallocate(CollectiveBufferType type_in, const size_t newNumElems, const double factor,
  SynchronousCollectiveScope scope) {
  SynchronousCollectives* syncColl = SynchronousCollectives::Object();
  GlobalGPUMgr* globalGPUMgr = GlobalGPUMgr::Object();
  const bool isMasterPe = globalGPUMgr->getIsMasterPe();

  size_t numElemRequested;
  if (isMasterPe) {
    numElemRequested = newNumElems;
  } else {
    numElemRequested = 0;
  }

  std::vector<size_t> numElemsVec = {numElemRequested};
  const size_t maxNumElemsRequested = syncColl->allReduce<size_t>(numElemsVec, CkReduction::max_ulong_long,
    scope)[0];

  reallocate_no_check(type_in, maxNumElemsRequested, factor);
}

template<typename T>
void CollectiveDeviceBuffer<T>::reallocate_no_check(CollectiveBufferType type_in, const size_t newNumElems, 
  const double factor) {

  GlobalGPUMgr* globalGPUMgr = GlobalGPUMgr::Object();
  const bool isMasterPe = globalGPUMgr->getIsMasterPe();
 
  if (isMasterPe) {
    if (type == CollectiveBufferType::Empty) {
      type = type_in;
    } else if (type != type_in) {
      NAMD_die("Reallocating buffer with different type");
    }

    const size_t maxNumElemsRequested = newNumElems;

    if (maxNumElemsRequested > numElemsAlloc) {
      const size_t newNumElemsAlloc = (size_t) ((double)maxNumElemsRequested * factor);
      deallocate();
      allocate_no_check(type_in, newNumElemsAlloc);
    }
  }
}

template<typename T>
void CollectiveDeviceBuffer<T>::deallocate() {
  // This is needed because different PEs can call functions on objects created by other PEs.
  // When this happens, both PEs could try to call deallocate leading to a double free and
  // a seg fault. This seems like a hacky way to fix
  GlobalGPUMgr* globalGPUMgr = GlobalGPUMgr::Object();

  numElemsAlloc = 0;

  if (!globalGPUMgr->getIsMasterPe()) return;

  if (buffer) {
    deallocate_device<T>(&buffer);
    buffer = nullptr;
  }
  if (d_peerBuffers) {
    deallocate_device<T*>(&d_peerBuffers);
    d_peerBuffers = nullptr;
  }
}

/*
 * Explicit instantiation
 */
template class CollectiveDeviceBuffer<char>;
template class CollectiveDeviceBuffer<uint64_t>;
template class CollectiveDeviceBuffer<int64_t>;
template class CollectiveDeviceBuffer<int>;
template class CollectiveDeviceBuffer<int4>;
template class CollectiveDeviceBuffer<float>;
template class CollectiveDeviceBuffer<float2>;
template class CollectiveDeviceBuffer<float4>;
template class CollectiveDeviceBuffer<double>;
template class CollectiveDeviceBuffer<double3>;
template class CollectiveDeviceBuffer<FullAtom>;
template class CollectiveDeviceBuffer<CudaLocalRecord>;
template class CollectiveDeviceBuffer<CudaForce>;
template class CollectiveDeviceBuffer<CudaBondStage>;
template class CollectiveDeviceBuffer<CudaAngleStage>;
template class CollectiveDeviceBuffer<CudaDihedralStage>;
template class CollectiveDeviceBuffer<CudaExclusionStage>;
template class CollectiveDeviceBuffer<CudaCrosstermStage>;

#endif  /* NAMD_CUDA || NAMD_HIP */