DeviceArray.h

#ifndef PSCF_DEVICE_ARRAY_H
#define PSCF_DEVICE_ARRAY_H
 
/*
* PSCF - Polymer Self-Consistent Field
*
* Copyright 2015 - 2025, The Regents of the University of Minnesota
* Distributed under the terms of the GNU General Public License.
*/
 
#include <util/misc/ReferenceCounter.h>   // member
#include <util/misc/CountedReference.h>   // member
 
// Forward declarations
namespace Util {
   template <typename Data> class DArray;
}
namespace Pscf {
   class DeviceMemory;
}
 
namespace Pscf {
 
   using namespace Util;

   template <typename Data>
   class DeviceArray
   {
 
   public:

      typedef Data ValueType;

      DeviceArray();

      DeviceArray(int capacity);

      DeviceArray(DeviceArray<Data> const & other);

      DeviceArray(DArray<Data> const & other);

      virtual ~DeviceArray();

      void allocate(int capacity);

      void deallocate();

      void associate(DeviceArray<Data>& arr, int beginId, int capacity);

      void associate(DeviceMemory& arr, int capacity);

      void dissociate();

      void addReference(CountedReference& reference);

      virtual
      DeviceArray<Data>& operator = (const DeviceArray<Data>& other);

      virtual
      DeviceArray<Data>& operator = (const DArray<Data>& other);

      int capacity() const;

      bool isAllocated() const;

      bool isOwner() const;

      bool isAssociated() const;

      Data* cArray();

      Data const * cArray() const;
 
   protected:

      Data* dataPtr_;

      int capacity_;

      ReferenceCounter refCounter_;

      CountedReference ref_;
 
   };
 
   // Inline member function definitions
 
   /*
   * Return array capacity.
   */
   template <typename Data> inline
   int DeviceArray<Data>::capacity() const
   {  return capacity_; }
 
   /*
   * Return true if this object has access to a memory block.
   */
   template <typename Data> inline
   bool DeviceArray<Data>::isAllocated() const
   {  return (bool) dataPtr_; }
 
   /*
   * Does this object own data?
   */
   template <typename Data> inline
   bool DeviceArray<Data>::isOwner() const
   {  return ((bool) dataPtr_ && !ref_.isAssociated()); }
 
   /*
   * Is this object associated with data it does not own?
   */
   template <typename Data> inline
   bool DeviceArray<Data>::isAssociated() const
   {  return ((bool) dataPtr_ && ref_.isAssociated()); }
 
}
 
#include "DeviceMemory.h"
#include "cudaErrorCheck.h"
#include <util/containers/DArray.h>
#include <util/global.h>
#include <cuda_runtime.h>
 
namespace Pscf {
 
   /*
   * Default constructor.
   */
   template <typename Data>
   DeviceArray<Data>::DeviceArray()
    : dataPtr_(nullptr),
      capacity_(0)
   {}
 
   /*
   * Allocating constructor.
   */
   template <typename Data>
   DeviceArray<Data>::DeviceArray(int capacity)
    : dataPtr_(nullptr),
      capacity_(0)
   {  allocate(capacity); }
 
   /*
   * Copy constructor.
   *
   * Allocates new memory and copies all elements by value.
   */
   template <typename Data>
   DeviceArray<Data>::DeviceArray(const DeviceArray<Data>& other)
    : dataPtr_(nullptr),
      capacity_(0)
   {
      if (!other.isAllocated()) {
         UTIL_THROW("Other array must be allocated.");
      }
 
      allocate(other.capacity_);
      cudaErrorCheck( cudaMemcpy(dataPtr_, other.cArray(),
                                 capacity_ * sizeof(Data),
                                 cudaMemcpyDeviceToDevice) );
   }
 
   /*
   * Copy constructor, deep copy from host DArray<Data>
   *
   * Allocates new memory and copies all elements by value.
   */
   template <typename Data>
   DeviceArray<Data>::DeviceArray(const DArray<Data>& other)
    : dataPtr_(nullptr),
      capacity_(0)
   {
      if (!other.isAllocated()) {
         UTIL_THROW("Other array must be allocated.");
      }
 
      allocate(other.capacity());
      cudaErrorCheck( cudaMemcpy(dataPtr_, other.cArray(),
                                 capacity_ * sizeof(Data),
                                 cudaMemcpyHostToDevice) );
   }
 
   /*
   * Destructor.
   */
   template <typename Data>
   DeviceArray<Data>::~DeviceArray()
   {
      if (isOwner()) {
         if (refCounter_.hasRefs()) {
            std::cout
                << "Error: Destroying DeviceArray that is referred"
                << "to by at least one CountedReference"
                << std::endl;
         }
         cudaFree(dataPtr_);
      }
      if (ref_.isAssociated()) {
         ref_.dissociate();
      }
   }
 
   /*
   * Allocate device memory, with this object then owns.
   */
   template <typename Data>
   void DeviceArray<Data>::allocate(int capacity)
   {
      if (capacity <= 0) {
         UTIL_THROW("Attempt to allocate with capacity <= 0");
      }
      if (isAllocated()) {
         UTIL_THROW("Attempt to re-allocate an array");
      }
      if (ref_.isAssociated()) {
         UTIL_THROW("Attempt to allocate array that uses external data");
      }
 
      cudaErrorCheck(cudaMalloc((void**) &dataPtr_,
                     capacity * sizeof(Data)));
      capacity_ = capacity;
   }
 
   /*
   * Deallocate device memory owned by this object, if any.
   */
   template <typename Data>
   void DeviceArray<Data>::deallocate()
   {
      UTIL_CHECK(dataPtr_);
      UTIL_CHECK(!ref_.isAssociated());
      UTIL_CHECK(!refCounter_.hasRefs());
      cudaErrorCheck( cudaFree(dataPtr_) );
      dataPtr_ = nullptr;
      capacity_ = 0;
   }
 
   /*
   * Associate this object with memory owned by a different DeviceArray.
   */
   template <typename Data>
   void DeviceArray<Data>::associate(DeviceArray<Data>& arr,
                                     int beginId, int capacity)
   {
      UTIL_CHECK(arr.isAllocated());
      UTIL_CHECK(arr.isOwner());
      UTIL_CHECK(beginId >= 0);
      UTIL_CHECK(capacity > 0);
      UTIL_CHECK(beginId + capacity <= arr.capacity());
      UTIL_CHECK(!dataPtr_);
      UTIL_CHECK(!ref_.isAssociated());
 
      // Copy data pointer and capacity
      dataPtr_ = arr.cArray() + beginId;
      capacity_ = capacity;
 
      // Associate ref_ member with reference counter of data owner
      arr.addReference(ref_);
   }
 
   /*
   * Associate this object with memory owned by a DeviceMemory object.
   */
   template <typename Data>
   void DeviceArray<Data>::associate(DeviceMemory& arr, int capacity)
   {
      UTIL_CHECK(arr.isAllocated());
      UTIL_CHECK(capacity > 0);
      UTIL_CHECK(capacity * sizeof(Data) <= arr.capacity());
      UTIL_CHECK(!dataPtr_);
      UTIL_CHECK(!ref_.isAssociated());
 
      // Copy data pointer and capacity
      dataPtr_ = (Data *) arr.cArray();
      capacity_ = capacity;
 
      // Associate ref_ member with reference counter of data owner
      arr.addReference(ref_);
   }
 
   /*
   * Dissociate this object from external device memory
   */
   template <typename Data>
   void DeviceArray<Data>::dissociate()
   {
      UTIL_CHECK(dataPtr_);
      UTIL_CHECK(ref_.isAssociated());
 
      dataPtr_ = nullptr; // reset to null
      capacity_ = 0;
      ref_.dissociate();
   }
 
   /*
   * Associate an external CountedReference with the ReferenceCounter.
   *
   * This function should called to create an association in which
   * this object is the data owner. The CountedReference parameter
   * should be owned by the data user.
   */
   template <typename Data>
   void DeviceArray<Data>::addReference(CountedReference& ref)
   {  ref.associate(refCounter_); }
 
   /*
   * Assignment from another DeviceArray<Data>.
   */
   template <typename Data>
   DeviceArray<Data>&
   DeviceArray<Data>::operator = (const DeviceArray<Data>& other)
   {
      // Check for self assignment
      if (this == &other) return *this;
 
      // Precondition - RHS DeviceArray<Data> must be allocated
      if (!other.isAllocated()) {
         UTIL_THROW("Other DeviceArray<Data> must be allocated.");
      }
 
      // If this is not allocated, then allocate
      if (!isAllocated()) {
         allocate(other.capacity());
      }
 
      // Require equal capacity values
      if (capacity_ != other.capacity_) {
         UTIL_THROW("Cannot assign arrays of unequal capacity");
      }
 
      // Copy elements
      cudaErrorCheck( cudaMemcpy(dataPtr_, other.cArray(),
                                 capacity_ * sizeof(Data),
                                 cudaMemcpyDeviceToDevice) );
 
      return *this;
   }
 
   /*
   * Assignment from RHS Util::DArray<Data>.
   */
   template <typename Data>
   DeviceArray<Data>&
   DeviceArray<Data>::operator = (const Util::DArray<Data>& other)
   {
      // Precondition
      if (!other.isAllocated()) {
         UTIL_THROW("RHS DArray<Data> must be allocated.");
      }
 
      // Allocate this if necessary
      if (!isAllocated()) {
         allocate(other.capacity());
      }
 
      // Require equal capacity values
      if (capacity_ != other.capacity()) {
         UTIL_THROW("Cannot assign arrays of unequal capacity");
      }
 
      // Copy elements
      cudaErrorCheck( cudaMemcpy(dataPtr_, other.cArray(),
                                 capacity_ * sizeof(Data),
                                 cudaMemcpyHostToDevice) );
 
      return *this;
   }
 
   /*
   * Get a pointer to the underlying C array.
   */
   template <typename Data>
   Data* DeviceArray<Data>::cArray()
   {
      UTIL_CHECK(dataPtr_);
      return dataPtr_;
   }
 
   /*
   * Get a pointer to const to the underlying C array.
   */
   template <typename Data>
   const Data* DeviceArray<Data>::cArray() const
   {
      UTIL_CHECK(dataPtr_);
      return dataPtr_;
   }
 
} // namespace Pscf
#endif