CollectiveDeviceBuffer.h

Go to the documentation of this file.

#ifndef COLLECTIVE_DEVICE_BUFFER_H
#define COLLECTIVE_DEVICE_BUFFER_H

#include <vector>

#include "SynchronousCollectives.h"

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

// The type of collective buffers
enum class CollectiveBufferType {
  // Empty. Used to indicate a buffer has not been allocated or assigned a type.
  Empty,
  // Single-process P2P load stores.
  SingleProcess,
  // Multi-process P2P load stores via IPC.
  IPC,
  // NVSHMEM Buffer.
  Nvshmem,
  // NCCL Registered User Buffer.
  Nccl
};

template<typename T>
class CollectiveDeviceBuffer {
public:

  void allocate(CollectiveBufferType type_in, const size_t numElemsIn, 
    SynchronousCollectiveScope scope = SynchronousCollectiveScope::all);

  void allocate_no_check(CollectiveBufferType type_in, const size_t numElemsIn);

  void reallocate(CollectiveBufferType type_in, const size_t newNumElems, const double factor,
    SynchronousCollectiveScope scope = SynchronousCollectiveScope::all);

  /*
   * @brief Reallocates a symmetric buffer on all devices without reducing buffer sizes
   */
  void reallocate_no_check(CollectiveBufferType type_in, const size_t newNumElems, const double factor);

  /*
   * @brief Deallocates buffer on all devices
   */
  void deallocate();

  T* const  getDevicePtr() const {
    return buffer;
  }

  T** const getDevicePeerPtr() const {
    return d_peerBuffers;
  }

  std::vector<T*> const& getHostPeer() const {
    return h_peerBuffers;
  }

  // Move operator
  CollectiveDeviceBuffer<T>& operator=(CollectiveDeviceBuffer<T>&& other) {
    type = other.type;

    numElemsAlloc = other.numElemsAlloc;
    buffer = other.buffer;
    d_peerBuffers = other.d_peerBuffers;

    // This object owns the underlying data now, so we modify other
    other.numElemsAlloc = 0;
    other.buffer = nullptr;
    other.d_peerBuffers = nullptr;

    h_peerBuffers = std::move(other.h_peerBuffers);

    return *this;
  }

  // Copy operator
  CollectiveDeviceBuffer<T>& operator=(const CollectiveDeviceBuffer<T>& other) {
    type = other.type;
    numElemsAlloc = other.numElemsAlloc;
    buffer = other.buffer;
    d_peerBuffers = other.d_peerBuffers;
    h_peerBuffers = other.h_peerBuffers;

    return *this;
  }

  // Copy Constructor
  CollectiveDeviceBuffer(const CollectiveDeviceBuffer& other) {
    type = other.type;
    numElemsAlloc = other.numElemsAlloc;
    buffer = other.buffer;
    d_peerBuffers = other.d_peerBuffers;
    h_peerBuffers = other.h_peerBuffers;
  }
  
  // Default constructor
  CollectiveDeviceBuffer() {}

private:
  // Device pointer to buffer
  T* buffer = nullptr;
  // Pointer to peer's device buffers on device
  T** d_peerBuffers = nullptr;
  // Vector with peer's device buffers on host
  std::vector<T*> h_peerBuffers;
  // Number of elements allocated
  size_t numElemsAlloc = 0;
  // Type of buffer
  CollectiveBufferType type = CollectiveBufferType::Empty;
};

#endif  /* NAMD_CUDA || NAMD_HIP */

#endif  // COLLECTIVE_DEVICE_BUFFER_H