// cudamatrix/cu-vector.h

// Copyright 2009-2012  Karel Vesely
//                      Johns Hopkins University (author: Daniel Povey)
//                      Lucas Ondel
//           2013       Xiaohui Zhang
//           2015       Guoguo Chen
//           2017       Daniel Galvez
//           2019       Yiwen Shao

// 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_VECTOR_H_
#define KALDI_CUDAMATRIX_CU_VECTOR_H_

#include "matrix/kaldi-vector.h"
#include "cudamatrix/cu-common.h"
#include "cudamatrix/cu-value.h"
#include "cudamatrix/cu-math.h"

namespace kaldi {

template<typename Real> class CuMatrixBase;

template<typename Real>
Real VecVec(const CuVectorBase<Real> &v1, const CuVectorBase<Real> &v2);

template<typename Real, typename OtherReal>
Real VecVec(const CuVectorBase<Real> &v1, const CuVectorBase<OtherReal> &v2);

/**
 * Vector for CUDA computing
 */
template<typename Real>
class CuVectorBase {
 public:
  friend class CuVectorBase<float>;
  friend class CuVectorBase<double>;
  friend class CuMatrixBase<Real>;
  friend class MatrixBase<Real>;
  friend class CuPackedMatrix<Real>;
  friend class CuSpMatrix<Real>;
  friend class CuTpMatrix<Real>;

  template <typename OtherReal>
  friend OtherReal VecVec(const CuVectorBase<OtherReal> &v1,
                          const CuVectorBase<OtherReal> &v2);
  friend void cu::Splice<Real>(const CuMatrixBase<Real> &src,
                               const CuArray<int32> &frame_offsets,
                               CuMatrixBase<Real> *tgt);
  friend class CuRand<Real>;

  /// Dimensions
  MatrixIndexT Dim() const { return dim_;  }

  /// Returns a pointer to the start of the vector's data.
  inline Real* Data() { return data_; }
  /// Returns a pointer to the start of the vector's data (const).
  inline const Real* Data() const { return data_; }

  /// Copy functions; these will crash if the dimension
  /// do not match.  The operator = in class CuVector will
  /// also change the sizes for you.
  void CopyFromVec(const CuVectorBase<Real> &src);

  template<typename OtherReal>
  void CopyFromVec(const CuVectorBase<OtherReal> &M);

  template<typename OtherReal>
  void CopyFromVec(const VectorBase<OtherReal> &src);

  template<typename OtherReal>
  void CopyToVec(VectorBase<OtherReal> *dst) const;

  void CopyRowsFromMat(const CuMatrixBase<Real> &M);

  void CopyRowsFromMat(const MatrixBase<Real> &M);

  /// Math operations
  void SetZero();
  void Set(Real value);
  void Add(Real value);
  void Scale(Real value);

  void AddVec(Real alpha, const CuVectorBase<Real> &vec, Real beta = 1.0);

  template<typename OtherReal>
  void AddVec(Real alpha, const CuVectorBase<OtherReal> &vec, Real beta = 1.0);

  /// Sum the rows of the matrix, add to vector
  void AddRowSumMat(Real alpha, const CuMatrixBase<Real> &mat, Real beta = 1.0);
  /// Sum the columns of the matrix, add to vector
  void AddColSumMat(Real alpha, const CuMatrixBase<Real> &mat, Real beta = 1.0);

  /// Add triangular matrix times vector: this <-- beta*this + alpha*M*v.
  /// Works even if rv == *this.
  void AddTpVec(const Real alpha, const CuTpMatrix<Real>&M,
                const MatrixTransposeType trans, const CuVectorBase<Real> &v,
                const Real beta);  // **beta previously defaulted to 0.0**

  /// Multiplies this vector by lower-triangular marix:  *this <-- *this *M
  void MulTp(const CuTpMatrix<Real> &M, const MatrixTransposeType trans);

  bool ApproxEqual(const CuVectorBase<Real> &other, float tol = 0.01) const;

  void InvertElements();


  /// Copies selected elements from 'mat' to *this.  Expects this->Dim()
  /// to equal elements.Dim(). If trans == kNoTrans,
  /// expects mat.NumRows() to equal this.Dim(), and for each i,
  /// copies mat(i, elements[i]) to (*this)(i).
  /// If trans == kTrans,
  /// expects mat.NumCols() to equal this.Dim(), and for each i,
  /// copies mat(elements[i], i) to (*this)(i).
  void CopyElements(const CuMatrixBase<Real> &mat,
                    const MatrixTransposeType trans,
                    const CuArrayBase<int32> &elements);

  void Floor(const CuVectorBase<Real> &src, Real floor_val, MatrixIndexT *floored_count = NULL);
  void Ceiling(const CuVectorBase<Real> &src, Real ceiling_val, MatrixIndexT *ceiled_count = NULL);
  void Pow(const CuVectorBase<Real> &src, Real power);
  
  inline void ApplyFloor(Real floor_val, MatrixIndexT *floored_count = NULL) {
    this -> Floor(*this, floor_val, floored_count); 
  };
  
  inline void ApplyCeiling(Real ceiling_val, MatrixIndexT *ceiled_count = NULL) {
    this -> Ceiling(*this, ceiling_val, ceiled_count);
  };
  
  inline void ApplyPow(Real power) {
    this -> Pow(*this, power);
  };

  void ApplySoftMax();
  void ApplyLogSoftMax();
  void ApplyExp();
  void ApplyLog();
  Real Sum() const;

  void SetRandn();
  void SetRandUniform();

  CuSubVector<Real> Range(const MatrixIndexT o, const MatrixIndexT l) {
    return CuSubVector<Real>(*this, o, l);
  }

  const CuSubVector<Real> Range(const MatrixIndexT o,
                                const MatrixIndexT l) const {
    return CuSubVector<Real>(*this, o, l);
  }

  void CopyColFromMat(const CuMatrixBase<Real> &mat, MatrixIndexT col);

  template<typename OtherReal>
  void CopyColFromMat(const CuMatrixBase<OtherReal> &mat, MatrixIndexT col);

  void AddMatVec(const Real alpha, const CuMatrixBase<Real> &M,
                 MatrixTransposeType trans, const CuVectorBase<Real> &v,
                 const Real beta);
  void AddVecVec(Real alpha, const CuVectorBase<Real> &v,
                 const CuVectorBase<Real> &r, Real beta);

  void AddSpVec(const Real alpha, const CuSpMatrix<Real> &S,
                const CuVectorBase<Real> &v, const Real beta);

  /// Add the diagonal of a matrix times itself:
  /// *this = diag(M M^T) +  beta * *this (if trans == kNoTrans), or
  /// *this = diag(M^T M) +  beta * *this (if trans == kTrans).
  void AddDiagMat2(Real alpha, const CuMatrixBase<Real> &M,
                   MatrixTransposeType trans, Real beta);

  /// Add the diagonal of a matrix product: *this = diag(M N), assuming the
  /// "trans" arguments are both kNoTrans; for transpose arguments, it behaves
  /// as you would expect.
  void AddDiagMatMat(Real alpha, const CuMatrixBase<Real> &M, MatrixTransposeType transM,
                     const CuMatrixBase<Real> &N, MatrixTransposeType transN,
                     Real beta = 1.0);

  inline CuValue<Real> operator() (MatrixIndexT i) {
    KALDI_PARANOID_ASSERT(static_cast<UnsignedMatrixIndexT>(i) <
                          static_cast<UnsignedMatrixIndexT>(dim_));
    return CuValue<Real>(data_ + i);
  }

  Real Norm(Real p); // Only works for p = 1 and p = 2.

  inline Real operator() (MatrixIndexT i) const {
    KALDI_PARANOID_ASSERT(static_cast<UnsignedMatrixIndexT>(i) <
                          static_cast<UnsignedMatrixIndexT>(dim_));
    return CuValue<Real>(data_ + i); // will be casted to Real.
  }

  /// Extracts the diagonal of a packed matrix M; works for Sp or Tp.
  void CopyDiagFromPacked(const CuPackedMatrix<Real> &M);

  /// Extracts the diagonal of a matrix.
  void CopyDiagFromMat(const CuMatrix<Real> &M);

  /// Returns the maximum value of any element, or -infinity for the empty vector.
  Real Max() const;

  /// Returns the minimum value of any element, or +infinity for the empty vector.
  Real Min() const;

  // Set each element to y = (x == orig ? changed : x).
  void ReplaceValue(Real orig, Real changed);

  // Multiplies (*this) by v elementwise: (*this)[i] *= v
  void MulElements(const CuVectorBase<Real> &v);

  // Divides (*this) by v elementwise: (*this)[i] /= v
  void DivElements(const CuVectorBase<Real> &v);

  // The following two functions should only be called if we did not compile
  // with CUDA or could not get a CUDA card; in that case the contents are
  // interpreted the same as a regular vector.
  // Do not use the following functions unless you know what you are doing!
  inline const VectorBase<Real> &Vec() const {
    return *(reinterpret_cast<const VectorBase<Real>* >(this));
  }
  inline VectorBase<Real> &Vec() {
    return *(reinterpret_cast<VectorBase<Real>* >(this));
  }

 protected:

  /// Default constructor: make it protected so the user cannot
  /// instantiate this class.
  CuVectorBase<Real>(): data_(NULL), dim_(0) { }

  Real *data_; ///< GPU data pointer (or regular data pointer
               ///< if CUDA is not compiled in or we have no GPU).
  MatrixIndexT dim_; ///< dimension of the vector

 private:
  KALDI_DISALLOW_COPY_AND_ASSIGN(CuVectorBase);
};

template<typename Real>
class CuVector: public CuVectorBase<Real> {
  friend class CuVectorBase<float>;
  friend class CuVectorBase<double>;
  friend class CuMatrixBase<Real>;
  friend class CuPackedMatrix<Real>;
  friend class CuSpMatrix<Real>;
  friend class CuTpMatrix<Real>;

 public:
  CuVector() { }
  CuVector(MatrixIndexT dim, MatrixResizeType t = kSetZero) { Resize(dim, t); }

  CuVector(const CuVectorBase<Real> &v);

  CuVector(const VectorBase<Real> &v);
  explicit CuVector(const CuVector<Real> &v) : CuVectorBase<Real>() {
    Resize(v.Dim(), kUndefined);
    this->CopyFromVec(v);
  }

  template<typename OtherReal>
  explicit CuVector(const CuVectorBase<OtherReal> &v) : CuVectorBase<Real>() {
    Resize(v.Dim(), kUndefined);
    this->CopyFromVec(v);
  }

  template<typename OtherReal>
  explicit CuVector(const VectorBase<OtherReal> &v) : CuVectorBase<Real>() {
    Resize(v.Dim(), kUndefined);
    this->CopyFromVec(Vector<Real>(v));
  }

  /// Allocate the memory
  void Resize(MatrixIndexT dim, MatrixResizeType t = kSetZero);

  ~CuVector() { Destroy(); }

  CuVector<Real> &operator = (const CuVectorBase<Real> &other) {
    Resize(other.Dim(), kUndefined);
    this->CopyFromVec(other);
    return *this;
  }

  CuVector<Real> &operator = (const CuVector<Real> &other) {
    Resize(other.Dim(), kUndefined);
    this->CopyFromVec(other);
    return *this;
  }
  CuVector<Real> &operator = (const VectorBase<Real> &other) {
    Resize(other.Dim());
    this->CopyFromVec(other);
    return *this;
  }

  void Swap(CuVector<Real> *vec);
  void Swap(Vector<Real> *vec);

  /// I/O
  void Read(std::istream &is, bool binary);
  void Write(std::ostream &is, bool binary) const;


 private:
  void Destroy();
};

// We'll fill out the following class if it's needed.
template<typename Real>
class CuSubVector: public CuVectorBase<Real> {
 public:
  CuSubVector(const CuVectorBase<Real> &t, const MatrixIndexT origin,
              const MatrixIndexT length) : CuVectorBase<Real>() {
    KALDI_ASSERT(static_cast<UnsignedMatrixIndexT>(origin)+
                 static_cast<UnsignedMatrixIndexT>(length) <=
                 static_cast<UnsignedMatrixIndexT>(t.Dim()));
    CuVectorBase<Real>::data_ = const_cast<Real*>(t.Data()+origin);
    CuVectorBase<Real>::dim_ = length;
  }
  /// Copy constructor
  /// this constructor needed for Range() to work in base class.
  CuSubVector(const CuSubVector &other) : CuVectorBase<Real> () {
    CuVectorBase<Real>::data_ = other.data_;
    CuVectorBase<Real>::dim_ = other.dim_;
  }

  CuSubVector(const Real* data, MatrixIndexT length) : CuVectorBase<Real> () {
    // 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.
    CuVectorBase<Real>::data_ = const_cast<Real*>(data);
    CuVectorBase<Real>::dim_ = length;
  }

  /// This operation does not preserve const-ness, so be careful.
  CuSubVector(const CuMatrixBase<Real> &matrix, MatrixIndexT row) {
    CuVectorBase<Real>::data_ = const_cast<Real*>(matrix.RowData(row));
    CuVectorBase<Real>::dim_ = matrix.NumCols();
  }


};

/// I/O
template<typename Real>
std::ostream &operator << (std::ostream &out, const CuVectorBase<Real> &vec);


template<typename Real>
bool ApproxEqual(const CuVectorBase<Real> &a,
                 const CuVectorBase<Real> &b, Real tol = 0.01) {
  return a.ApproxEqual(b, tol);
}

template<typename Real>
inline void AssertEqual(const CuVectorBase<Real> &a,
                        const CuVectorBase<Real> &b, Real tol = 0.01) {
  KALDI_ASSERT(a.ApproxEqual(b, tol));
}

template<typename Real>
template<typename OtherReal>
void CuVectorBase<Real>::CopyFromVec(const CuVectorBase<OtherReal> &v) {
  v.CopyToVec(&this);
}

template<typename Real>
template<typename OtherReal>
void VectorBase<Real>::CopyFromVec(const CuVectorBase<OtherReal> &cu) {
  cu.CopyToVec(this);
}

// declare template specializations.
template <>
template <>
void CuVectorBase<double>::CopyFromVec<float>(const CuVectorBase<float> &src);

template<>
template <>
void CuVectorBase<float>::CopyFromVec<double>(const CuVectorBase<double> &src);

template<typename Real>
template<typename OtherReal>
Vector<Real>::Vector(const CuVectorBase<OtherReal> &cu) {
  Init(cu.Dim());
  cu.CopyToVec(this);
}

/// Returns \f$ v_1^T M v_2  \f$ .
template<typename Real>
Real VecMatVec(const CuVectorBase<Real> &v1, const CuMatrixBase<Real> &M,
               const CuVectorBase<Real> &v2);


} // namespace

#endif