// cudamatrix/cu-rand.cc

// Copyright 2016-2017  Brno University of Technology (author 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.

#include "cudamatrix/cu-rand.h"

namespace kaldi {

#if HAVE_CUDA == 1
/// Wrappers of curand functions to interface both float and double as 1 function,

/// Wrapper of curandGenerateUniform(), curandGenerateUniformDouble(),
template<typename Real>
curandStatus_t curandGenerateUniformWrap(curandGenerator_t gen, Real *ptr, size_t num);
//
template<>
curandStatus_t curandGenerateUniformWrap(curandGenerator_t gen, float *ptr, size_t num) {
  return curandGenerateUniform(gen, ptr, num);
}
template<>
curandStatus_t curandGenerateUniformWrap(curandGenerator_t gen, double *ptr, size_t num) {
  return curandGenerateUniformDouble(gen, ptr, num);
}

/// Wrapper of curandGenerateNormal(), curandGenerateNormalDouble(),
template<typename Real>
curandStatus_t curandGenerateNormalWrap(
    curandGenerator_t gen, Real *ptr, size_t num);
//
template<>
curandStatus_t curandGenerateNormalWrap<float>(
    curandGenerator_t gen, float *ptr, size_t num) {
  return curandGenerateNormal(gen, ptr, num, 0.0 /*mean*/, 1.0 /*stddev*/);
}
template<>
curandStatus_t curandGenerateNormalWrap<double>(
    curandGenerator_t gen, double *ptr, size_t num) {
  return curandGenerateNormalDouble(gen, ptr, num, 0.0 /*mean*/, 1.0 /*stddev*/);
}
/// End of wrappers.
#endif


template<typename Real>
void CuRand<Real>::RandUniform(CuMatrixBase<Real> *tgt) {
#if HAVE_CUDA == 1
  if (CuDevice::Instantiate().Enabled()) {
    CuTimer tim;
    // Better use 'tmp' matrix, 'tgt' can be a window into a larger matrix,
    // so we should not use it to generate random numbers over whole stride.
    // Use the option kStrideEqualNumCols to ensure consistency
    // (because when memory is nearly exhausted, the stride of CudaMallocPitch
    // may vary).
    CuMatrix<Real> tmp(tgt->NumRows(), tgt->NumCols(), kUndefined,
                       kStrideEqualNumCols);
    size_t s = static_cast<size_t>(tmp.NumRows()) * static_cast<size_t>(tmp.Stride());
    CURAND_SAFE_CALL(curandGenerateUniformWrap(
          GetCurandHandle(), tmp.Data(), s));
    tgt->CopyFromMat(tmp);
    CuDevice::Instantiate().AccuProfile(__func__, tim);
  } else
#endif
  {
    tgt->Mat().SetRandUniform();
  }
}

template<typename Real>
void CuRand<Real>::RandUniform(CuMatrix<Real> *tgt) {
#if HAVE_CUDA == 1
  if (CuDevice::Instantiate().Enabled()) {
    CuTimer tim;
    // Here we don't need to use 'tmp' matrix,
    size_t s = static_cast<size_t>(tgt->NumRows()) * static_cast<size_t>(tgt->Stride());
    CURAND_SAFE_CALL(curandGenerateUniformWrap(
          GetCurandHandle(), tgt->Data(), s));
    CuDevice::Instantiate().AccuProfile(__func__, tim);
  } else
#endif
  {
    tgt->Mat().SetRandUniform();
  }
}

template<typename Real>
void CuRand<Real>::RandUniform(CuVectorBase<Real> *tgt) {
#if HAVE_CUDA == 1
  if (CuDevice::Instantiate().Enabled()) {
    CuTimer tim;
    CURAND_SAFE_CALL(curandGenerateUniformWrap(
          GetCurandHandle(), tgt->Data(), tgt->Dim()));
    CuDevice::Instantiate().AccuProfile(__func__, tim);
  } else
#endif
  {
    tgt->Vec().SetRandUniform();
  }
}

template<typename Real>
void CuRand<Real>::RandGaussian(CuMatrixBase<Real> *tgt) {
#if HAVE_CUDA == 1
  if (CuDevice::Instantiate().Enabled()) {
    CuTimer tim;
    // Better use 'tmp' matrix, 'tgt' can be a window into a larger matrix,
    // so we should not use it to generate random numbers over whole stride.
    // Also, we ensure to have 'even' number of elements for calling 'curand'
    // by possibly adding one column. Even number of elements is required by
    // curandGenerateUniform(), curandGenerateUniformDouble().
    // Use the option kStrideEqualNumCols to ensure consistency
    // (because when memory is nearly exhausted, the stride of CudaMallocPitch
    // may vary).
    MatrixIndexT num_cols_even = tgt->NumCols() + (tgt->NumCols() % 2); // + 0 or 1,
    CuMatrix<Real> tmp(tgt->NumRows(), num_cols_even, kUndefined,
                       kStrideEqualNumCols);
    CURAND_SAFE_CALL(curandGenerateNormalWrap(
          GetCurandHandle(), tmp.Data(), tmp.NumRows()*tmp.Stride()));
    tgt->CopyFromMat(tmp.ColRange(0,tgt->NumCols()));
    CuDevice::Instantiate().AccuProfile(__func__, tim);
  } else
#endif
  {
    tgt->Mat().SetRandn();
  }
}

template<typename Real>
void CuRand<Real>::RandGaussian(CuMatrix<Real> *tgt) {
#if HAVE_CUDA == 1
  if (CuDevice::Instantiate().Enabled()) {
    CuTimer tim;
    // Here we don't need to use 'tmp' matrix, if the number of elements is even,
    MatrixIndexT num_elements = tgt->NumRows() * tgt->Stride();
    if (0 == (num_elements % 2)) {
      CURAND_SAFE_CALL(curandGenerateNormalWrap(
            GetCurandHandle(), tgt->Data(), num_elements));
    } else {
      // We use 'tmp' matrix with one column added, this guarantees an even
      // number of elements.  Use the option kStrideEqualNumCols to ensure
      // consistency (because when memory is nearly exhausted, the stride of
      // CudaMallocPitch may vary).
      MatrixIndexT num_cols_even = tgt->NumCols() + (tgt->NumCols() % 2); // + 0 or 1,
      CuMatrix<Real> tmp(tgt->NumRows(), num_cols_even, kUndefined,
                         kStrideEqualNumCols);
      CURAND_SAFE_CALL(curandGenerateNormalWrap(
            GetCurandHandle(), tmp.Data(), tmp.NumRows() * tmp.Stride()));
      tgt->CopyFromMat(tmp.ColRange(0,tgt->NumCols()));
    }
    CuDevice::Instantiate().AccuProfile(__func__, tim);
  } else
#endif
  {
    tgt->Mat().SetRandn();
  }
}

template<typename Real>
void CuRand<Real>::RandGaussian(CuVectorBase<Real> *tgt) {
#if HAVE_CUDA == 1
  if (CuDevice::Instantiate().Enabled()) {
    CuTimer tim;
    // To ensure 'even' number of elements, we use 'tmp' vector of even length.
    // Even number of elements is required by 'curand' functions:
    // curandGenerateUniform(), curandGenerateUniformDouble().
    MatrixIndexT num_elements = tgt->Dim();
    if (0 == (num_elements % 2)) {
      CURAND_SAFE_CALL(curandGenerateNormalWrap(
            GetCurandHandle(), tgt->Data(), tgt->Dim()));
    } else {
      MatrixIndexT dim_even = tgt->Dim() + (tgt->Dim() % 2); // + 0 or 1,
      CuVector<Real> tmp(dim_even, kUndefined);
      CURAND_SAFE_CALL(curandGenerateNormalWrap(
            GetCurandHandle(), tmp.Data(), tmp.Dim()));
      tgt->CopyFromVec(tmp.Range(0,tgt->Dim()));
    }
    CuDevice::Instantiate().AccuProfile(__func__, tim);
  } else
#endif
  {
    tgt->Vec().SetRandn();
  }
}

/// convert probabilities binary values,
template<typename Real>
void CuRand<Real>::BinarizeProbs(const CuMatrix<Real> &probs, CuMatrix<Real> *states) {
  CuMatrix<Real> tmp(probs.NumRows(), probs.NumCols());
  this->RandUniform(&tmp);  // [0..1]
  tmp.Scale(-1.0);  // [-1..0]
  tmp.AddMat(1.0, probs);  // [-1..+1]
  states->Heaviside(tmp);  // negative
}

/// add gaussian noise to each element
template<typename Real>
void CuRand<Real>::AddGaussNoise(CuMatrix<Real> *tgt, Real gscale) {
  // Use the option kStrideEqualNumCols to ensure consistency (because when
  // memory is nearly exhausted, the stride of CudaMallocPitch may vary).
  CuMatrix<Real> tmp(tgt->NumRows(), tgt->NumCols(),
                     kUndefined, kStrideEqualNumCols);
  this->RandGaussian(&tmp);
  tgt->AddMat(gscale, tmp);
}

// explicit instantiation,
template class CuRand<float>;
template class CuRand<double>;

}  // namespace,