// cudamatrix/cu-block-matrix.cc
  
  // Copyright 2013      Johns Hopkins University (author: Daniel Povey)
  
  // 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.
  
  
  #if HAVE_CUDA == 1
  #include <cuda_runtime_api.h>
  #include <cublas_v2.h>
  #endif
  
  #include <algorithm>
  #include "base/timer.h"
  #include "cudamatrix/cu-block-matrix.h"
  #include "cudamatrix/cu-matrix.h"
  #include "cudamatrix/cu-device.h"
  
  namespace kaldi {
  
  template<class Real>
  CuBlockMatrix<Real>::CuBlockMatrix() {
  #if HAVE_CUDA == 1
    cu_data_ = NULL;
  #endif
  }
  
  template<class Real>
  CuBlockMatrix<Real>::CuBlockMatrix(const std::vector<CuMatrix<Real> >&data) {
  #if HAVE_CUDA == 1
    cu_data_ = NULL;
  #endif
    block_data_.resize(data.size());
    MatrixIndexT row_offset = 0, col_offset = 0, max_num_rows = 0;
    for (size_t b = 0; b < data.size(); b++) {
      MatrixIndexT num_rows = data[b].NumRows(), num_cols = data[b].NumCols();
      KALDI_ASSERT(num_rows > 0 && num_cols > 0);
      BlockMatrixData block_data;
      block_data.num_rows = num_rows;
      block_data.num_cols = num_cols;
      block_data.row_offset = row_offset;
      block_data.col_offset = col_offset;
      row_offset += num_rows;
      col_offset += num_cols;
      max_num_rows = std::max(max_num_rows, num_rows);
      block_data_[b] = block_data;
    }
    num_rows_ = row_offset;
    data_.Resize(max_num_rows, col_offset);
    for (int32 b = 0; b < NumBlocks(); b++)
      Block(b).CopyFromMat(data[b]);
    SetCudaData();
  }
  
  
  template<class Real>
  const CuSubMatrix<Real> CuBlockMatrix<Real>::Block(int32 b) const {
    KALDI_ASSERT(static_cast<size_t>(b) < block_data_.size());
    const BlockMatrixData &block_data = block_data_[b];
    return CuSubMatrix<Real>(data_, 0, block_data.num_rows,
                             block_data.col_offset, block_data.num_cols);
  }
  
  template<class Real>
  CuSubMatrix<Real> CuBlockMatrix<Real>::Block(int32 b) {
    KALDI_ASSERT(static_cast<size_t>(b) < block_data_.size());
    BlockMatrixData &block_data = block_data_[b];
    return CuSubMatrix<Real>(data_, 0, block_data.num_rows,
                             block_data.col_offset, block_data.num_cols);
  }
  
  
  template<class Real>
  CuBlockMatrix<Real>::CuBlockMatrix(const CuBlockMatrix<Real> &other):
      data_(other.data_), block_data_(other.block_data_), num_rows_(other.num_rows_) {
  #if HAVE_CUDA == 1
    cu_data_ = NULL;
  #endif
    SetCudaData();
  }
  
  template<class Real>
  CuBlockMatrix<Real> &CuBlockMatrix<Real>::operator =(const CuBlockMatrix<Real> &other) {
    FreeCudaData();
    data_ = other.data_;
    block_data_ = other.block_data_;
    num_rows_ = other.num_rows_;
    SetCudaData();
    return *this;
  }
  
  template<class Real>
  void CuBlockMatrix<Real>::FreeCudaData() {
  #if HAVE_CUDA == 1
    if (cu_data_ != NULL) {
      if (CuDevice::Instantiate().Enabled()) {
        CuDevice::Instantiate().Free(cu_data_);
        cu_data_ = NULL;
      } else {
        KALDI_ERR << "CuBlockMatrix: you have CUDA data pointer but "
                  << "no GPU is enabled: likely code error.";
      }
    }
  #endif
  }
  
  
  template<class Real>
  void CuBlockMatrix<Real>::SetCudaData() {
  #if HAVE_CUDA == 1
    KALDI_ASSERT(cu_data_ == NULL);
    if (block_data_.size() == 0) return; // Nothing to do.
    if (CuDevice::Instantiate().Enabled()) {
      CuTimer tim;
      std::vector<CuBlockMatrixData> tmp_cu_data(NumBlocks());
      int32 row_offset = 0, col_offset = 0;
      for (size_t b = 0; b < NumBlocks(); b++) {
        CuSubMatrix<Real> this_mat = Block(b);
        CuBlockMatrixData &this_cu_data = tmp_cu_data[b];
        this_cu_data.row_offset = row_offset;
        this_cu_data.col_offset = col_offset;
        this_cu_data.matrix_dim = this_mat.Dim();
        this_cu_data.matrix_data = static_cast<void*>(this_mat.Data());
        row_offset += this_mat.NumRows();
        col_offset += this_mat.NumCols();
      }
      size_t size = NumBlocks() * sizeof(CuBlockMatrixData);
      cu_data_ = static_cast<CuBlockMatrixData*>(
          CuDevice::Instantiate().Malloc(size));
      CU_SAFE_CALL(cudaMemcpyAsync(cu_data_, &(tmp_cu_data[0]), size, 
                                   cudaMemcpyHostToDevice, cudaStreamPerThread));
      CU_SAFE_CALL(cudaStreamSynchronize(cudaStreamPerThread));
      CuDevice::Instantiate().AccuProfile(__func__, tim);    
    }
  #endif
  }
  
  template<class Real>
  void CuBlockMatrix<Real>::Swap(CuBlockMatrix<Real> *other) {
    data_.Swap(&other->data_);
    block_data_.swap(other->block_data_);
    std::swap(num_rows_, other->num_rows_);
  #if HAVE_CUDA == 1
    std::swap(cu_data_, other->cu_data_);
  #endif
  }
  
  template<class Real>
  void CuBlockMatrix<Real>::Write(std::ostream &os, bool binary) const {
    WriteToken(os, binary, "<CuBlockMatrix>");
    int32 num_blocks = NumBlocks();
    WriteBasicType(os, binary, num_blocks);
    for (int32 b = 0; b < num_blocks; b++)
      this->Block(b).Write(os, binary);
    WriteToken(os, binary, "</CuBlockMatrix>");  
  }
  
  
  template<class Real>
  void CuBlockMatrix<Real>::Read(std::istream &is, bool binary) {
    Destroy();
    int i = Peek(is, binary);
    std::vector<CuMatrix<Real> > data;
    if (i != static_cast<int>('<')) {
      // back-compatibility code so we can read the older format of
      // MixtureProbComponent.  This code should be deleted eventually.
      int32 size;
      ReadBasicType(is, binary, &size);
      KALDI_ASSERT(size >= 0);
      data.resize(size);
      for (int32 i = 0; i < size; i++)
        data[i].Read(is, binary);
    } else {
      ExpectToken(is, binary, "<CuBlockMatrix>");
      int32 size;
      ReadBasicType(is, binary, &size);
      KALDI_ASSERT(size >= 0);
      data.resize(size);
      for (int32 i = 0; i < size; i++)
        data[i].Read(is, binary);
      ExpectToken(is, binary, "</CuBlockMatrix>");    
    }
  
    CuBlockMatrix<Real> block_mat(data); // initializer from std::vector<CuMatrix<Real> > does
    // the main job of initialization.
    this->Swap(&block_mat);
  }
  
  template<class Real>
  void CuBlockMatrix<Real>::Destroy() {
    data_.Resize(0, 0);
    block_data_.clear();
    num_rows_ = 0;  
    FreeCudaData();
  }
  
  // Does *this = alpha A B + beta * *this, discarding elements outside
  // the block structure of the *this matrix. 
  template<class Real>
  void CuBlockMatrix<Real>::AddMatMat(
      BaseFloat alpha,
      const CuMatrix<Real> &A, MatrixTransposeType transA,
      const CuMatrix<Real> &B, MatrixTransposeType transB,
      BaseFloat beta) {
    MatrixIndexT A_num_rows = A.NumRows(), A_num_cols = A.NumCols(),
        A_row_stride = A.Stride(), A_col_stride = 1,
        B_num_rows = B.NumRows(), B_num_cols = B.NumCols(),
        B_row_stride = B.Stride(), B_col_stride = 1;
    if (transA == kTrans) {
      std::swap(A_num_rows, A_num_cols);
      std::swap(A_row_stride, A_col_stride);
    }
    if (transB == kTrans) {
      std::swap(B_num_rows, B_num_cols);
      std::swap(B_row_stride, B_col_stride);
    }
    KALDI_ASSERT(A_num_rows == NumRows() && B_num_cols == NumCols()
                 && A_num_cols == B_num_rows);
    if (NumBlocks() == 0) return; // empty matrix.
  #if HAVE_CUDA == 1
    if (CuDevice::Instantiate().Enabled()) {
      CuTimer tim;
  
      // (x,y,z) dimensions are (block-id, row-of-block, col-of-block)
      // First some logic to choose block dims...
      // we assume (which we can, safely) that CU1DBLOCK is <= the max threads per block.
      int32 x_blocksize = std::min(CU1DBLOCK, NumBlocks()); // x dim corresponds to block-idx.
      int32 max_block_rows = MaxBlockRows(), max_block_cols = MaxBlockCols();
      int32 y_blocksize = max_block_rows;
      while (y_blocksize * x_blocksize > CU1DBLOCK || y_blocksize > CU2DBLOCK)
        y_blocksize--;
      int32 z_blocksize = max_block_cols;
      while (z_blocksize * x_blocksize * y_blocksize > CU1DBLOCK || z_blocksize > CU2DBLOCK)
        z_blocksize--;
      
      dim3 dimBlock(x_blocksize, y_blocksize, z_blocksize);
      dim3 dimGrid(n_blocks(NumBlocks(), x_blocksize),
                   n_blocks(max_block_rows, y_blocksize),
                   n_blocks(max_block_cols, z_blocksize));
      cuda_block_add_mat_mat(dimGrid, dimBlock, cu_data_, NumBlocks(),
                             A.Data(), A_num_cols, A_row_stride, A_col_stride,
                             B.Data(), B_row_stride, B_col_stride, alpha, beta);
      CU_SAFE_CALL(cudaGetLastError());    
      CuDevice::Instantiate().AccuProfile(__func__, tim);    
    } else
  #endif
    {
      int32 row_offset = 0, col_offset = 0;    
      for (MatrixIndexT b = 0; b < NumBlocks(); b++) {
        CuSubMatrix<Real> this_block = Block(b);
        MatrixIndexT this_num_rows = this_block.NumRows(),
            this_num_cols = this_block.NumCols();
        CuSubMatrix<Real> A_part = (transA == kNoTrans ?
                                    A.Range(row_offset, this_num_rows,
                                            0, A.NumCols()) :
                                    A.Range(0, A.NumRows(),
                                            row_offset, this_num_rows)),
            B_part = (transB == kNoTrans ?
                      B.Range(0, B.NumRows(),
                              col_offset, this_num_cols) :
                      B.Range(col_offset, this_num_cols,
                              0, B.NumCols()));
        this_block.AddMatMat(alpha, A_part, transA, B_part, transB, beta);
        row_offset += this_num_rows;
        col_offset += this_num_cols;
      }
      KALDI_ASSERT(row_offset == NumRows() && col_offset == NumCols());
    }
  }
  
  template<class Real>
  MatrixIndexT CuBlockMatrix<Real>::MaxBlockCols() const {
    MatrixIndexT max_cols = 0;
    for (size_t i = 0; i < block_data_.size(); i++)
      max_cols = std::max(max_cols, block_data_[i].num_cols);
    return max_cols;
  }
  
  template<class Real>
  MatrixIndexT CuBlockMatrix<Real>::MaxBlockRows() const {
    return data_.NumRows();
  }
  
  template<class Real>
  void CuBlockMatrix<Real>::CopyFromMat(const CuMatrix<Real> &M) {
    KALDI_ASSERT(NumRows() == M.NumRows() && NumCols() == M.NumCols());
    MatrixIndexT row_offset = 0, col_offset = 0;
    for (MatrixIndexT b = 0; b < NumBlocks(); b++) {
      CuSubMatrix<Real> this_block = Block(b);
      MatrixIndexT this_num_rows = this_block.NumRows(),
          this_num_cols = this_block.NumCols();
      const CuSubMatrix<Real> src(M, row_offset, this_num_rows,
                                  col_offset, this_num_cols);
      this_block.CopyFromMat(src);
      row_offset += this_num_rows;
      col_offset += this_num_cols;
    }
    KALDI_ASSERT(row_offset == NumRows() && col_offset == NumCols());
  }
  
  /**
   * Print the matrix to stream
   */
  template<typename Real>
  std::ostream &operator << (std::ostream &out, const CuBlockMatrix<Real> &mat) {
    bool binary = false;
    mat.Write(out, binary);
    return out;
  }
  // instantiate the template
  template
  std::ostream &operator << (std::ostream &out, const CuBlockMatrix<float> &mat);
  template 
  std::ostream &operator << (std::ostream &out, const CuBlockMatrix<double> &mat);
  
  // Instantiate the class for float and double.
  template class CuBlockMatrix<float>;
  template class CuBlockMatrix<double>;
  
  } // namespace kaldi