// cudamatrix/cu-packed-matrix.cc
  
  // Copyright 2009-2013  Johns Hopkins University (author: Daniel Povey)
  //                      Karel Vesely
  
  // 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 "base/timer.h"
  #include "cudamatrix/cu-common.h"
  #include "cudamatrix/cu-vector.h"
  #include "cudamatrix/cu-device.h"
  #include "cudamatrix/cu-kernels.h"
  #include "cudamatrix/cu-math.h"
  #include "cudamatrix/cu-packed-matrix.h"
  #include "cudamatrix/cublas-wrappers.h"
  
  namespace kaldi {
  
  template<typename Real>
  void CuPackedMatrix<Real>::Resize(MatrixIndexT rows,
                                    MatrixResizeType resize_type) {
    // This code does not currently support the other resize_type options.
    KALDI_ASSERT(resize_type == kSetZero || resize_type == kUndefined);
  
    if (this->num_rows_ == rows) {
      if (resize_type == kSetZero) this->SetZero();
      return;
    }
  
    if (this->num_rows_ != 0)
      this->Destroy();
    if (rows == 0) return;
  #if HAVE_CUDA == 1
    CuDevice &device = CuDevice::Instantiate();
    if (device.Enabled()) {
      CuTimer tim;
      this->num_rows_ = rows;
      size_t nr = static_cast<size_t>(num_rows_),
          num_bytes = ((nr * (nr+1)) / 2) * sizeof(Real);
      this->data_ = static_cast<Real*>(device.Malloc(num_bytes));
  
      if (resize_type == kSetZero) this->SetZero();
      device.AccuProfile("CuPackedMatrix::Resize", tim);
    } else
  #endif
    { // Let the initializer of SpMatrix<Real> handle the allocation,
      // and then just do Swap which will switch the pointers.
      // This wastes a few instructions but is simple to code.
      SpMatrix<Real> mat(rows, resize_type);
      this->Swap(&mat);
    }
  }
  
  template<typename Real>
  void CuPackedMatrix<Real>::SetRandn() {
    if (num_rows_ != 0) {
      MatrixIndexT size = num_rows_ * (num_rows_ + 1) / 2;
      CuSubVector<Real> tmp(data_, size);
      CuRand<Real> rand;
      rand.RandGaussian(&tmp);
    }
  }
  
  template<typename Real>
  void CuPackedMatrix<Real>::Destroy() {
  #if HAVE_CUDA == 1
    if (CuDevice::Instantiate().Enabled()) {
      if (this->data_ != NULL) {
        CuDevice::Instantiate().Free(this->data_);
      }
    } else
  #endif
    {
      if (this->data_ != NULL) KALDI_MEMALIGN_FREE(this->data_);
    }
    this->data_ = NULL;
    this->num_rows_ = 0;
  }
  
  template<typename Real>
  void CuPackedMatrix<Real>::Swap(PackedMatrix<Real> *mat) {
  #if HAVE_CUDA == 1
    if (CuDevice::Instantiate().Enabled()) {
      if (this->num_rows_ == 0) {
        if (mat->num_rows_ != 0) {
          // *this is empty, but mat is nonempty.
          Resize(mat->num_rows_, kUndefined);
          CopyFromPacked(*mat);
          mat->Resize(0);
        }
        // else both are empty.
      } else { // *this is nonempty.
        if (mat->num_rows_ != 0) {
          // Both *this and *mat are nonempty.  Recurse to simpler cases.
          // this could be done more efficiently in the case where
          // the size does not change.
          PackedMatrix<Real> temp;
          this->Swap(&temp); // now temp is full, *this is empty.
          mat->Swap(&temp); // now mat has data from *this, temp has
          // data from mat.
          this->Swap(&temp); // copy data in mat to *this, which is now empty.
        } else { // *this is full but *mat is empty.
          mat->Resize(this->num_rows_, kUndefined);
          this->CopyToPacked(mat);
          this->Destroy();
        }
      }
    } else
  #endif
    {
      std::swap(mat->data_, this->data_);
      std::swap(mat->num_rows_, this->num_rows_);
    }
  }
  
  template<typename Real>
  void CuPackedMatrix<Real>::CopyFromPacked(const CuPackedMatrix<Real> &src) {
    KALDI_ASSERT(src.NumRows() == num_rows_);
  #if HAVE_CUDA == 1
    if (CuDevice::Instantiate().Enabled()) {
      if (num_rows_ == 0) return; // Nothing to do.
      CuTimer tim;
      size_t nr = static_cast<size_t>(num_rows_),
          num_bytes = ((nr * (nr+1)) / 2) * sizeof(Real);
  
      CU_SAFE_CALL(
        cudaMemcpyAsync(data_, src.data_, num_bytes, cudaMemcpyDeviceToDevice,
                        cudaStreamPerThread));
      CuDevice::Instantiate().AccuProfile("CuPackedMatrix::CopyFromPacked1",
                                          tim);
    } else
  #endif
    {
      Mat().CopyFromPacked(src.Mat());
    }
  }
  
  template<typename Real>
  void CuPackedMatrix<Real>::CopyFromPacked(const PackedMatrix<Real> &src) {
    KALDI_ASSERT(src.NumRows() == num_rows_);
  #if HAVE_CUDA == 1
    if (CuDevice::Instantiate().Enabled()) {
      if (num_rows_ == 0) return; // Nothing to do.
      CuTimer tim;
      CU_SAFE_CALL(cudaMemcpyAsync(data_, src.data_, src.SizeInBytes(),
                                   cudaMemcpyHostToDevice, cudaStreamPerThread));
      CU_SAFE_CALL(cudaStreamSynchronize(cudaStreamPerThread));
      CuDevice::Instantiate().AccuProfile("CuPackedMatrix::CopyFromPacked2", tim);
    } else
  #endif
    {
      Mat().CopyFromPacked(src);
      //memcpy(data_, src.Data(), SizeInBytes());
    }
  }
  
  template<typename Real>
  void CuPackedMatrix<Real>::CopyToPacked(PackedMatrix<Real> *dst) const {
    KALDI_ASSERT(dst->NumRows() == NumRows());
  
  #if HAVE_CUDA == 1
    if (CuDevice::Instantiate().Enabled()) {
      if (num_rows_ == 0) return; // Nothing to do.
      CuTimer tim;
      size_t nr = static_cast<size_t>(num_rows_),
        num_bytes = ((nr * (nr+1)) / 2) * sizeof(Real);
  
      CU_SAFE_CALL(cudaMemcpyAsync(dst->data_, data_, num_bytes,
                                   cudaMemcpyDeviceToHost, cudaStreamPerThread));
      CU_SAFE_CALL(cudaStreamSynchronize(cudaStreamPerThread));
      CuDevice::Instantiate().AccuProfile("CuPackedMatrix::CopyToPackedD2H", tim);
    } else
  #endif
    {
      //memcpy(data_, dst->Data(), SizeInBytes());
      dst->CopyFromPacked(Mat());
    }
  }
  
  /*
  template<typename Real>
  void CuPackedMatrix<Real>::CopyRowsFromPacked(int32 r, const CuPackedMatrix<Real> &src, int32 src_ro, int32 dst_ro) {
    KALDI_ASSERT(r+src_ro <= src.NumRows());
    KALDI_ASSERT(r+dst_ro <= NumRows());
    KALDI_ASSERT(NumCols() == src.NumCols());
  
    #if HAVE_CUDA == 1
    if (CuDevice::Instantiate().Enabled()) {
      CuTimer tim;
  
      MatrixIndexT dst_pitch = stride_*sizeof(Real);
      MatrixIndexT src_pitch = src.Stride()*sizeof(Real);
      MatrixIndexT width = src.NumCols()*sizeof(Real);
  
      const Real *p_src = src.Data() + src_ro*src.Stride();
      Real *p_dst = data_ + dst_ro*stride_;
  
      CU_SAFE_CALL(cudaMemcpy2D(p_dst, dst_pitch, p_src, src_pitch, width, r, cudaMemcpyDeviceToDevice));
  
      CuDevice::Instantiate().AccuProfile("CuMatrix::CopyRowsD2D", tim);
    } else
    #endif
    {
      memcpy(Data()+dst_ro*stride_, src.Data()+src_ro*src.Stride(), r*stride_*sizeof(Real));
    }
  } */
  
  
  
  template<typename Real>
  void CuPackedMatrix<Real>::Read(std::istream &is, bool binary) {
    PackedMatrix<Real> temp;
    temp.Read(is, binary);
    Destroy();
    Swap(&temp);
  }
  
  template<typename Real>
  void CuPackedMatrix<Real>::Write(std::ostream &os, bool binary) const {
    PackedMatrix<Real> temp(this->num_rows_, kUndefined);
    this->CopyToPacked(&temp);
    temp.Write(os, binary);
  }
  
  template<typename Real>
  void CuPackedMatrix<Real>::SetZero() {
    #if HAVE_CUDA == 1
    if (CuDevice::Instantiate().Enabled()) {
      CuTimer tim;
      size_t nr = static_cast<size_t>(num_rows_),
        num_bytes = ((nr * (nr+1)) / 2) * sizeof(Real);
  
      CU_SAFE_CALL(cudaMemsetAsync(reinterpret_cast<void*>(this->data_), 0, 
            num_bytes, cudaStreamPerThread));
      CuDevice::Instantiate().AccuProfile("CuPackedMatrix::SetZero", tim);
    } else
    #endif
    {
      Mat().SetZero();
    }
  }
  
  template<typename Real>
  Real CuPackedMatrix<Real>::Trace() const {
    Real result = 0.0;
  #if HAVE_CUDA == 1
    if (CuDevice::Instantiate().Enabled()) {
      if (num_rows_ == 0) return 0.0;
      CuVector<Real> tmp(num_rows_, kUndefined);
      tmp.CopyDiagFromPacked(*this);
      return tmp.Sum();
    } else
  #endif
    {
      result = Mat().Trace();
    }
    return result;
  }
  
  template<typename Real>
  void CuPackedMatrix<Real>::SetDiag(Real alpha) {
  #if HAVE_CUDA == 1
    if (CuDevice::Instantiate().Enabled()) {
      if (num_rows_ == 0) return;
      CuTimer tim;
      int dimBlock(CU1DBLOCK);
      int dimGrid(n_blocks(NumRows(),CU1DBLOCK));
      cuda_set_diag_packed(dimGrid,dimBlock,data_,alpha,num_rows_);
      CU_SAFE_CALL(cudaGetLastError());
      CuDevice::Instantiate().AccuProfile("CuPackedMatrix::SetDiag", tim);
    } else
  #endif
    {
      Mat().SetDiag(alpha);
    }
  }
  
  template<typename Real>
  void CuPackedMatrix<Real>::Scale(Real alpha) {
  #if HAVE_CUDA == 1
    if (CuDevice::Instantiate().Enabled()) {
      CuTimer tim;
      size_t nr = static_cast<size_t>(num_rows_),
          num_elements = ((nr * (nr+1)) / 2);
      CUBLAS_SAFE_CALL(cublas_scal(GetCublasHandle(), num_elements, alpha, data_, 1));
  
      CuDevice::Instantiate().AccuProfile("CuPackedMatrix::Scale", tim);
    } else
  #endif
    {
      Mat().Scale(alpha);
    }
  }
  
  template<typename Real>
  void CuPackedMatrix<Real>::ScaleDiag(Real alpha) {
  #if HAVE_CUDA == 1
    if (CuDevice::Instantiate().Enabled()) {
      CuTimer tim;
      int dimBlock(CU1DBLOCK);
      int dimGrid(n_blocks(NumRows(),CU1DBLOCK));
      cuda_scale_diag_packed(dimGrid,dimBlock,data_,alpha,num_rows_);
      CU_SAFE_CALL(cudaGetLastError());
      CuDevice::Instantiate().AccuProfile("CuPackedMatrix::ScaleDiag", tim);
    } else
  #endif
    {
      Mat().ScaleDiag(alpha);
    }
  }
  
  template<typename Real>
  void CuPackedMatrix<Real>::AddPacked(const Real alpha, const CuPackedMatrix<Real> &M) {
    KALDI_ASSERT(num_rows_ == M.NumRows());
  #if HAVE_CUDA == 1
    if (CuDevice::Instantiate().Enabled()) {
      if (num_rows_ == 0) return;
      CuTimer tim;
      size_t nr = num_rows_,
          sz = (nr * (nr + 1)) / 2;
      cublas_axpy(GetCublasHandle(), sz, alpha, M.Data(), 1, data_, 1);
      CuDevice::Instantiate().AccuProfile("CuPackedMatrix::AddPacked", tim);
    } else
  #endif
    {
      Mat().AddPacked(alpha, M.Mat());
    }
  }
  
  template<typename Real>
  void CuPackedMatrix<Real>::AddToDiag(Real r) {
  #if HAVE_CUDA == 1
    if (CuDevice::Instantiate().Enabled()) {
      if (num_rows_ == 0) return;
      CuTimer tim;
      int dimBlock(CU1DBLOCK);
      int dimGrid(n_blocks(NumRows(),CU1DBLOCK));
      cuda_add_diag_packed(dimGrid,dimBlock,data_,r,num_rows_);
      CU_SAFE_CALL(cudaGetLastError());
      CuDevice::Instantiate().AccuProfile("CuPackedMatrix::AddToDiag", tim);
    } else
  #endif
    {
      // TODO
      Mat().AddToDiag(r);
    }
  }
  
  template<typename Real>
  void CuPackedMatrix<Real>::SetUnit() {
  #if HAVE_CUDA == 1
    if (CuDevice::Instantiate().Enabled()) {
      this->SetZero();
      this->SetDiag(1.0);
    } else
  #endif
    {
      Mat().SetUnit();
    }
  }
  
  /**
   * Print the matrix to stream
   */
  template<typename Real>
  std::ostream &operator << (std::ostream &out, const CuPackedMatrix<Real> &mat) {
    PackedMatrix<Real> temp(mat.NumRows());
    mat.CopyToPacked(&temp);
    out << temp;
    return out;
  }
  
  // instantiate the template
  template
  std::ostream &operator << (std::ostream &out, const CuPackedMatrix<float> &mat);
  template
  std::ostream &operator << (std::ostream &out, const CuPackedMatrix<double> &mat);
  
  
  // Instantiate class CuPackedMatrix for float and double.
  template class CuPackedMatrix<float>;
  template class CuPackedMatrix<double>;
  
  
  } // namespace kaldi