// cudamatrix/cu-sparse-matrix.h

// Copyright      2015  Johns Hopkins University (author: Daniel Povey)
//                2015  Guoguo Chen
//                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_SPARSE_MATRIX_H_
#define KALDI_CUDAMATRIX_CU_SPARSE_MATRIX_H_

#include <sstream>
#include <vector>

#include "cudamatrix/cu-matrixdim.h"
#include "cudamatrix/cu-common.h"
#include "cudamatrix/cu-value.h"
#include "matrix/matrix-common.h"
#include "matrix/kaldi-matrix.h"
#include "matrix/sparse-matrix.h"
#include "cudamatrix/cu-array.h"
#include "cudamatrix/cu-math.h"
#include "cudamatrix/cu-rand.h"

namespace kaldi {

template <typename Real>
Real TraceMatSmat(const CuMatrixBase<Real> &A,
                  const CuSparseMatrix<Real> &B,
                  MatrixTransposeType trans = kNoTrans);

template<class Real>
class CuSparseMatrix {
public:
  friend class CuMatrixBase<float> ;
  friend class CuMatrixBase<double> ;
  friend class CuMatrixBase<Real> ;
  friend class CuVectorBase<float> ;
  friend class CuVectorBase<double> ;
  friend class CuVectorBase<Real> ;

  friend Real TraceMatSmat<Real>(const CuMatrixBase<Real> &A,
                                 const CuSparseMatrix<Real> &B,
                                 MatrixTransposeType trans);

  MatrixIndexT NumRows() const;

  MatrixIndexT NumCols() const;

  MatrixIndexT NumElements() const;

  template<typename OtherReal>
  void CopyToMat(CuMatrixBase<OtherReal> *dest, MatrixTransposeType trans =
                     kNoTrans) const;

  Real Sum() const;

  Real FrobeniusNorm() const;

  /// Copy from CPU-based matrix.
  CuSparseMatrix<Real> &operator =(const SparseMatrix<Real> &smat);

  /// Copy from possibly-GPU-based matrix.
  CuSparseMatrix<Real> &operator =(const CuSparseMatrix<Real> &smat);

  /// Copy from CPU-based matrix.  We will add the transpose option later when it
  /// is necessary.  Resizes *this as needed.
  template<typename OtherReal>
  void CopyFromSmat(const SparseMatrix<OtherReal> &smat);

  /// Copy from GPU-based matrix, supporting transposition.  Resizes *this
  /// as needed.
  void CopyFromSmat(const CuSparseMatrix<Real> &smat,
                    MatrixTransposeType trans = kNoTrans);

  /// Select a subset of the rows of a CuSparseMatrix.
  /// Sets *this to only the rows of 'smat_other' that are listed
  /// in 'row_indexes'.
  /// 'row_indexes' must satisfy 0 <= row_indexes[i] < smat_other.NumRows().
  void SelectRows(const CuArray<int32> &row_indexes,
                  const CuSparseMatrix<Real> &smat_other);

  /// Copy to CPU-based matrix. We will add the transpose option later when it
  /// is necessary.
  template<typename OtherReal>
  void CopyToSmat(SparseMatrix<OtherReal> *smat) const;

  /// Copy elements to CuVector. It is the caller's responsibility to resize
  /// <*vec>.
  void CopyElementsToVec(CuVectorBase<Real> *vec) const;

  /// Swap with CPU-based matrix.
  void Swap(SparseMatrix<Real> *smat);

  /// Swap with possibly-CPU-based matrix.
  void Swap(CuSparseMatrix<Real> *smat);

  /// Sets up to a pseudo-randomly initialized matrix, with each element zero
  /// with probability zero_prob and else normally distributed- mostly for
  /// purposes of testing.
  void SetRandn(BaseFloat zero_prob);

  void Write(std::ostream &os, bool binary) const;

  void Read(std::istream &is, bool binary);

  /// Default constructor
  CuSparseMatrix() :
      num_rows_(0), num_cols_(0), nnz_(0), csr_row_ptr_col_idx_(NULL), csr_val_(
          NULL) {
  }

  /// Constructor from CPU-based sparse matrix.
  explicit CuSparseMatrix(const SparseMatrix<Real> &smat) :
      num_rows_(0), num_cols_(0), nnz_(0), csr_row_ptr_col_idx_(NULL), csr_val_(
      NULL) {
    this->CopyFromSmat(smat);
  }

  /// Constructor from GPU-based sparse matrix (supports transposition).
  CuSparseMatrix(const CuSparseMatrix<Real> &smat, MatrixTransposeType trans =
                     kNoTrans) :
      num_rows_(0), num_cols_(0), nnz_(0), csr_row_ptr_col_idx_(NULL), csr_val_(
      NULL) {
    this->CopyFromSmat(smat, trans);
  }

  /// Constructor from an array of indexes.
  /// If trans == kNoTrans, construct a sparse matrix
  /// with num-rows == indexes.Dim() and num-cols = 'dim'.
  /// 'indexes' is expected to contain elements in the
  /// range [0, dim - 1].  Each row 'i' of *this after
  /// calling the constructor will contain  a single
  /// element at column-index indexes[i] with value 1.0.
  ///
  /// If trans == kTrans, the result will be the transpose
  /// of the sparse matrix described above.
  CuSparseMatrix(const CuArray<int32> &indexes, int32 dim,
                 MatrixTransposeType trans = kNoTrans);

  /// Constructor from an array of indexes and an array of
  /// weights; requires indexes.Dim() == weights.Dim().
  /// If trans == kNoTrans, construct a sparse matrix
  /// with num-rows == indexes.Dim() and num-cols = 'dim'.
  /// 'indexes' is expected to contain elements in the
  /// range [0, dim - 1].  Each row 'i' of *this after
  /// calling the constructor will contain a single
  /// element at column-index indexes[i] with value weights[i].
  /// If trans == kTrans, the result will be the transpose
  /// of the sparse matrix described above.
  CuSparseMatrix(const CuArray<int32> &indexes,
                 const CuVectorBase<Real> &weights, int32 dim,
                 MatrixTransposeType trans = kNoTrans);

  ~CuSparseMatrix() {
    Destroy();
  }

protected:
  // 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 sparse matrix.
  inline const SparseMatrix<Real> &Smat() const {
    return *(reinterpret_cast<const SparseMatrix<Real>*>(this));
  }
  inline SparseMatrix<Real> &Smat() {
    return *(reinterpret_cast<SparseMatrix<Real>*>(this));
  }

  /// Users of this class won't normally have to use Resize.
  /// 'nnz' should be determined beforehand when calling this API.
  void Resize(const MatrixIndexT num_rows, const MatrixIndexT num_cols,
              const MatrixIndexT nnz, MatrixResizeType resize_type = kSetZero);

  /// Returns pointer to the data array of length nnz_ that holds all nonzero
  /// values in zero-based CSR format
  const Real* CsrVal() const {
    return csr_val_;
  }
  Real* CsrVal() {
    return csr_val_;
  }

  /// Returns pointer to the integer array of length NumRows()+1 that holds
  /// indices of the first nonzero element in the i-th row, while the last entry
  /// contains nnz_, as zero-based CSR format is used.
  const int* CsrRowPtr() const {
    return csr_row_ptr_col_idx_;
  }
  int* CsrRowPtr() {
    return csr_row_ptr_col_idx_;
  }

  /// Returns pointer to the integer array of length nnz_ that contains
  /// the column indices of the corresponding elements in array CsrVal()
  const int* CsrColIdx() const {
    return csr_row_ptr_col_idx_ + num_rows_ + 1;
  }
  int* CsrColIdx() {
    return csr_row_ptr_col_idx_ + num_rows_ + 1;
  }

private:
  void Destroy();

private:
  // This member is only used if we did not compile for the GPU, or if the GPU
  // is not enabled.  It needs to be first because we reinterpret_cast this
  std::vector<SparseVector<Real> > cpu_rows_;

  // This is where the data lives if we are using a GPU.
  // The sparse matrix is stored in CSR format, as documented here.
  // http://docs.nvidia.com/cuda/cusparse/index.html#compressed-sparse-row-format-csr
  // The 3 arrays are stored in 2 allocated blocks of memory.
  // Row ptr and col idx are both int arrays, thus stored in one block pointed
  // 'by csr_row_ptr_col_idx_'
  // Val are Real array, pointed by `csr_val_`

  // matrix size num_rows_ x num_cols_
  MatrixIndexT num_rows_;
  MatrixIndexT num_cols_;

  // number of non-zeros
  MatrixIndexT nnz_;

  // csr row ptrs and col indices in a single int array
  // of the length (num_rows_ + 1 + nnz_)
  int* csr_row_ptr_col_idx_;

  // csr value array of the length nnz_
  Real* csr_val_;
};


}  // namespace

#endif  // KALDI_CUDAMATRIX_CU_SPARSE_MATRIX_H_