YAKL_ArrayBase.h

Go to the documentation of this file.

 
#pragma once
// Included by YAKL_Array.h
 
__YAKL_NAMESPACE_WRAPPER_BEGIN__
namespace yakl {
 
  // This implements all functionality used by all dynamically allocated arrays
  template <class T, int rank, int myMem, int myStyle>
  class ArrayBase {
  public:
 
    typedef typename std::remove_cv<T>::type       type;
    typedef          T                             value_type;
    typedef typename std::add_const<type>::type    const_value_type;
    typedef typename std::remove_const<type>::type non_const_value_type;
 
    T       * myData;         // Pointer to the flattened internal data
    index_t dimension[rank];  // Sizes of the 8 possible dimensions
    int     * refCount;       // Pointer shared by multiple copies of this Array to keep track of allcation / free
    #ifdef YAKL_DEBUG
 
      char const * myname;    // Label for debug printing. Only stored if debugging is turned on
    #endif
    #ifdef YAKL_ENABLE_STREAMS
      StreamList stream_dependencies;
    #else
      // This only exists to void extra data in the Array classes
      struct StreamListDummy {
        static bool constexpr empty() { return true; }
        void push_back( Stream stream ) { }
        int  size() { return 0; }
        Stream operator[] (int i) { return Stream(); }
      };
      StreamListDummy stream_dependencies;
    #endif
 
 
    void add_stream_dependency(Stream stream) {
      if constexpr (streams_enabled) {
        if (use_pool()) stream_dependencies.push_back(stream);
      }
    }
 
 
    void add_stream_dependencies(std::vector<Stream> streams) {
      if constexpr (streams_enabled) {
        if (use_pool()) {
          for (int i=0; i < streams.size(); i++) { stream_dependencies.push_back(streams[i]); }
        }
      }
    }
 
 
    void create_inform() {
      #ifdef YAKL_VERBOSE
        std::string msg = "Allocating ";
        if constexpr (myMem == memHost) {
          msg += std::string("host, ");
        } else {
          msg += std::string("device, ");
        }
        if constexpr (myStyle == styleC) {
          msg += std::string("C-style, ");
        } else {
          msg += std::string("Fortran-style, ");
        }
        msg += std::string("rank ") + std::to_string(rank) + std::string(" Array");
        msg += std::string(" of size ") + std::to_string(totElems()*sizeof(T)) + std::string(" bytes");
        verbose_inform(msg,this->label());
      #endif
    }
 
 
    void destroy_inform() {
      #ifdef YAKL_VERBOSE
        std::string msg = "Deallocating ";
        if constexpr (myMem == memHost) {
          msg += std::string("host, ");
        } else {
          msg += std::string("device, ");
        }
        if constexpr (myStyle == styleC) {
          msg += std::string("C-style, ");
        } else {
          msg += std::string("Fortran-style, ");
        }
        msg += std::string("rank ") + std::to_string(rank) + std::string(" Array");
        verbose_inform(msg,this->label());
      #endif
    }
 
 
    template <class ARR>
    void copy_inform(ARR const &dest) const {
      #ifdef YAKL_VERBOSE
        std::string msg = "Initiating ";
        if (myMem == memHost  ) msg += std::string("host to ");
        if (myMem == memDevice) msg += std::string("device to ");
        if (dest.get_memory_space() == memHost  ) msg += std::string("host memcpy of ");
        if (dest.get_memory_space() == memDevice) msg += std::string("device memcpy of ");
        msg += std::to_string(totElems()*sizeof(T)) + std::string(" bytes");
        if (label() != "") msg += " from Array labeled \"" + std::string(label()) + std::string("\"");
        if (dest.label() != "") msg += " to Array labeled \"" + std::string(dest.label()) + std::string("\"");
        verbose_inform(msg);
      #endif
    }
 
 
    // Deep copy this array's contents to another array that's on the host
    template <int theirRank, int theirStyle>
    inline void deep_copy_to(Array<typename std::remove_cv<T>::type,theirRank,memHost,theirStyle> const &lhs , Stream stream = Stream()) const {
      #ifdef YAKL_VERBOSE
        copy_inform(lhs);
      #endif
      #ifdef YAKL_DEBUG
        if (this->totElems() != lhs.totElems()) { yakl_throw("ERROR: deep_copy_to with different number of elements"); }
        if (this->myData == nullptr || lhs.myData == nullptr) { yakl_throw("ERROR: deep_copy_to with nullptr"); }
      #endif
      if (myMem == memHost) { memcpy_host_to_host  ( lhs.myData , this->myData , this->totElems()          ); }
      else                  { memcpy_device_to_host( lhs.myData , this->myData , this->totElems() , stream ); }
      #ifdef YAKL_AUTO_FENCE
        fence();
      #endif
    }
 
 
    // Deep copy this array's contents to another array that's on the device
    template <int theirRank, int theirStyle>
    inline void deep_copy_to(Array<typename std::remove_cv<T>::type,theirRank,memDevice,theirStyle> const &lhs , Stream stream = Stream()) const {
      #ifdef YAKL_VERBOSE
        copy_inform(lhs);
      #endif
      #ifdef YAKL_DEBUG
        if (this->totElems() != lhs.totElems()) { yakl_throw("ERROR: deep_copy_to with different number of elements"); }
        if (this->myData == nullptr || lhs.myData == nullptr) { yakl_throw("ERROR: deep_copy_to with nullptr"); }
      #endif
      if (myMem == memHost) { memcpy_host_to_device  ( lhs.myData , this->myData , this->totElems() , stream ); }
      else                  { memcpy_device_to_device( lhs.myData , this->myData , this->totElems() , stream ); }
      #ifdef YAKL_AUTO_FENCE
        fence();
      #endif
    }
 
 
    /* ACCESSORS */
    YAKL_INLINE int get_rank() const { return rank; }
    YAKL_INLINE index_t get_totElems() const {
      index_t tot = this->dimension[0];
      for (int i=1; i<rank; i++) { tot *= this->dimension[i]; }
      return tot;
    }
    YAKL_INLINE index_t get_elem_count() const { return get_totElems(); }
    YAKL_INLINE index_t totElems() const { return get_totElems(); }
    YAKL_INLINE index_t size() const { return get_totElems(); }
    YAKL_INLINE T *data() const { return this->myData; }
    YAKL_INLINE T *get_data() const { return this->myData; }
    YAKL_INLINE T *begin() const { return this->myData; }
    YAKL_INLINE T *end() const { return begin() + size(); }
    YAKL_INLINE bool span_is_contiguous() const { return true; }
    YAKL_INLINE bool initialized() const { return this->myData != nullptr; }
    YAKL_INLINE int get_memory_space() const { return myMem == memHost ? memHost : memDevice; }
    const char* label() const {
      #ifdef YAKL_DEBUG
        return this->myname;
      #else
        return "\"Unlabeled: YAKL_DEBUG CPP macro not defined\"";
      #endif
    }
    inline int use_count() const {
      if (this->refCount != nullptr) { return *(this->refCount); }
      else                           { return 0;                 }
    }
 
 
    // Allocate the array and the reference counter (if owned)
    template <class TLOC=T, typename std::enable_if< ! std::is_const<TLOC>::value , int >::type = 0>
    inline void allocate() {
      // static_assert( std::is_arithmetic<T>() || myMem == memHost , 
      //                "ERROR: You cannot use non-arithmetic types inside owned Arrays on the device" );
      yakl_mtx_lock();
      this->refCount = new int;
      (*(this->refCount)) = 1;
      if (myMem == memDevice) {
        this->myData = (T *) alloc_device( this->totElems()*sizeof(T) , this->label() );
      } else {
        this->myData = new T[this->totElems()];
      }
      #ifdef YAKL_VERBOSE
        this->create_inform();
      #endif
      yakl_mtx_unlock();
    }
 
 
    // Decrement the reference counter (if owned), and if it's zero after decrement, then deallocate the data
    // For const types, the pointer must be const casted to a non-const type before deallocation
    template <class TLOC=T, typename std::enable_if< std::is_const<TLOC>::value , int >::type = 0>
    inline void deallocate() {
      yakl_mtx_lock();
      typedef typename std::remove_cv<T>::type T_non_const;
      T_non_const *data = const_cast<T_non_const *>(this->myData);
      if (this->refCount != nullptr) {
        (*(this->refCount))--;
 
        if (*this->refCount == 0) {
          #ifdef YAKL_VERBOSE
            destroy_inform();
          #endif
          delete this->refCount;
          this->refCount = nullptr;
          if (this->totElems() > 0) {
            if (myMem == memDevice) {
              if (streams_enabled && use_pool() && get_yakl_instance().device_allocators_are_default && (! stream_dependencies.empty()) ) {
                std::vector<Event> event_dependencies;
                for (int i=0; i < stream_dependencies.size(); i++) {
                  event_dependencies.push_back( record_event(stream_dependencies[i]) );
                }
                get_yakl_instance().pool.free_with_event_dependencies( data , event_dependencies , this->label() );
              } else {
                free_device(data,this->label());
              }
            } else {
              delete[] data;
            }
            this->myData = nullptr;
          }
        }
 
      }
      yakl_mtx_unlock();
    }
 
 
    // Decrement the reference counter (if owned), and if it's zero after decrement, then deallocate the data
    template <class TLOC=T, typename std::enable_if< ! std::is_const<TLOC>::value , int >::type = 0>
    inline void deallocate() {
      yakl_mtx_lock();
      if (this->refCount != nullptr) {
        (*(this->refCount))--;
 
        if (*this->refCount == 0) {
          #ifdef YAKL_VERBOSE
            destroy_inform();
          #endif
          delete this->refCount;
          this->refCount = nullptr;
          if (this->totElems() > 0) {
            if (myMem == memDevice) {
              if (streams_enabled && use_pool() && get_yakl_instance().device_allocators_are_default && (! stream_dependencies.empty()) ) {
                std::vector<Event> event_dependencies;
                for (int i=0; i < stream_dependencies.size(); i++) {
                  event_dependencies.push_back( record_event(stream_dependencies[i]) );
                }
                get_yakl_instance().pool.free_with_event_dependencies( this->myData , event_dependencies , this->label() );
              } else {
                free_device(this->myData,this->label());
              }
            } else {
              delete[] this->myData;
            }
            this->myData = nullptr;
          }
        }
 
      }
      yakl_mtx_unlock();
    }
 
 
    // Print the array contents
    inline friend std::ostream &operator<<(std::ostream& os, Array<T,rank,myMem,myStyle> const &v) {
      os << "For Array labeled: " << v.label() << "\n";
      os << "Number of Dimensions: " << rank << "\n";
      os << "Total Number of Elements: " << v.totElems() << "\n";
      os << "Dimension Sizes: ";
      for (int i=0; i<rank; i++) {
        os << v.dimension[i] << ", ";
      }
      os << "\n";
      const_value_type     *local = v.myData;
      non_const_value_type *from_dev;
      if (myMem == memDevice) {
        from_dev = new non_const_value_type[v.totElems()];
        #ifdef YAKL_ENABLE_STREAMS
          fence();
        #endif
        memcpy_device_to_host( from_dev , v.myData , v.totElems() );
        fence();
        local = from_dev;
      }
      for (index_t i=0; i<v.totElems(); i++) {
        os << local[i] << " ";
      }
      if (myMem == memDevice) {
        delete[] from_dev;
      }
      os << "\n";
      return os;
    }
 
 
  };
 
}
__YAKL_NAMESPACE_WRAPPER_END__