// cudamatrix/cu-math.cc

// Copyright 2009-2012  Karel Vesely
//                      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.

#include "base/timer.h"
#include "cudamatrix/cu-common.h"
#include "cudamatrix/cu-matrix.h"
#include "cudamatrix/cu-device.h"
#include "cudamatrix/cu-kernels.h"

namespace kaldi {

namespace cu {

/*
 * templated functions wrapping the ANSI-C CUDA kernel functions
 */


template<typename Real>
void RegularizeL1(CuMatrixBase<Real> *weight, CuMatrixBase<Real> *grad, Real l1, Real lr) {
  KALDI_ASSERT(SameDim(*weight, *grad));
#if HAVE_CUDA == 1
  if (CuDevice::Instantiate().Enabled()) {
    CuTimer tim;

    dim3 dimBlock(CU2DBLOCK, CU2DBLOCK);
    dim3 dimGrid(n_blocks(weight->NumCols(), CU2DBLOCK), n_blocks(weight->NumRows(), CU2DBLOCK));

    cuda_regularize_l1(dimGrid, dimBlock, weight->Data(), grad->Data(), l1, lr,
                       weight->Dim(), grad->Stride());
    CU_SAFE_CALL(cudaGetLastError());

    CuDevice::Instantiate().AccuProfile(__func__, tim);
  } else
  #endif
  {
    MatrixBase<Real> &weight2 = weight->Mat();
    MatrixBase<Real> &grad2 = grad->Mat();
    for(MatrixIndexT r=0; r<weight2.NumRows(); r++) {
      for(MatrixIndexT c=0; c<weight2.NumCols(); c++) {

        if(weight2(r,c)==0.0) continue; // skip L1 if zero weightght!

        Real l1_signed = l1;
        if (weight2(r, c) < 0.0)
          l1_signed = -l1;

        Real before = weight2(r, c);
        Real after = weight2(r, c) - lr*grad2(r, c) - l1_signed;
        if ((after > 0.0) ^ (before > 0.0)) {
          weight2(r, c) = 0.0;
          grad2(r, c) = 0.0;
        } else {
          weight2(r, c) -= l1_signed;
        }
      }
    }
  }
}


template<typename Real>
void Randomize(const CuMatrixBase<Real> &src,
               const CuArray<int32> &copy_from_idx,
               CuMatrixBase<Real> *tgt) {

  KALDI_ASSERT(src.NumCols() == tgt->NumCols());
  KALDI_ASSERT(src.NumRows() == tgt->NumRows());
  KALDI_ASSERT(copy_from_idx.Dim() <= tgt->NumRows());

  #if HAVE_CUDA == 1
  if (CuDevice::Instantiate().Enabled()) {
    CuTimer tim;

    /*
    Note: default 16x16 block-size limits the --cachesize to matrix size 16*65535 x 16*65535
    dim3 dimBlock(CU2DBLOCK, CU2DBLOCK);
    dim3 dimGrid(n_blocks(tgt->NumCols(), CU2DBLOCK), n_blocks(copy_from_idx.Dim(), CU2DBLOCK));
    */

    /*
     * Let's use blocksize 4 x 128 (512 threads/block)
     * and extend the randomizable matrices to: col 4*65535, row 128*65535
     * (ie. max-cols:262140 (dim), max-rows:8388480 (datapoints))
     */
    dim3 dimBlock(4, 128);
    dim3 dimGrid(n_blocks(tgt->NumCols(), 4), n_blocks(copy_from_idx.Dim(), 128));
    /*
     */

    MatrixDim dimsrc = src.Dim(); dimsrc.rows=copy_from_idx.Dim();
    MatrixDim dimtgt = tgt->Dim(); dimtgt.rows=copy_from_idx.Dim();

    cuda_randomize(dimGrid, dimBlock, tgt->Data(), src.Data(),
                   copy_from_idx.Data(), dimtgt, dimsrc);
    CU_SAFE_CALL(cudaGetLastError());

    CuDevice::Instantiate().AccuProfile(__func__, tim);
  } else
  #endif
  {
    // randomize in CPU
    const MatrixBase<Real> &srcmat = src.Mat();
    const int32 *copy_from_idxvec = copy_from_idx.Data();
    MatrixBase<Real> &tgtmat = tgt->Mat();
    for(int32 i=0; i<copy_from_idx.Dim(); i++) {
      tgtmat.Row(i).CopyFromVec(srcmat.Row(copy_from_idxvec[i]));
    }
  }
}



template<typename Real>
void Splice(const CuMatrixBase<Real> &src, const CuArray<int32> &frame_offsets,
            CuMatrixBase<Real> *tgt) {

  KALDI_ASSERT(src.NumCols()*frame_offsets.Dim() == tgt->NumCols());
  KALDI_ASSERT(src.NumRows() == tgt->NumRows());

  #if HAVE_CUDA == 1
  if (CuDevice::Instantiate().Enabled()) {
    CuTimer tim;

    dim3 dimBlock(CU2DBLOCK, CU2DBLOCK);
    dim3 dimGrid(n_blocks(tgt->NumCols(), CU2DBLOCK), n_blocks(tgt->NumRows(), CU2DBLOCK));

    cuda_splice(dimGrid, dimBlock, tgt->Data(), src.Data(),
                frame_offsets.Data(), tgt->Dim(), src.Dim());
    CU_SAFE_CALL(cudaGetLastError());

    CuDevice::Instantiate().AccuProfile(__func__, tim);
  } else
  #endif
  {
    // expand in CPU
    const MatrixBase<Real> &srcmat = src.Mat();
    const int32 *frame_offsetvec = frame_offsets.Data();
    int32 dim = frame_offsets.Dim();
    MatrixBase<Real> &tgtmat = tgt->Mat();
    //
    for(int32 r=0; r < tgtmat.NumRows(); r++) {
      for(int32 off=0; off < dim; off++) {
        int32 r_off = r + frame_offsetvec[off];
        if(r_off < 0) r_off = 0;
        if(r_off >= srcmat.NumRows()) r_off = srcmat.NumRows()-1;
        memcpy(tgtmat.RowData(r)+off*srcmat.NumCols(),srcmat.RowData(r_off),sizeof(Real)*srcmat.NumCols());
      }
    }
  }
}



template<typename Real>
void Copy(const CuMatrixBase<Real> &src, const CuArray<int32> &copy_from_indices,
          CuMatrixBase<Real> *tgt) {

  KALDI_ASSERT(copy_from_indices.Dim() == tgt->NumCols());
  KALDI_ASSERT(src.NumRows() == tgt->NumRows());

  #if HAVE_CUDA == 1
  if (CuDevice::Instantiate().Enabled()) {
    CuTimer tim;

    dim3 dimBlock(CU2DBLOCK, CU2DBLOCK);
    dim3 dimGrid(n_blocks(tgt->NumCols(), CU2DBLOCK), n_blocks(tgt->NumRows(), CU2DBLOCK));

    cuda_copy(dimGrid, dimBlock, tgt->Data(), src.Data(),
              copy_from_indices.Data(), tgt->Dim(), src.Dim());
    CU_SAFE_CALL(cudaGetLastError());

    CuDevice::Instantiate().AccuProfile(__func__, tim);
  } else
  #endif
  {
    // expand in CPU
    const MatrixBase<Real> &srcmat = src.Mat();
    const int32 *copy_from_indicesvec = copy_from_indices.Data();
    int32 dim = copy_from_indices.Dim();
    MatrixBase<Real> &tgtmat = tgt->Mat();
    //
    for(int32 r = 0; r < tgtmat.NumRows(); r++) {
      for(int32 c = 0; c < dim; c++) {
        tgtmat(r,c) = srcmat(r,copy_from_indicesvec[c]);
      }
    }
  }
}

template <typename Real>
void EnsureNonzero(const CuMatrixBase<Real> &src,
                   Real epsilon,
                   CuMatrixBase<Real> *dest) {
  KALDI_ASSERT(SameDim(*dest, src) && epsilon > 0.0);
#if HAVE_CUDA == 1
  if (CuDevice::Instantiate().Enabled()) {
    CuTimer tim;
    dim3 dimGrid, dimBlock;
    GetBlockSizesForSimpleMatrixOperation(src.NumRows(), src.NumCols(),
                                          &dimGrid, &dimBlock);
    cuda_ensure_nonzero(dimGrid, dimBlock, src.Data(), src.Dim(),
                        epsilon, dest->Stride(), dest->Data());
    CU_SAFE_CALL(cudaGetLastError());
    CuDevice::Instantiate().AccuProfile(__func__, tim);
  } else
#endif
  {
    int32 num_rows = src.NumRows(), num_cols = src.NumCols();
    for (int32 r = 0; r < num_rows; r++) {
      const Real *src_data = src.RowData(r);
      Real *dest_data = dest->RowData(r);
      for (int32 c = 0; c < num_cols; c++) {
        Real x = src_data[c], y;
        if (x <= -epsilon || x >= epsilon) y = x;
        else if (x >= 0.0) y = epsilon;
        else y = -epsilon;
        dest_data[c] = y;
      }
    }
  }
}


// instantiate the templates.
template
void RegularizeL1(CuMatrixBase<float> *weight, CuMatrixBase<float> *grad, float l1, float lr);
template
void RegularizeL1(CuMatrixBase<double> *weight, CuMatrixBase<double> *grad, double l1, double lr);

template
void Splice(const CuMatrixBase<float> &src, const CuArray<int32> &frame_offsets,
            CuMatrixBase<float> *tgt);
template
void Splice(const CuMatrixBase<double> &src, const CuArray<int32> &frame_offsets,
            CuMatrixBase<double> *tgt);
template
void Copy(const CuMatrixBase<float> &src, const CuArray<int32> &copy_from_indices,
          CuMatrixBase<float> *tgt);
template
void Copy(const CuMatrixBase<double> &src, const CuArray<int32> &copy_from_indices,
          CuMatrixBase<double> *tgt);

template
void Randomize(const CuMatrixBase<float> &src,
               const CuArray<int32> &copy_from_idx,
               CuMatrixBase<float> *tgt);
template
void Randomize(const CuMatrixBase<double> &src,
               const CuArray<int32> &copy_from_idx,
               CuMatrixBase<double> *tgt);

// The output y_i = scale * x_i,
// and we want to RMS value of the y_i to equal target_rms,
// so y^t y = D * target_rms^2 (if y is one row of the input).
// we need to have scale = 1.0 / sqrt(x^t x / (D * target_rms^2)).
// there is also flooring involved, to avoid division-by-zero
// problems.  It's important for the backprop, that the floor's
// square root is exactly representable as float.
// If add_log_stddev_ is true, log(max(epsi, sqrt(x^t x / D)))
// is an extra dimension of the output.
template<typename Real>
void NormalizePerRow(const CuMatrixBase<Real>& in, const Real target_rms,
                     const bool add_log_stddev, CuMatrixBase<Real>* out) {
  const Real kSquaredNormFloor = 1.3552527156068805425e-20; // 2^-66
  if (add_log_stddev) {
    KALDI_ASSERT(in.NumRows() == out->NumRows());
    KALDI_ASSERT(in.NumCols() + 1 == out->NumCols());
  } else {
    KALDI_ASSERT(SameDim(in, *out));
  }

#if HAVE_CUDA == 1
  if (CuDevice::Instantiate().Enabled()) {
    CuTimer tim;
    size_t dimBlock = CU1DBLOCK;
    size_t dimGrid = out->NumRows();
    cuda_normalize_per_row(dimGrid, dimBlock, out->Data(), out->Stride(),
                           in.Data(), in.Dim(), target_rms, add_log_stddev);
    CU_SAFE_CALL(cudaGetLastError());
    CuDevice::Instantiate().AccuProfile(__func__, tim);
  } else
#endif
  {
    CuSubMatrix<Real> out_no_log(*out, 0, out->NumRows(), 0, in.NumCols());
    if (in.Data() != out_no_log.Data())
      out_no_log.CopyFromMat(in);
    CuVector<Real> in_norm(in.NumRows());
    Real d_scaled = in.NumCols() * target_rms * target_rms;
    in_norm.AddDiagMat2(1.0 / d_scaled, in, kNoTrans, 0.0);
    in_norm.ApplyFloor(kSquaredNormFloor);
    in_norm.ApplyPow(-0.5);
    out_no_log.MulRowsVec(in_norm);
    if (add_log_stddev) {
      in_norm.ApplyLog();
      in_norm.Scale(-1.0);
      in_norm.Add(log(target_rms));
      out->CopyColFromVec(in_norm, in.NumCols());
    }
  }
}

template
void NormalizePerRow(const CuMatrixBase<float>& in, const float target_rms,
                     const bool add_log_stddev, CuMatrixBase<float>* out);
template
void NormalizePerRow(const CuMatrixBase<double>& in, const double target_rms,
                     const bool add_log_stddev, CuMatrixBase<double>* out);


// A note on the derivative of NormalizeComponent...
// let both row_in and row_out be vectors of dimension D.
// Let p = row_in^T row_in / (D * target_rms^2), and let
// f = 1.0 / sqrt(max(kSquaredNormFloor, p)), and we compute row_out as:
// row_out = f row_in.
// Suppose we have a quantity deriv_out which is the derivative
// of the objective function w.r.t. row_out.  We want to compute
// deriv_in which is the derivative of the objective function w.r.t.
// row_in.  Let the objective function be F.  One term is obvious: we have
// deriv_in = f deriv_out + ....
// next we have to take into account the derivative that gets back-propagated
// through f.  Obviously, dF/df = deriv_out^T row_in.
// And df/dp = (p <= kSquaredNormFloor ? 0.0 : -0.5 p^{-1.5}) = (f == 1.0 / sqrt(kSquaredNormFloor) ? 0.0 : -0.5 f^3),
// and dp/d(row_in) = 2/(D * target_rms^2) row_in. [it's vector_valued].
// So this term in dF/d(row_in) equals:
// dF/df df/dp dp/d(row_in)   =    2/(D * target_rms^2) (f == 1.0 / sqrt(kSquaredNormFloor)  ? 0.0 : -0.5 f^3) (deriv_out^T row_in) row_in
// So
// deriv_in = f deriv_out + (f == 1.0 ? 0.0 : -f^3  / (D * target_rms^2) ) (deriv_out^T row_in) row_in
//  if add_log_stddev_ true, the deriv_in has another term as
// dF/dx_i = dF/df . df/dx_i => df/dx_i = x_i/(x^T x)
template<typename Real>
void DiffNormalizePerRow(const CuMatrixBase<Real> &in_value,
                         const CuMatrixBase<Real> &out_deriv,
                         const Real target_rms, const bool add_log_stddev,
                         CuMatrixBase<Real>* in_deriv) {
  const Real kSquaredNormFloor = 1.3552527156068805425e-20; // 2^-66
#if HAVE_CUDA == 1
  if (CuDevice::Instantiate().Enabled()) {
    CuTimer tim;
    size_t dimBlock = CU1DBLOCK;
    size_t dimGrid = in_deriv->NumRows();
    cuda_diff_normalize_per_row(dimGrid, dimBlock, in_deriv->Data(),
                                in_deriv->Stride(), in_value.Data(),
                                in_value.Dim(), out_deriv.Data(),
                                out_deriv.Stride(), target_rms, add_log_stddev);
    CU_SAFE_CALL(cudaGetLastError());
    CuDevice::Instantiate().AccuProfile(__func__, tim);
  } else
#endif
  {
    const CuSubMatrix<Real> out_deriv_no_log(out_deriv, 0, out_deriv.NumRows(),
                                             0, in_value.NumCols());
    CuVector<Real> dot_products(out_deriv.NumRows());
    dot_products.AddDiagMatMat(1.0, out_deriv_no_log, kNoTrans, in_value,
                               kTrans, 0.0);
    CuVector<Real> in_norm(in_value.NumRows());
    Real d_scaled = (in_value.NumCols() * target_rms * target_rms);
    in_norm.AddDiagMat2(1.0, in_value, kNoTrans, 0.0);

    if (add_log_stddev) {
      CuVector<Real> log_stddev_deriv(in_norm), // log_stddev deriv as dF/dy .* (x^T x)^-1
      out_deriv_for_stddev(out_deriv.NumRows(), kUndefined);
      // f = log(sqrt(max(epsi, x^T x / D)))
      // df/dx = epsi^2 * D < x^T x ? (1/(x^T x)) * x  : 0.
      // we don't compute this exactly below for the case when x^2 x is very
      // small, but we do make sure that the deriv isn't infinity when the input
      // is zero.
      log_stddev_deriv.ApplyFloor(in_value.NumCols() * kSquaredNormFloor);
      log_stddev_deriv.ApplyPow(-1.0);
      out_deriv_for_stddev.CopyColFromMat(out_deriv, (out_deriv.NumCols() - 1));
      log_stddev_deriv.MulElements(out_deriv_for_stddev);
      if (in_deriv)
        in_deriv->AddDiagVecMat(1.0, log_stddev_deriv, in_value, kNoTrans, 1.0);
    }
    in_norm.Scale(1.0 / d_scaled);
    in_norm.ApplyFloor(kSquaredNormFloor);
    in_norm.ApplyPow(-0.5);
    if (in_deriv) {
      if (in_deriv->Data() != out_deriv_no_log.Data())
        in_deriv->AddDiagVecMat(1.0, in_norm, out_deriv_no_log, kNoTrans, 1.0);
      else
        in_deriv->MulRowsVec(in_norm);
      in_norm.ReplaceValue(1.0 / sqrt(kSquaredNormFloor), 0.0);
      in_norm.ApplyPow(3.0);
      dot_products.MulElements(in_norm);

      in_deriv->AddDiagVecMat(-1.0 / d_scaled, dot_products, in_value, kNoTrans,
                              1.0);
    }
  }
}

template
void DiffNormalizePerRow(const CuMatrixBase<float> &in_value,
                         const CuMatrixBase<float> &out_deriv,
                         const float target_rms, const bool add_log_stddev,
                         CuMatrixBase<float>* in_deriv);
template
void DiffNormalizePerRow(const CuMatrixBase<double> &in_value,
                         const CuMatrixBase<double> &out_deriv,
                         const double target_rms, const bool add_log_stddev,
                         CuMatrixBase<double>* in_deriv);


// not calling this Sigmoid to reduce the chance of future collisions.
template<typename Real>
static inline Real ScalarSigmoid(Real a) {
  if (a > Real(0)) {
    return Real(1) / (Real(1) + Exp(-a));
  } else {
    Real x = Exp(a);
    return x / (x + Real(1));
  }
}

template<typename Real>
static inline Real ScalarTanh(Real a) {
  if (a > Real(0)) {
    Real inv_expa = Exp(-a);
    return -Real(1) + Real(2) / (Real(1) + inv_expa * inv_expa);
  } else {
    Real expa = Exp(a);
    return Real(1) - Real(2) / (Real(1) + expa * expa);
  }
}

template<typename Real>
void CpuComputeLstmNonlinearity(const MatrixBase<Real> &input_mat,
                                const MatrixBase<Real> &params_mat,
                                MatrixBase<Real> *output) {
  int32 num_rows = input_mat.NumRows(),
      input_cols = input_mat.NumCols(),
        cell_dim = input_cols / 5;
  KALDI_ASSERT(input_cols == (cell_dim * 5) || input_cols == (cell_dim * 5) + 3);
  KALDI_ASSERT(output->NumRows() == num_rows);
  KALDI_ASSERT(params_mat.NumRows() == 3);
  KALDI_ASSERT(params_mat.NumCols() == cell_dim);
  KALDI_ASSERT(output->NumCols() == 2 * cell_dim);

  MatrixBase<Real> &output_mat = *output;
  const Real *params_data = params_mat.Data();
  int32 params_stride = params_mat.Stride();
  for (int32 r = 0; r < num_rows; r++) {
    const Real *input_row = input_mat.RowData(r);
    // i_scale and f_scale relate to dropout, they will normally be 1.0.
    Real i_scale = (input_cols == cell_dim*5 ? 1.0:input_row[cell_dim*5]),
         f_scale = (input_cols == cell_dim*5 ? 1.0:input_row[cell_dim*5 + 1]),
         o_scale = (input_cols == cell_dim*5 ? 1.0:input_row[cell_dim*5 + 2]);

    Real *output_row = output_mat.RowData(r);
    for (int32 c = 0; c < cell_dim; c++) {
      Real i_part = input_row[c];
      Real f_part = input_row[c + cell_dim];
      Real c_part = input_row[c + 2 * cell_dim];
      Real o_part = input_row[c + 3 * cell_dim];
      Real c_prev = input_row[c + 4 * cell_dim];
      Real w_ic = params_data[c];
      Real w_fc = params_data[c + params_stride];
      Real w_oc = params_data[c + params_stride * 2];
      Real i_t = ScalarSigmoid(i_part + w_ic * c_prev);
      Real f_t = ScalarSigmoid(f_part + w_fc * c_prev);
      Real c_t = f_t * f_scale * c_prev + i_t * i_scale * ScalarTanh(c_part);
      Real o_t = ScalarSigmoid(o_part + w_oc * c_t);
      Real m_t = o_t * o_scale * ScalarTanh(c_t);
      output_row[c] = c_t;
      output_row[c + cell_dim] = m_t;
    }
  }
}

template<typename Real>
void ComputeLstmNonlinearity(const CuMatrixBase<Real> &input,
                             const CuMatrixBase<Real> &params,
                             CuMatrixBase<Real> *output) {
  int32 num_rows = input.NumRows(),
      input_cols = input.NumCols(),
        cell_dim = input_cols / 5;
  KALDI_ASSERT(input_cols == (cell_dim * 5) || input_cols == (cell_dim * 5) + 3);
  KALDI_ASSERT(output->NumRows() == num_rows);
  KALDI_ASSERT(params.NumRows() == 3);
  KALDI_ASSERT(params.NumCols() == cell_dim);
  KALDI_ASSERT(output->NumCols() == 2 * cell_dim);

#if HAVE_CUDA == 1
  if (CuDevice::Instantiate().Enabled()) {
    CuTimer tim;

    int have_dropout_mask = (input_cols == (cell_dim * 5) + 3);

    // Each thread block is working on 1 row of the data.
    // It's best that cell dim is a multiple fo CU1DBLOCK
    dim3 dimBlock(CU1DBLOCK);
    dim3 dimGrid(num_rows);

    cuda_lstm_nonlinearity(dimGrid, dimBlock, input.Data(), input.Stride(),
                           params.Data(), params.Stride(), output->Stride(),
                           cell_dim, have_dropout_mask, num_rows, output->Data());
    CU_SAFE_CALL(cudaGetLastError());

    CuDevice::Instantiate().AccuProfile(__func__, tim);
  } else
#endif
  {
    CpuComputeLstmNonlinearity(input.Mat(), params.Mat(), &output->Mat());
  }
}

template
void CpuComputeLstmNonlinearity(const MatrixBase<float> &input_mat,
                                const MatrixBase<float> &params_mat,
                                MatrixBase<float> *output);
template
void CpuComputeLstmNonlinearity(const MatrixBase<double> &input_mat,
                                const MatrixBase<double> &params_mat,
                                MatrixBase<double> *output);
template
void ComputeLstmNonlinearity(const CuMatrixBase<float> &input,
                             const CuMatrixBase<float> &params,
                             CuMatrixBase<float> *output);
template
void ComputeLstmNonlinearity(const CuMatrixBase<double> &input,
                             const CuMatrixBase<double> &params,
                             CuMatrixBase<double> *output);

template<typename Real>
void CpuBackpropLstmNonlinearity(const MatrixBase<Real> &input,
                                 const MatrixBase<Real> &params,
                                 const MatrixBase<Real> &output_deriv,
                                 const MatrixBase<double> &deriv_sum_in,
                                 const VectorBase<Real> &self_repair_config,
                                 double count_in,
                                 MatrixBase<Real> *input_deriv,
                                 MatrixBase<Real> *params_deriv,
                                 MatrixBase<double> *value_sum_out,
                                 MatrixBase<double> *deriv_sum_out,
                                 MatrixBase<Real> *self_repair_sum_out) {
  int32 num_rows = input.NumRows(),
      input_cols = input
                   .NumCols(),
        cell_dim = input.NumCols() / 5;
  // Check dimensions.
  KALDI_ASSERT(input_cols == (cell_dim * 5) || input_cols == (cell_dim * 5) + 3);
  KALDI_ASSERT(params.NumRows() == 3);
  KALDI_ASSERT(params.NumCols() == cell_dim);
  KALDI_ASSERT(output_deriv.NumRows() == num_rows);
  KALDI_ASSERT(output_deriv.NumCols() == 2 * cell_dim);
  KALDI_ASSERT(deriv_sum_in.NumRows() == 5);
  KALDI_ASSERT(deriv_sum_in.NumCols() == cell_dim);
  KALDI_ASSERT(self_repair_config.Dim() == 10);
  if (input_deriv != NULL) {
    KALDI_ASSERT(SameDim(input, *input_deriv));
  }
  if (params_deriv == NULL) {
    KALDI_ASSERT(value_sum_out == NULL);
    KALDI_ASSERT(deriv_sum_out == NULL);
    KALDI_ASSERT(self_repair_sum_out == NULL);
  } else {
    KALDI_ASSERT(value_sum_out != NULL);
    KALDI_ASSERT(deriv_sum_out != NULL);
    KALDI_ASSERT(self_repair_sum_out != NULL);
    KALDI_ASSERT(SameDim(params, *params_deriv));
    KALDI_ASSERT(value_sum_out->NumRows() == 5);
    KALDI_ASSERT(value_sum_out->NumCols() == cell_dim);
    KALDI_ASSERT(SameDim(*value_sum_out, *deriv_sum_out));
    KALDI_ASSERT(self_repair_sum_out->NumRows() == 5);
    KALDI_ASSERT(self_repair_sum_out->NumCols() == cell_dim);
  }

  const MatrixBase<Real> &input_mat = input;
  const MatrixBase<Real> &params_mat = params;
  const MatrixBase<Real> &output_deriv_mat = output_deriv;
  const MatrixBase<double> &deriv_sum_in_mat = deriv_sum_in;
  const VectorBase<Real> &sr_config = self_repair_config;
  MatrixBase<Real> *input_deriv_mat = (
      input_deriv == NULL ? NULL : input_deriv);
  MatrixBase<Real> *params_deriv_mat = NULL;
  MatrixBase<Real> *self_repair_sum_out_mat = NULL;
  MatrixBase<double> *value_sum_out_mat = NULL;
  MatrixBase<double> *deriv_sum_out_mat = NULL;
  if (params_deriv != NULL) {
    params_deriv_mat = params_deriv;
    value_sum_out_mat = value_sum_out;
    deriv_sum_out_mat = deriv_sum_out;
    self_repair_sum_out_mat = self_repair_sum_out;
  }


  // We add 1.0 (i.e. a small value) to the count to avoid division by zero.
  Real count = 1.0 + count_in;
  for (int32 c = 0; c < cell_dim; c++) {
    // parameters
    Real w_ic = params_mat(0, c);
    Real w_fc = params_mat(1, c);
    Real w_oc = params_mat(2, c);
    // derivative sums w.r.t. parameters.
    Real w_ic_deriv_sum = 0.0;
    Real w_fc_deriv_sum = 0.0;
    Real w_oc_deriv_sum = 0.0;

    // average derivatives, for self-repair.
    // The 5 nonlinearities that are subject to self-repair are written as:
    //  Sigmoid(i_t_input), Sigmoid(f_t_input),
    //  Tanh(c_part), Sigmoid(o_t_input),  Tanh(c_t)
    Real i_t_self_repair = (
        deriv_sum_in_mat(0, c) / count < sr_config(0) ? sr_config(5) : 0.0);
    Real f_t_self_repair = (
        deriv_sum_in_mat(1, c) / count < sr_config(1) ? sr_config(6) : 0.0);
    Real c_part_self_repair = (
        deriv_sum_in_mat(2, c) / count < sr_config(2) ? sr_config(7) : 0.0);
    Real o_t_self_repair = (
        deriv_sum_in_mat(3, c) / count < sr_config(3) ? sr_config(8) : 0.0);
    Real c_t_self_repair = (
        deriv_sum_in_mat(4, c) / count < sr_config(4) ? sr_config(9) : 0.0);
    // Note on how we add self-repair for sigmoids/tanh's.  If self-repair
    // is activated for this unit, then...
    // For sigmoids we'd add -self_repair_scale * (2 * sigmoid(x) - 1.0)
    // ... to the input-deriv;
    // For tanh's we'd add -self_repair_scale * tanh(x)
    // If self-repair is not activated, the 'self_repair' scales are set to zero.

    // The following variables are for the accumulation of stats on the
    // sigmoid and tanh units.
    Real i_t_value_sum = 0.0, i_t_deriv_sum = 0.0;
    Real f_t_value_sum = 0.0, f_t_deriv_sum = 0.0;
    Real c_part_value_sum = 0.0, c_part_deriv_sum = 0.0;
    Real o_t_value_sum = 0.0, o_t_deriv_sum = 0.0;
    Real c_t_value_sum = 0.0, c_t_deriv_sum = 0.0;


    for (int32 r = 0; r < num_rows; r++) {
      Real i_part = input_mat(r, c),
          f_part = input_mat(r, c + cell_dim),
          c_part = input_mat(r, c + 2 * cell_dim),
          o_part = input_mat(r, c + 3 * cell_dim),
          c_prev = input_mat(r, c + 4 * cell_dim);

      Real i_scale = (input_cols == cell_dim * 5 ? 1.0 :
                      input_mat(r, cell_dim * 5)),
           f_scale = (input_cols == cell_dim * 5 ? 1.0 :
                      input_mat(r, cell_dim * 5 + 1)),
           o_scale = (input_cols == cell_dim * 5 ? 1.0 :
                      input_mat(r, cell_dim * 5 + 2));

      // For greater clarity, we give some of the quantities in the
      // forward equations their own names.
      Real i_t_input = i_part + w_ic * c_prev,
          i_t = ScalarSigmoid(i_t_input),
          f_t_input = f_part + w_fc * c_prev,
          f_t = ScalarSigmoid(f_t_input),
          tanh_c_part = ScalarTanh(c_part),
          c_t = f_t * f_scale * c_prev + i_t * i_scale * tanh_c_part,
          o_t_input = o_part + w_oc * c_t,
          o_t = ScalarSigmoid(o_t_input),
          tanh_c_t = ScalarTanh(c_t);
      // we'd also compute, in the forward pass,
      //   m_t = o_t * tanh_c_t;
      // but this variable is not needed.

      // Accumulate nonlinearity value and derivative stats.
      // Note:
      //    tanh'(x)  = sech^2(x) = -(tanh(x)+1) (tanh(x)-1) = 1 - tanh^2(x)
      //  sigmoid'(x) = sigmoid(x) * (1 - sigmoid(x)).
      i_t_value_sum += i_t;
      i_t_deriv_sum += i_t * (1.0F - i_t);
      f_t_value_sum += f_t;
      f_t_deriv_sum += f_t * (1.0F - f_t);
      c_part_value_sum += tanh_c_part;
      c_part_deriv_sum += 1.0F - tanh_c_part * tanh_c_part;
      o_t_value_sum += o_t;
      o_t_deriv_sum += o_t * (1.0F - o_t);
      c_t_value_sum += tanh_c_t;
      c_t_deriv_sum += 1.0F - tanh_c_t * tanh_c_t;


      // the derivative of the objective function w.r.t. a particular quantity
      // will be written by prepending "d" to the name.
      // We compute these derivatives in the reverse of the order in which
      // we computed the original quantities.
      // dc_t_out is the part of the derivative w.r.t. c_t that
      // comes directly from the output of this function.
      Real dc_t_out = output_deriv_mat(r, c);
      Real dm_t = output_deriv_mat(r, c + cell_dim);
      Real dtanh_c_t = o_t * o_scale * dm_t;
      Real do_t = o_scale * tanh_c_t * dm_t;
      Real do_t_input = (o_t * (1.0F - o_t) * do_t
          - (2.0F * o_t - 1.0F) * o_t_self_repair);
      Real dc_t = ((1.0F - tanh_c_t * tanh_c_t) * dtanh_c_t + dc_t_out
          + do_t_input * w_oc) - tanh_c_t * c_t_self_repair;
      Real dtanh_c_part = i_t * i_scale * dc_t;
      Real df_t = dc_t * f_scale * c_prev;
      Real df_t_input = ((df_t * f_t * (1.0F - f_t)
                          - (2.0F * f_t - 1.0F) * f_t_self_repair));
      Real di_t = dc_t * i_scale * tanh_c_part;
      Real di_t_input = ((di_t * i_t * (1.0F - i_t)
                          - (2.0F * i_t - 1.0F) * i_t_self_repair));

      w_ic_deriv_sum += c_prev * di_t_input;
      w_fc_deriv_sum += c_prev * df_t_input;
      w_oc_deriv_sum += c_t * do_t_input;

      Real dc_prev = w_ic * di_t_input + w_fc * df_t_input + f_t * f_scale * dc_t;
      Real do_part = do_t_input;
      Real dc_part = ((1.0F - tanh_c_part * tanh_c_part) * dtanh_c_part
          - tanh_c_part * c_part_self_repair);
      Real df_part = df_t_input;
      Real di_part = di_t_input;

      if (input_deriv_mat != NULL) {
        (*input_deriv_mat)(r, c) = di_part;
        (*input_deriv_mat)(r, c + cell_dim) = df_part;
        (*input_deriv_mat)(r, c + 2 * cell_dim) = dc_part;
        (*input_deriv_mat)(r, c + 3 * cell_dim) = do_part;
        (*input_deriv_mat)(r, c + 4 * cell_dim) = dc_prev;
      }
    }

    if (params_deriv != NULL) {
      // note: for optimizing things you can assume that params_deriv and
      // input_deriv_mat are non-NULL (i.e. all the output matrices are
      // non-NULL).  The situations when some of the output matrices are NULL
      // does not happen often (mainly only in testing code).

      (*params_deriv_mat)(0, c) = w_ic_deriv_sum;
      (*params_deriv_mat)(1, c) = w_fc_deriv_sum;
      (*params_deriv_mat)(2, c) = w_oc_deriv_sum;

      (*value_sum_out_mat)(0, c) += i_t_value_sum;
      (*value_sum_out_mat)(1, c) += f_t_value_sum;
      (*value_sum_out_mat)(2, c) += c_part_value_sum;
      (*value_sum_out_mat)(3, c) += o_t_value_sum;
      (*value_sum_out_mat)(4, c) += c_t_value_sum;

      // need to update self_repair_sum_out before deriv_sum_out, because
      // deriv_sum_out and deriv_sum_in might point to the same memory.
      for (int32 i = 0; i < 5; i++)
        (*self_repair_sum_out_mat)(i, c) =
            (deriv_sum_in_mat(i, c) / count < sr_config(i) ? num_rows : 0);

      (*deriv_sum_out_mat)(0, c) += i_t_deriv_sum;
      (*deriv_sum_out_mat)(1, c) += f_t_deriv_sum;
      (*deriv_sum_out_mat)(2, c) += c_part_deriv_sum;
      (*deriv_sum_out_mat)(3, c) += o_t_deriv_sum;
      (*deriv_sum_out_mat)(4, c) += c_t_deriv_sum;
    }
  }
}



template<typename Real>
void BackpropLstmNonlinearity(const CuMatrixBase<Real> &input,
                              const CuMatrixBase<Real> &params,
                              const CuMatrixBase<Real> &output_deriv,
                              const CuMatrixBase<double> &deriv_sum_in,
                              const CuVectorBase<Real> &self_repair_config,
                              double count_in,
                              CuMatrixBase<Real> *input_deriv,
                              CuMatrixBase<Real> *params_deriv,
                              CuMatrixBase<double> *value_sum_out,
                              CuMatrixBase<double> *deriv_sum_out,
                              CuMatrixBase<Real> *self_repair_sum_out) {
  int32 num_rows = input.NumRows(),
        cell_dim = input.NumCols() / 5,
      input_cols = input.NumCols();
  // Check dimensions.
  KALDI_ASSERT(input_cols == (cell_dim * 5) || input_cols == (cell_dim*5) + 3);
  KALDI_ASSERT(params.NumRows() == 3);
  KALDI_ASSERT(params.NumCols() == cell_dim);
  KALDI_ASSERT(output_deriv.NumRows() == num_rows);
  KALDI_ASSERT(output_deriv.NumCols() == 2 * cell_dim);
  KALDI_ASSERT(deriv_sum_in.NumRows() == 5);
  KALDI_ASSERT(deriv_sum_in.NumCols() == cell_dim);
  KALDI_ASSERT(self_repair_config.Dim() == 10);
  if (input_deriv != NULL) {
    KALDI_ASSERT(SameDim(input, *input_deriv));
  }
  if (params_deriv == NULL) {
    KALDI_ASSERT(value_sum_out == NULL);
    KALDI_ASSERT(deriv_sum_out == NULL);
    KALDI_ASSERT(self_repair_sum_out == NULL);
  } else {
    KALDI_ASSERT(value_sum_out != NULL);
    KALDI_ASSERT(deriv_sum_out != NULL);
    KALDI_ASSERT(self_repair_sum_out != NULL);
    KALDI_ASSERT(SameDim(params, *params_deriv));
    KALDI_ASSERT(value_sum_out->NumRows() == 5);
    KALDI_ASSERT(value_sum_out->NumCols() == cell_dim);
    KALDI_ASSERT(SameDim(*value_sum_out, *deriv_sum_out));
    KALDI_ASSERT(self_repair_sum_out->NumRows() == 5);
    KALDI_ASSERT(self_repair_sum_out->NumCols() == cell_dim);
  }


#if HAVE_CUDA == 1
  if (CuDevice::Instantiate().Enabled()) {
    CuTimer tim;
    // Each thread block is working on 1 row of the data.
    // It's best that cell dim is a multiple fo CU1DBLOCK

    int have_dropout_mask = (input_cols == (cell_dim * 5) + 3);

    // Use 2D block (8x32 threads) as we need to compute column sum.
    // Use 1D grid to cover the data matrix width `cell_dim`.
    const int kWarpSize = 32;
    dim3 dimBlock(kWarpSize, CU1DBLOCK / kWarpSize);
//    dim3 dimGrid(n_blocks(cell_dim, dimBlock.x),
//                 n_blocks(num_rows, dimBlock.y));
//    if (dimGrid.x * dimGrid.y > 1024) {
//      dimGrid.y = std::max(1024 / dimGrid.x, 1);
//    }
    dim3 dimGrid(n_blocks(cell_dim, dimBlock.x));
    if (input_deriv == NULL) {
      if (params_deriv == NULL) {
        cuda_diff_lstm_nonlinearity(dimGrid, dimBlock, cell_dim,
                                    have_dropout_mask, num_rows,
                                    input.Data(), input.Stride(), params.Data(),
                                    params.Stride(), output_deriv.Data(),
                                    output_deriv.Stride(), deriv_sum_in.Data(),
                                    deriv_sum_in.Stride(),
                                    self_repair_config.Data(), count_in + 1,
                                    NULL,
                                    0,
                                    NULL,
                                    0,
                                    NULL,
                                    0,
                                    NULL,
                                    0,
                                    NULL,
                                    0);

      } else {
        cuda_diff_lstm_nonlinearity(dimGrid, dimBlock, cell_dim,
                                    have_dropout_mask, num_rows,
                                    input.Data(), input.Stride(), params.Data(),
                                    params.Stride(), output_deriv.Data(),
                                    output_deriv.Stride(), deriv_sum_in.Data(),
                                    deriv_sum_in.Stride(),
                                    self_repair_config.Data(), count_in + 1,
                                    NULL,
                                    0, params_deriv->Data(),
                                    params_deriv->Stride(),
                                    value_sum_out->Data(),
                                    value_sum_out->Stride(),
                                    deriv_sum_out->Data(),
                                    deriv_sum_out->Stride(),
                                    self_repair_sum_out->Data(),
                                    self_repair_sum_out->Stride());
      }
    } else {
      if (params_deriv == NULL) {
        cuda_diff_lstm_nonlinearity(dimGrid, dimBlock, cell_dim,
                                    have_dropout_mask, num_rows,
                                    input.Data(), input.Stride(), params.Data(),
                                    params.Stride(), output_deriv.Data(),
                                    output_deriv.Stride(), deriv_sum_in.Data(),
                                    deriv_sum_in.Stride(),
                                    self_repair_config.Data(), count_in + 1,
                                    input_deriv->Data(), input_deriv->Stride(),
                                    NULL,
                                    0, NULL, 0, NULL, 0, NULL, 0);
      } else {
        cuda_diff_lstm_nonlinearity(dimGrid, dimBlock, cell_dim,
                                    have_dropout_mask, num_rows,
                                    input.Data(), input.Stride(), params.Data(),
                                    params.Stride(), output_deriv.Data(),
                                    output_deriv.Stride(), deriv_sum_in.Data(),
                                    deriv_sum_in.Stride(),
                                    self_repair_config.Data(), count_in + 1,
                                    input_deriv->Data(), input_deriv->Stride(),
                                    params_deriv->Data(),
                                    params_deriv->Stride(),
                                    value_sum_out->Data(),
                                    value_sum_out->Stride(),
                                    deriv_sum_out->Data(),
                                    deriv_sum_out->Stride(),
                                    self_repair_sum_out->Data(),
                                    self_repair_sum_out->Stride());
      }
    }

    CU_SAFE_CALL(cudaGetLastError());

    CuDevice::Instantiate().AccuProfile(__func__, tim);
  } else
#endif
  {
    CpuBackpropLstmNonlinearity(input.Mat(), params.Mat(), output_deriv.Mat(),
                                deriv_sum_in.Mat(), self_repair_config.Vec(),
                                count_in, &(input_deriv->Mat()),
                                &(params_deriv->Mat()), &(value_sum_out->Mat()),
                                &(deriv_sum_out->Mat()),
                                &(self_repair_sum_out->Mat()));
  }
}

template <typename Real>
void EnsureNonzero(const CuVectorBase<Real> &src,
                   Real epsilon,
                   CuVectorBase<Real> *dest) {
  KALDI_ASSERT(src.Dim() == dest->Dim());
  int32 dim = src.Dim();
  // fake it with a 1-row matrix.
  CuSubMatrix<Real> src_mat(src.Data(), 1,  dim, dim),
      dest_mat(dest->Data(), 1, dim, dim);
  EnsureNonzero(src_mat, epsilon, &dest_mat);
}

// Instantiate the templates we defined above.

template
void EnsureNonzero(const CuMatrixBase<float> &src,
                   float epsilon,
                   CuMatrixBase<float> *dest);
template
void EnsureNonzero(const CuMatrixBase<double> &src,
                   double epsilon,
                   CuMatrixBase<double> *dest);

template
void EnsureNonzero(const CuVectorBase<float> &src,
                   float epsilon,
                   CuVectorBase<float> *dest);
template
void EnsureNonzero(const CuVectorBase<double> &src,
                   double epsilon,
                   CuVectorBase<double> *dest);

template
void CpuBackpropLstmNonlinearity(const MatrixBase<float> &input,
                                 const MatrixBase<float> &params,
                                 const MatrixBase<float> &output_deriv,
                                 const MatrixBase<double> &deriv_sum_in,
                                 const VectorBase<float> &self_repair_config,
                                 double count_in,
                                 MatrixBase<float> *input_deriv,
                                 MatrixBase<float> *params_deriv,
                                 MatrixBase<double> *value_sum_out,
                                 MatrixBase<double> *deriv_sum_out,
                                 MatrixBase<float> *self_repair_sum_out);
template
void CpuBackpropLstmNonlinearity(const MatrixBase<double> &input,
                                 const MatrixBase<double> &params,
                                 const MatrixBase<double> &output_deriv,
                                 const MatrixBase<double> &deriv_sum_in,
                                 const VectorBase<double> &self_repair_config,
                                 double count_in,
                                 MatrixBase<double> *input_deriv,
                                 MatrixBase<double> *params_deriv,
                                 MatrixBase<double> *value_sum_out,
                                 MatrixBase<double> *deriv_sum_out,
                                 MatrixBase<double> *self_repair_sum_out);
template
void BackpropLstmNonlinearity(const CuMatrixBase<float> &input,
                              const CuMatrixBase<float> &params,
                              const CuMatrixBase<float> &output_deriv,
                              const CuMatrixBase<double> &deriv_sum_in,
                              const CuVectorBase<float> &self_repair_config,
                              double count_in,
                              CuMatrixBase<float> *input_deriv,
                              CuMatrixBase<float> *params_deriv,
                              CuMatrixBase<double> *value_sum_out,
                              CuMatrixBase<double> *deriv_sum_out,
                              CuMatrixBase<float> *self_repair_sum_out);
template
void BackpropLstmNonlinearity(const CuMatrixBase<double> &input,
                              const CuMatrixBase<double> &params,
                              const CuMatrixBase<double> &output_deriv,
                              const CuMatrixBase<double> &deriv_sum_in,
                              const CuVectorBase<double> &self_repair_config,
                              double count_in,
                              CuMatrixBase<double> *input_deriv,
                              CuMatrixBase<double> *params_deriv,
                              CuMatrixBase<double> *value_sum_out,
                              CuMatrixBase<double> *deriv_sum_out,
                              CuMatrixBase<double> *self_repair_sum_out);



} //namespace cu

} //namespace kaldi