cu-sparse-matrix.h

Go to the documentation of this file.

// 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;

  CuSparseMatrix<Real> &operator =(const SparseMatrix<Real> &smat);

  CuSparseMatrix<Real> &operator =(const CuSparseMatrix<Real> &smat);

  template<typename OtherReal>
  void CopyFromSmat(const SparseMatrix<OtherReal> &smat);

  void CopyFromSmat(const CuSparseMatrix<Real> &smat,
                    MatrixTransposeType trans = kNoTrans);

  void SelectRows(const CuArray<int32> &row_indexes,
                  const CuSparseMatrix<Real> &smat_other);

  template<typename OtherReal>
  void CopyToSmat(SparseMatrix<OtherReal> *smat) const;

  void CopyElementsToVec(CuVectorBase<Real> *vec) const;

  void Swap(SparseMatrix<Real> *smat);

  void Swap(CuSparseMatrix<Real> *smat);

  void SetRandn(BaseFloat zero_prob);

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

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

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

  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);
  }

  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);
  }

  CuSparseMatrix(const CuArray<int32> &indexes, int32 dim,
                 MatrixTransposeType trans = kNoTrans);

  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));
  }

  void Resize(const MatrixIndexT num_rows, const MatrixIndexT num_cols,
              const MatrixIndexT nnz, MatrixResizeType resize_type = kSetZero);

  const Real* CsrVal() const {
    return csr_val_;
  }
  Real* CsrVal() {
    return csr_val_;
  }

  const int* CsrRowPtr() const {
    return csr_row_ptr_col_idx_;
  }
  int* CsrRowPtr() {
    return csr_row_ptr_col_idx_;
  }

  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_