// cudamatrix/cu-array.h
  
  // Copyright 2009-2012  Karel Vesely
  //                2013  Johns Hopkins University (author: Daniel Povey)
  //                2017  Shiyin Kang
  
  // See ../../COPYING for clarification regarding multiple authors
  //
  // Licensed under the Apache License, Version 2.0 (the "License");
  // you may not use this file except in compliance with the License.
  // You may obtain a copy of the License at
  //
  //  http://www.apache.org/licenses/LICENSE-2.0
  //
  // THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  // KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
  // WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
  // MERCHANTABLITY OR NON-INFRINGEMENT.
  // See the Apache 2 License for the specific language governing permissions and
  // limitations under the License.
  
  
  
  #ifndef KALDI_CUDAMATRIX_CU_ARRAY_H_
  #define KALDI_CUDAMATRIX_CU_ARRAY_H_
  
  #include "matrix/kaldi-vector.h"
  
  
  namespace kaldi {
  
  template <typename T> class CuArray;
  template <typename T> class CuSubArray;
  
  
  /**
     Class CuArrayBase, CuSubArray and CuArray are analogues of classes
     CuVectorBase, CuSubVector and CuVector, except that they are intended to
     store things other than float/double: they are intended to store integers or
     small structs.  Their CPU-based equivalents are std::vector, and we provide
     ways to copy to/from a std::vector of the same type.
  */
  template<typename T>
  class CuArrayBase {
    friend class CuArray<T>;
    friend class CuSubArray<T>;
   public:
    /// Return the vector dimension
    MatrixIndexT Dim() const { return dim_;  }
  
    /// Get raw pointer
    const T* Data() const { return data_; }
  
    T* Data() { return data_; }
  
    /// Sets the memory for the object to zero, via memset.  You should verify
    /// that this makes sense for type T.
    void SetZero();
  
    /// The caller is responsible to ensure dim is equal between *this and src.
    /// Note: copying to GPU is done via memcpy,
    /// and any constructors or assignment operators are not called.
    void CopyFromArray(const CuArrayBase<T> &src);
  
    /// The caller is responsible to ensure dim is equal between *this and src.
    /// Note: copying to GPU is done via memcpy,
    /// and any constructors or assignment operators are not called.
    void CopyFromVec(const std::vector<T> &src);
  
    /// This function resizes *dst if needed.  On resize of "dst", the STL vector
    /// may call copy-constructors, initializers, and assignment operators for
    /// existing objects (which will be overwritten), but the copy from GPU to CPU
    /// is done via memcpy.  So be very careful calling this function if your
    /// objects are more than plain structs.
    void CopyToVec(std::vector<T> *dst) const;
  
    /// Version of the above function that copies contents to a host array
    /// (i.e. to regular memory, not GPU memory, assuming we're using a GPU).
    /// This function requires *dst to be allocated before calling. The allocated
    /// size should be dim_ * sizeof(T)
    void CopyToHost(T *dst) const;
  
  
    /// Set to a constant value.  Note: any copying is done as if using memcpy, and
    /// assignment operators or destructors are not called.  This is NOT IMPLEMENTED
    /// YET except for T == int32 (the current implementation will just crash).
    void Set(const T &value);
  
    /// Fill with the sequence [base ... base + Dim())
    /// This is not implemented except for T=int32
    void Sequence(const T base);
  
    /// Add a constant value. This is NOT IMPLEMENTED YET except for T == int32
    /// (the current implementation will just crash).
    void Add(const T &value);
  
    /// Get minimum value (for now implemented on CPU, reimplement if slow).
    /// Asserts the vector is non-empty, otherwise crashes.
    T Min() const;
  
    /// Get minimum value (for now implemented on CPU, reimplement if slow).
    /// Asserts the vector is non-empty, otherwise crashes.
    T Max() const;
  
   protected:
    /// Default constructor: make it protected so the user cannot
    /// instantiate this class.
    CuArrayBase<T>(): data_(NULL), dim_(0) { }
  
  
    T *data_;  ///< GPU data pointer (if GPU not available,
               ///< will point to CPU memory).
    MatrixIndexT dim_;     ///< dimension of the vector
  
  };
  
  /**
     Class CuArray represents a vector of an integer or struct of type T.  If we
     are using a GPU then the memory is on the GPU, otherwise it's on the CPU.
     This class owns the data that it contains from a memory allocation
     perspective; see also CuSubArrary which does not own the data it contains.
   */
  template<typename T>
  class CuArray: public CuArrayBase<T> {
   public:
  
    /// Default constructor, initialized data_ to NULL and dim_ to 0 via
    /// constructor of CuArrayBase.
    CuArray<T>() { }
  
    /// Constructor with memory initialisation.  resize_type may be kSetZero or
    /// kUndefined.
    explicit CuArray<T>(MatrixIndexT dim, MatrixResizeType resize_type = kSetZero)
       { Resize(dim, resize_type); }
  
    /// Constructor from CPU-based int vector
    explicit CuArray<T>(const std::vector<T> &src) { CopyFromVec(src); }
  
    /// Copy constructor.  We don't make this explicit because we want to be able
    /// to create a std::vector<CuArray>.
    CuArray<T>(const CuArray<T> &src) { CopyFromArray(src); }
  
    /// Destructor
    ~CuArray() { Destroy(); }
  
    /// Allocate the memory.  resize_type may be kSetZero or kUndefined.
    /// kCopyData not yet supported (can be implemented if needed).
    void Resize(MatrixIndexT dim, MatrixResizeType resize_type = kSetZero);
  
    /// Deallocate the memory and set dim_ and data_ to zero.  Does not call any
    /// destructors of the objects stored.
    void Destroy();
  
    /// This function resizes if needed.  Note: copying to GPU is done via memcpy,
    /// and any constructors or assignment operators are not called.
    void CopyFromVec(const std::vector<T> &src);
  
    /// This function resizes if needed.
    void CopyFromArray(const CuArrayBase<T> &src);
  
    CuArray<T> &operator= (const CuArray<T> &in) {
      this->CopyFromArray(in); return *this;
    }
  
    CuArray<T> &operator= (const std::vector<T> &in) {
      this->CopyFromVec(in); return *this;
    }
  
    /// Shallow swap with another CuArray<T>.
    void Swap(CuArray<T> *other);
  
    /// I/O
    void Read(std::istream &is, bool binary);
    void Write(std::ostream &is, bool binary) const;
  
  };
  
  
  template<typename T>
  class CuSubArray: public CuArrayBase<T> {
   public:
    /// Constructor as a range of an existing CuArray or CuSubArray.  Note: like
    /// similar constructors in class CuVector and others, it can be used to evade
    /// 'const' constraints; don't do that.
    explicit CuSubArray<T>(const CuArrayBase<T> &src,
                           MatrixIndexT offset, MatrixIndexT dim);
  
    /// Construct from raw pointers
    CuSubArray(const T* data, MatrixIndexT length) {
      // Yes, we're evading C's restrictions on const here, and yes, it can be used
      // to do wrong stuff; unfortunately the workaround would be very difficult.
      CuArrayBase<T>::data_ = const_cast<T*>(data);
      CuArrayBase<T>::dim_ = length;
    }
  };
  
  
  
  /// I/O
  template<typename T>
  std::ostream &operator << (std::ostream &out, const CuArray<T> &vec);
  
  } // namespace
  
  #include "cudamatrix/cu-array-inl.h"
  
  #endif