ffCudaNn.cu

#include <assert.h>
#include <cuda_runtime.h>

#include <chrono>
#include <random>

#include "ffCudaNn.h"

#define K_THREAD_PER_BLOCK 512
#define K_SMALL_THREAD_PER_BLOCK 64

namespace ff {
///////////////////////////////////////////////////////////////////////
// std::default_random_engine g_generator;
std::default_random_engine g_generator(static_cast<int>(
    std::chrono::steady_clock::now().time_since_epoch().count()));
std::uniform_real_distribution<float> g_uniformDistribution;
static std::normal_distribution<float> g_normalDistribution(0.0f, 1.0f);

CudaTensor::CudaTensor()
    : _d0(0),
      _d1(0),
      _d2(0),
      _d3(0),
      _dataSize(0),
      _dataGpu(nullptr),
      _dataGpuSize(0) {}

CudaTensor::CudaTensor(int d0, int d1, int d2, int d3)
    : _dataGpu(nullptr), _dataGpuSize(0) {
  ResetTensor(d0, d1, d2, d3);
}

CudaTensor::~CudaTensor() {
  if (nullptr != _dataGpu) cudaFree(_dataGpu);
}

CudaTensor::CudaTensor(const CudaTensor& rhs) {
  assert(0);  // DISALLOW_COPY_AND_ASSIGN
}

CudaTensor& CudaTensor::operator=(const CudaTensor& rhs) {
  assert(0);  // DISALLOW_COPY_AND_ASSIGN
  return *this;
}

void CudaTensor::ResetTensor(int d0, int d1, int d2, int d3) {
  _d0 = d0;
  _d1 = d1;
  _d2 = d2;
  _d3 = d3;
  _dataSize = _d0 * _d1 * _d2 * _d3;
  _data.resize(_dataSize);
  if (_dataGpuSize < _dataSize) {
    _dataGpuSize = _dataSize;
    if (_dataGpu) cudaFree(_dataGpu);
    cudaError_t err = cudaMalloc(&_dataGpu, _dataGpuSize * sizeof(float));
    assert(err == cudaSuccess);
  }
}

void CudaTensor::Reshape(int d0, int d1, int d2, int d3) {
  assert((d0 * d1 * d2 * d3) == _dataSize);
  _d0 = d0;
  _d1 = d1;
  _d2 = d2;
  _d3 = d3;
}

void CudaTensor::SetRandom(const float multiplier) {
  for (int i = 0; i < _dataSize; ++i) {
    _data[i] = (g_uniformDistribution(g_generator) * 2.0f - 1.0f) * multiplier;
  }
  PushToGpu();
}

void CudaTensor::SetZero() {
  memset(&_data[0], 0, _data.size() * sizeof(float));
  PushToGpu();
}

void CudaTensor::SetDropoutMask(float zeroRatio) {
  assert(zeroRatio > 0.0f && zeroRatio < 1.0f);
  float s = 1.0f / (1.0f - zeroRatio);
  for (int i = 0; i < _dataSize; ++i) {
    _data[i] = 0.0f;
    if (g_uniformDistribution(g_generator) > zeroRatio) _data[i] = s;
  }
  PushToGpu();
}

void CudaTensor::PushToGpu() {
  cudaError_t err = cudaMemcpy(_dataGpu, &_data[0], _dataSize * sizeof(float),
                               cudaMemcpyKind::cudaMemcpyHostToDevice);
  assert(err == cudaSuccess);
}

void CudaTensor::PullFromGpu() {
  cudaError_t err = cudaMemcpy(&_data[0], _dataGpu, _dataSize * sizeof(float),
                               cudaMemcpyKind::cudaMemcpyDeviceToHost);
  assert(err == cudaSuccess);
}

void CudaTensor::WriteDataToFile(std::ofstream& file) {
  PullFromGpu();
  for (int i = 0; i < _dataSize; i++) {
    file << std::fixed << _data[i] << "\n";
  }
}

void CudaTensor::LoadDataFromFile(FILE* file) {
  for (int i = 0; i < _dataSize; i++) {
    fscanf(file, "%f", &(_data[i]));
  }
  PushToGpu();
}

///////////////////////////////////////////////////////////////////////
FcLayer::FcLayer(CudaNn* nn, int inDim, int outDim)
    : CudaLayer(nn), _pX(nullptr) {
  _w.ResetTensor(outDim, inDim);
  _w.SetRandom(1.0f / sqrtf((float)inDim));
  _wG.ResetTensor(outDim, inDim);
  _wG_m.ResetTensor(outDim, inDim);
  _wG_m.SetZero();
  _wG_v.ResetTensor(outDim, inDim);
  _wG_v.SetZero();
  _b.ResetTensor(outDim);
  _b.SetRandom(1.0f / sqrtf((float)inDim));
  _bG.ResetTensor(outDim);
  _bG_m.ResetTensor(outDim);
  _bG_m.SetZero();
  _bG_v.ResetTensor(outDim);
  _bG_v.SetZero();
}

__global__ void LinearTransform_Cuda(float* y, const float* x, const float* w,
                                     const float* b, int nColX, int nColW,
                                     int nJobs) {
  const int yIndex = blockIdx.x * blockDim.x + threadIdx.x;
  if (yIndex >= nJobs) return;
  int rowX = yIndex / nColW;
  int colW = yIndex % nColW;

  float v = 0.0f;
  int xBaseIndex = rowX * nColX;
  for (int i = 0; i < nColX; ++i) {
    v += x[i + xBaseIndex] * w[colW + i * nColW];
  }
  y[yIndex] = v + b[colW];
}

const CudaTensor* FcLayer::Forward(const CudaTensor* x) {
  int x_d0 = x->_d0;
  int x_d1 = x->_d1;
  if (x_d0 != _w._d1) {
    x_d0 = x->_d0 * x->_d1 * x->_d2;
    x_d1 = x->_d3;
  }
  assert(x_d0 == _w._d1);

  _pX = x;
  _y.ResetTensor(_w._d0, x_d1);

  // y = xw+b
  int nJobs = x_d1 * _w._d0;
  int numBlocks = (nJobs + K_THREAD_PER_BLOCK - 1) / K_THREAD_PER_BLOCK;
  dim3 block(numBlocks), threads(K_THREAD_PER_BLOCK);
  LinearTransform_Cuda<<<block, threads>>>(_y._dataGpu, _pX->_dataGpu,
                                           _w._dataGpu, _b._dataGpu, x_d0,
                                           _w._d0, nJobs);
  assert(cudaGetLastError() == cudaSuccess);
  return &_y;
}

__global__ void BackwardFc_Wg_Cuda(float* wG, const float* x, const float* yG,
                                   int nXcol, int nXrow, int nWcol, int nJobs) {
  const int wGindex = blockIdx.x * K_THREAD_PER_BLOCK + threadIdx.x;
  if (wGindex >= nJobs) return;
  int r = wGindex / nWcol;
  int c = wGindex % nWcol;

  // wG = x.T * yG
  float val = 0.0f;
  for (int i = 0; i < nXrow; ++i) {
    val += x[r + i * nXcol] * yG[c + i * nWcol];
  }
  // wG[wGindex] = val / nXrow;
  wG[wGindex] = val;
}

__global__ void BackwardFc_Bg_Cuda(float* bG, const float* yG, int nYgCol,
                                   int nYgRow) {
  const int c = blockIdx.x * blockDim.x + threadIdx.x;
  if (c >= nYgCol) return;

  float val = 0.0f;
  for (int i = 0; i < nYgRow; ++i) {
    val += yG[c + i * nYgCol];
  }
  // bG[c] = val / nYgRow;
  bG[c] = val;
}

__global__ void BackwardFc_Xg_Cuda(float* xG, const float* yG, const float* w,
                                   int yGw, int wTh, int xGw, int nJobs) {
  const int xGindex = blockIdx.x * blockDim.x + threadIdx.x;
  if (xGindex >= nJobs) return;
  int r = xGindex / xGw;
  int c = xGindex % xGw;

  // xG = yG * w.T
  float val = 0.0f;
  int yGbaseIndex = r * yGw;
  int wBaseIndex = c * wTh;
  for (int i = 0; i < yGw; ++i) {
    val += yG[i + yGbaseIndex] * w[i + wBaseIndex];
  }
  xG[xGindex] = val;
}

const CudaTensor* FcLayer::Backward(const CudaTensor* yG,
                                    const int layerIndex) {
  assert(yG->_d0 == _wG._d0);
  int x_d0 = _pX->_d0;
  int x_d1 = _pX->_d1;
  if (x_d0 != _w._d1) {
    x_d0 = _pX->_d0 * _pX->_d1 * _pX->_d2;
    x_d1 = _pX->_d3;
  }

  {
    // wG = x.T * yG
    int nJobs = _wG._d1 * _wG._d0;
    int numBlocks = (nJobs + K_THREAD_PER_BLOCK - 1) / K_THREAD_PER_BLOCK;
    dim3 block(numBlocks), threads(K_THREAD_PER_BLOCK);
    BackwardFc_Wg_Cuda<<<block, threads>>>(
        _wG._dataGpu, _pX->_dataGpu, yG->_dataGpu, x_d0, x_d1, _wG._d0, nJobs);
    assert(cudaGetLastError() == cudaSuccess);
  }
  {
    int numBlocks = (_b._d0 + K_THREAD_PER_BLOCK - 1) / K_THREAD_PER_BLOCK;
    dim3 block(numBlocks), threads(K_THREAD_PER_BLOCK);
    BackwardFc_Bg_Cuda<<<block, threads>>>(_bG._dataGpu, yG->_dataGpu, yG->_d0,
                                           yG->_d1);
    assert(cudaGetLastError() == cudaSuccess);
  }

  if (layerIndex > 0) {
    assert(yG->_d1 == x_d1);
    _xG.ResetTensor(x_d0, x_d1);
    // xG = yG * w.T
    int nJobs = _xG._d1 * _xG._d0;
    int numBlocks = (nJobs + K_THREAD_PER_BLOCK - 1) / K_THREAD_PER_BLOCK;
    dim3 block(numBlocks), threads(K_THREAD_PER_BLOCK);
    BackwardFc_Xg_Cuda<<<block, threads>>>(_xG._dataGpu, yG->_dataGpu,
                                           _w._dataGpu, yG->_d0, _w._d0,
                                           _xG._d0, nJobs);
    assert(cudaGetLastError() == cudaSuccess);
  }
  return &_xG;
}

__global__ void UpdateWs_Cuda(float learningRate, float beta1, float beta2,
                              float beta1t, float beta2t, float* w,
                              const float* wG, float* wG_m, float* wG_v,
                              int nJobs) {
  int index = blockIdx.x * blockDim.x + threadIdx.x;
  if (index >= nJobs) return;

  // Vanilla
  // w[index] -= wG[index] * learningRate;

  // Adam
  float g = wG[index];
  float m = wG_m[index];
  float v = wG_v[index];
  m = beta1 * m + (1.0f - beta1) * g;
  v = beta2 * v + (1.0f - beta2) * g * g;
  float unbiased_m = m / (1.0f - beta1t);
  float unbiased_v = v / (1.0f - beta2t);
  wG_m[index] = m;
  wG_v[index] = v;
  float impv = learningRate * unbiased_m / (sqrtf(unbiased_v) + 1e-8f);
  w[index] -= impv;
}

void FcLayer::SaveWeights(std::ofstream& file) {
  _w.WriteDataToFile(file);
  _b.WriteDataToFile(file);
}

void FcLayer::LoadWeights(FILE* file) {
  _w.LoadDataFromFile(file);
  _b.LoadDataFromFile(file);
}

void FcLayer::UpdateWs(float learningRate, float beta1, float beta2,
                       float beta1t, float beta2t) {
  {
    int nJobs = _w._d1 * _w._d0;
    int numBlocks = (nJobs + K_THREAD_PER_BLOCK - 1) / K_THREAD_PER_BLOCK;
    dim3 block(numBlocks), threads(K_THREAD_PER_BLOCK);
    UpdateWs_Cuda<<<block, threads>>>(learningRate, beta1, beta2, beta1t,
                                      beta2t, _w._dataGpu, _wG._dataGpu,
                                      _wG_m._dataGpu, _wG_v._dataGpu, nJobs);
    assert(cudaGetLastError() == cudaSuccess);
  }
  {
    int numBlocks = (_b._d0 + K_THREAD_PER_BLOCK - 1) / K_THREAD_PER_BLOCK;
    dim3 block(numBlocks), threads(K_THREAD_PER_BLOCK);
    UpdateWs_Cuda<<<block, threads>>>(learningRate, beta1, beta2, beta1t,
                                      beta2t, _b._dataGpu, _bG._dataGpu,
                                      _bG_m._dataGpu, _bG_v._dataGpu, _b._d0);
    assert(cudaGetLastError() == cudaSuccess);
  }
}

void FcLayer::Pull() {
  _xG.PullFromGpu();
  _w.PullFromGpu();
  _wG.PullFromGpu();
  _wG_m.PullFromGpu();
  _wG_v.PullFromGpu();
  _b.PullFromGpu();
  _bG.PullFromGpu();
  _bG_m.PullFromGpu();
  _bG_v.PullFromGpu();
  _y.PullFromGpu();
}

Conv2dLayer::Conv2dLayer(CudaNn* nn, int kernelSize, int nInChannel,
                         int nOutChannel, int stride, int padding)
    : CudaLayer(nn),
      _kernelSize(kernelSize),
      _stride(stride),
      _padding(padding) {
  _w.ResetTensor(_kernelSize, _kernelSize, nInChannel, nOutChannel);
  _w.SetRandom(1.0f / sqrtf(float(_kernelSize * _kernelSize * nInChannel)));
  _wG.ResetTensor(_kernelSize, _kernelSize, nInChannel, nOutChannel);
  _wG_m.ResetTensor(_kernelSize, _kernelSize, nInChannel, nOutChannel);
  _wG_m.SetZero();
  _wG_v.ResetTensor(_kernelSize, _kernelSize, nInChannel, nOutChannel);
  _wG_v.SetZero();
  _b.ResetTensor(nOutChannel);
  _b.SetRandom(1.0f / sqrtf(float(_kernelSize * _kernelSize * nInChannel)));
  _bG.ResetTensor(nOutChannel);
  _bG_m.ResetTensor(nOutChannel);
  _bG_m.SetZero();
  _bG_v.ResetTensor(nOutChannel);
  _bG_v.SetZero();
}

__global__ void ForwardConv2d_Cuda(float* y, const float* x, const float* w,
                                   const float* b, int nOutChannel, int nRowY,
                                   int nColY, int nInChannel, int nRowX,
                                   int nColX, int kernelSize, int stride,
                                   int padding, int nJobs) {
  int index = blockIdx.x * blockDim.x + threadIdx.x;
  if (index >= nJobs) return;

  const int yIndex = index;
  int image = index / (nOutChannel * nRowY * nColY);
  index -= image * (nOutChannel * nRowY * nColY);
  int outChannel = index / (nRowY * nColY);
  index -= outChannel * (nRowY * nColY);
  int rowY = index / nColY;
  int colY = index % nColY;

  float val = 0.0f;
  int startRowX = rowY * stride - padding;
  int startColX = colY * stride - padding;
  int startRowKernel = max(0, padding - (rowY * stride));
  int startColKernel = max(0, padding - (colY * stride));
  for (int rowX = startRowX + startRowKernel; rowX < startRowX + kernelSize;
       ++rowX) {
    if (rowX >= nRowX) break;
    for (int colX = startColX + startColKernel; colX < startColX + kernelSize;
         ++colX) {
      if (colX >= nColX) break;
      for (int inChannel = 0; inChannel < nInChannel; ++inChannel) {
        int xBaseIndex = image * (nInChannel * nRowX * nColX) +
                         inChannel * (nRowX * nColX) + rowX * nColX;
        int wBaseIndex = outChannel * (nInChannel * kernelSize * kernelSize) +
                         inChannel * (kernelSize * kernelSize) +
                         (rowX - startRowX) * kernelSize;
        val += x[xBaseIndex + colX] * w[wBaseIndex + (colX - startColX)];
      }
    }
  }
  y[yIndex] = val + b[outChannel];
}

const CudaTensor* Conv2dLayer::Forward(const CudaTensor* x) {
  assert(x->_d2 == _w._d2);  // Check # of input channels

  // (N - F) / stride + 1
  _pX = x;
  int nColY = (_pX->_d0 + _padding * 2 - _w._d0) / _stride + 1;
  int nRowY = (_pX->_d1 + _padding * 2 - _w._d1) / _stride + 1;

  // x x y x nOutChannel x nImages
  _y.ResetTensor(nColY, nRowY, _w._d3, _pX->_d3);

  int nJobs = _y._dataSize;
  int numBlocks = (nJobs + K_THREAD_PER_BLOCK - 1) / K_THREAD_PER_BLOCK;
  dim3 blocks(numBlocks), threads(K_THREAD_PER_BLOCK);
  ForwardConv2d_Cuda<<<blocks, threads>>>(
      _y._dataGpu, _pX->_dataGpu, _w._dataGpu, _b._dataGpu, _y._d2, _y._d1,
      _y._d0, _pX->_d2, _pX->_d1, _pX->_d0, _kernelSize, _stride, _padding,
      nJobs);
  assert(cudaGetLastError() == cudaSuccess);
  return &_y;
}

__global__ void BackwardConv2d_Wg_Cuda(float* wG, const float* x,
                                       const float* yG, int nOutChannel,
                                       int nInChannel, int nImages, int nRowY,
                                       int nColY, int nRowX, int nColX,
                                       int kernelSize, int stride, int padding,
                                       int nJobs) {
  int index = blockIdx.x * blockDim.x + threadIdx.x;
  if (index >= nJobs) return;

  const int wIndex = index;
  int outChannel = index / (nInChannel * kernelSize * kernelSize);
  index -= outChannel * (nInChannel * kernelSize * kernelSize);
  int inChannel = index / (kernelSize * kernelSize);
  index -= inChannel * (kernelSize * kernelSize);
  int rowW = index / kernelSize;
  int colW = index % kernelSize;

  float val = 0.0f;
  int startRowY = max(int(ceil((padding - rowW) / float(stride))), 0);
  int startColY = max(int(ceil((padding - colW) / float(stride))), 0);
  for (int image = 0; image < nImages; ++image) {
    int yBaseIndex =
        image * (nOutChannel * nRowY * nColY) + outChannel * (nRowY * nColY);
    int xBaseIndex =
        image * (nInChannel * nRowX * nColX) + inChannel * (nRowX * nColX);
    for (int rowY = startRowY; rowY < nRowY; ++rowY) {
      int rowX = rowY * stride - padding + rowW;
      if (rowX >= nRowX) break;
      for (int colY = startColY; colY < nColY; ++colY) {
        int colX = colY * stride - padding + colW;
        if (colX >= nColX) break;
        int yIndex = yBaseIndex + rowY * nColY + colY;
        int xIndex = xBaseIndex + rowX * nColX + colX;
        val += x[xIndex] * yG[yIndex];
      }
    }
  }
  // wG[wIndex] = val / nImages;
  wG[wIndex] = val;
}

__global__ void BackwardConv2d_Bg_Cuda(float* bG, const float* yG, int nImages,
                                       int nOutChannel, int nRowY, int nColY) {
  const int outChannel = blockIdx.x * blockDim.x + threadIdx.x;
  if (outChannel >= nOutChannel) return;

  float val = 0.0f;
  for (int image = 0; image < nImages; ++image) {
    for (int rowY = 0; rowY < nRowY; ++rowY) {
      int yBaseIndex = image * (nOutChannel * nRowY * nColY) +
                       outChannel * (nRowY * nColY) + rowY * nColY;
      for (int colY = 0; colY < nColY; ++colY) {
        val += yG[yBaseIndex + colY];
      }
    }
  }
  // bG[outChannel] = val / nImages;
  bG[outChannel] = val;
}

__global__ void BackwardConv2d_Xg_Cuda(float* xG, const float* w,
                                       const float* yG, int nImages,
                                       int nInChannel, int nRowX, int nColX,
                                       int nOutChannel, int nRowY, int nColY,
                                       int kernelSize, int stride, int padding,
                                       int nJobs) {
  int index = blockIdx.x * blockDim.x + threadIdx.x;
  if (index >= nJobs) return;

  int xIndex = index;
  int image = index / (nInChannel * nRowX * nColX);
  index -= image * (nInChannel * nRowX * nColX);
  int inChannel = index / (nRowX * nColX);
  index -= inChannel * (nRowX * nColX);
  int rowX = index / nColX;
  int colX = index % nColX;

  float val = 0.0f;
  // Note(dongwook): Iterate all y's which use the current x
  int startRowY =
      max(0, int(floor(float(rowX + padding - kernelSize) / stride)) + 1);
  int startColY =
      max(0, int(floor(float(colX + padding - kernelSize) / stride)) + 1);
  for (int rowY = startRowY; rowY < nRowY; ++rowY) {
    int upperX = rowY * stride - padding;
    if (rowX < upperX) break;
    int rowW = rowX - upperX;
    for (int colY = startColY; colY < nColY; ++colY) {
      int leftX = colY * stride - padding;
      if (colX < leftX) break;
      int colW = colX - leftX;
      for (int outChannel = 0; outChannel < nOutChannel; ++outChannel) {
        int yIndex = image * (nOutChannel * nRowY * nColY) +
                     outChannel * (nRowY * nColY) + rowY * nColY + colY;
        int wIndex = outChannel * (nInChannel * kernelSize * kernelSize) +
                     inChannel * (kernelSize * kernelSize) + rowW * kernelSize +
                     colW;
        val += w[wIndex] * yG[yIndex];
      }
    }
  }
  xG[xIndex] = val;
}

const CudaTensor* Conv2dLayer::Backward(const CudaTensor* yG,
                                        const int layerIndex) {
  assert(_y._dataSize == yG->_dataSize);
  _xG.ResetTensor(_pX->_d0, _pX->_d1, _pX->_d2, _pX->_d3);

  {
    assert(_y._d3 <= 256);
    int nJobs = _wG._dataSize;
    int numBlocks =
        (nJobs + K_SMALL_THREAD_PER_BLOCK - 1) / K_SMALL_THREAD_PER_BLOCK;
    dim3 blocks(numBlocks), threads(K_SMALL_THREAD_PER_BLOCK);
    BackwardConv2d_Wg_Cuda<<<blocks, threads>>>(
        _wG._dataGpu, _pX->_dataGpu, yG->_dataGpu, _wG._d3, _wG._d2, _y._d3,
        _y._d1, _y._d0, _pX->_d1, _pX->_d0, _kernelSize, _stride, _padding,
        nJobs);
    assert(cudaGetLastError() == cudaSuccess);
  }
  {
    assert(_bG._d0 == _y._d2);
    int nJobs = _bG._d0;
    int numBlocks = (nJobs + K_THREAD_PER_BLOCK - 1) / K_THREAD_PER_BLOCK;
    dim3 blocks(numBlocks), threads(K_THREAD_PER_BLOCK);
    BackwardConv2d_Bg_Cuda<<<blocks, threads>>>(_bG._dataGpu, yG->_dataGpu,
                                                _y._d3, _y._d2, _y._d1, _y._d0);
    assert(cudaGetLastError() == cudaSuccess);
  }
  if (layerIndex > 0) {
    int nJobs = _xG._dataSize;
    int numBlocks = (nJobs + K_THREAD_PER_BLOCK - 1) / K_THREAD_PER_BLOCK;
    dim3 blocks(numBlocks), threads(K_THREAD_PER_BLOCK);
    BackwardConv2d_Xg_Cuda<<<blocks, threads>>>(
        _xG._dataGpu, _w._dataGpu, yG->_dataGpu, _xG._d3, _xG._d2, _xG._d1,
        _xG._d0, _y._d2, _y._d1, _y._d0, _kernelSize, _stride, _padding, nJobs);
    assert(cudaGetLastError() == cudaSuccess);
  }
  return &_xG;
}

void Conv2dLayer::UpdateWs(float learningRate, float beta1, float beta2,
                           float beta1t, float beta2t) {
  {
    int nJobs = _w._dataSize;
    int numBlocks = (nJobs + K_THREAD_PER_BLOCK - 1) / K_THREAD_PER_BLOCK;
    dim3 block(numBlocks), threads(K_THREAD_PER_BLOCK);
    UpdateWs_Cuda<<<block, threads>>>(learningRate, beta1, beta2, beta1t,
                                      beta2t, _w._dataGpu, _wG._dataGpu,
                                      _wG_m._dataGpu, _wG_v._dataGpu, nJobs);
    assert(cudaGetLastError() == cudaSuccess);
  }
  {
    int numBlocks = (_b._d0 + K_THREAD_PER_BLOCK - 1) / K_THREAD_PER_BLOCK;
    dim3 block(numBlocks), threads(K_THREAD_PER_BLOCK);
    UpdateWs_Cuda<<<block, threads>>>(learningRate, beta1, beta2, beta1t,
                                      beta2t, _b._dataGpu, _bG._dataGpu,
                                      _bG_m._dataGpu, _bG_v._dataGpu, _b._d0);
    assert(cudaGetLastError() == cudaSuccess);
  }
}

void Conv2dLayer::SaveWeights(std::ofstream& file) {
  _w.WriteDataToFile(file);
  _b.WriteDataToFile(file);
}

void Conv2dLayer::LoadWeights(FILE* file) {
  _w.LoadDataFromFile(file);
  _b.LoadDataFromFile(file);
}

void Conv2dLayer::Pull() {
  _xG.PullFromGpu();
  _w.PullFromGpu();
  _wG.PullFromGpu();
  _wG_m.PullFromGpu();
  _wG_v.PullFromGpu();
  _b.PullFromGpu();
  _bG.PullFromGpu();
  _bG_m.PullFromGpu();
  _bG_v.PullFromGpu();
  _y.PullFromGpu();
}

__global__ void ForwardMaxPool_Cuda(float* y, float* maxIndex, const float* x,
                                    int nImages, int nInChannel, int nRow,
                                    int nCol) {
  int index = blockIdx.x * blockDim.x + threadIdx.x;
  int nJobs = nImages * nInChannel * nRow * nCol;
  if (nJobs <= index) return;

  const int yIndex = index;
  int image = index / (nInChannel * nRow * nCol);
  index -= image * (nInChannel * nRow * nCol);
  int inChannel = index / (nRow * nCol);
  index -= inChannel * (nRow * nCol);
  int row = index / nCol;
  int col = index % nCol;

  const int offset[4][4] = {0,
                            1,
                            nCol * 2,
                            nCol * 2 + 1,
                            1,
                            nCol * 2,
                            nCol * 2 + 1,
                            0,
                            nCol * 2,
                            nCol * 2 + 1,
                            0,
                            1,
                            nCol * 2 + 1,
                            0,
                            1,
                            nCol * 2};
  int s = nJobs % 4;
  int xIndex = image * (nInChannel * nRow * nCol * 4) +
               inChannel * (nRow * nCol * 4) + row * nCol * 4 + col * 2;
  float maxVal = x[xIndex + offset[s][0]], maxIdx = (float)offset[s][0];
  if (maxVal < x[xIndex + offset[s][1]]) {
    maxIdx = (float)offset[s][1];
    maxVal = x[xIndex + offset[s][1]];
  }
  if (maxVal < x[xIndex + offset[s][2]]) {
    maxIdx = (float)offset[s][2];
    maxVal = x[xIndex + offset[s][2]];
  }
  if (maxVal < x[xIndex + offset[s][3]]) {
    maxIdx = (float)offset[s][3];
    maxVal = x[xIndex + offset[s][3]];
  }
  maxIndex[yIndex] = maxIdx;
  y[yIndex] = maxVal;
}

const CudaTensor* MaxPoolLayer::Forward(const CudaTensor* x) {
  assert(0 == (x->_d0 % 2) && 0 == (x->_d1 % 2));

  _pX = x;
  _y.ResetTensor(x->_d0 / 2, x->_d1 / 2, x->_d2, x->_d3);
  _maxIndex.ResetTensor(_y._d0, _y._d1, _y._d2, _y._d3);

  int nJobs = _y._dataSize;
  int numBlocks = (nJobs + K_THREAD_PER_BLOCK - 1) / K_THREAD_PER_BLOCK;
  dim3 block(numBlocks), threads(K_THREAD_PER_BLOCK);
  ForwardMaxPool_Cuda<<<block, threads>>>(_y._dataGpu, _maxIndex._dataGpu,
                                          _pX->_dataGpu, _y._d3, _y._d2, _y._d1,
                                          _y._d0);
  assert(cudaGetLastError() == cudaSuccess);
  return &_y;
}

__global__ void BackwardMaxPool_Cuda(float* xG, const float* yG,
                                     const float* maxIndex, int nImages,
                                     int nOutChannel, int nRow, int nCol) {
  int index = blockIdx.x * blockDim.x + threadIdx.x;
  int nJobs = nImages * nOutChannel * nRow * nCol;
  if (nJobs <= index) return;

  int yGindex = index;
  int image = index / (nOutChannel * nRow * nCol);
  index -= image * (nOutChannel * nRow * nCol);
  int outChannel = index / (nRow * nCol);
  index -= outChannel * (nRow * nCol);
  int row = index / nCol;
  int col = index % nCol;

  int xGindex = image * (nOutChannel * nRow * nCol * 4) +
                outChannel * (nRow * nCol * 4) + row * nCol * 4 + col * 2;
  xG[xGindex] = 0.0f;
  xG[xGindex + 1] = 0.0f;
  xG[xGindex + nCol * 2] = 0.0f;
  xG[xGindex + nCol * 2 + 1] = 0.0f;
  xG[xGindex + (int)maxIndex[yGindex]] = yG[yGindex];
}

const CudaTensor* MaxPoolLayer::Backward(const CudaTensor* yG,
                                         const int layerIndex) {
  assert(_y._dataSize == yG->_dataSize);
  _xG.ResetTensor(2 * _y._d0, 2 * _y._d1, _y._d2, _y._d3);

  if (layerIndex > 0) {
    int nJobs = _y._dataSize;
    int numBlocks = (nJobs + K_THREAD_PER_BLOCK - 1) / K_THREAD_PER_BLOCK;
    dim3 block(numBlocks), threads(K_THREAD_PER_BLOCK);
    BackwardMaxPool_Cuda<<<block, threads>>>(_xG._dataGpu, yG->_dataGpu,
                                             _maxIndex._dataGpu, _y._d3, _y._d2,
                                             _y._d1, _y._d0);
    assert(cudaGetLastError() == cudaSuccess);
  }
  return &_xG;
}

void MaxPoolLayer::Pull() {
  _maxIndex.PullFromGpu();
  _xG.PullFromGpu();
  _y.PullFromGpu();
}

__global__ void ForwardRelu_Cuda(float* relu_x, const float* x, int nJobs) {
  int jobIndex = blockIdx.x * blockDim.x + threadIdx.x;
  if (jobIndex >= nJobs) return;

  relu_x[jobIndex] = fmaxf(x[jobIndex], 0.0f);
}

const CudaTensor* ReluLayer::Forward(const CudaTensor* x) {
  _pX = x;
  _xRelu.ResetTensor(_pX->_d0, _pX->_d1, _pX->_d2, _pX->_d3);
  int nJobs = _xRelu._dataSize;
  int numBlocks = (nJobs + K_THREAD_PER_BLOCK - 1) / K_THREAD_PER_BLOCK;
  dim3 block(numBlocks), threads(K_THREAD_PER_BLOCK);
  ForwardRelu_Cuda<<<block, threads>>>(_xRelu._dataGpu, _pX->_dataGpu, nJobs);
  assert(cudaGetLastError() == cudaSuccess);
  return &_xRelu;
}

__global__ void BackwardRelu_Cuda(float* xG, const float* yG, const float* x,
                                  int nJobs) {
  int jobIndex = blockIdx.x * blockDim.x + threadIdx.x;
  if (jobIndex >= nJobs) return;

  xG[jobIndex] = x[jobIndex] < 0.0f ? 0.0f : yG[jobIndex];
}

const CudaTensor* ReluLayer::Backward(const CudaTensor* yG,
                                      const int layerIndex) {
  if (layerIndex > 0) {
    _xG.ResetTensor(_xRelu._d0, _xRelu._d1, _xRelu._d2, _xRelu._d3);
    int nJobs = _xRelu._dataSize;
    int numBlocks = (nJobs + K_THREAD_PER_BLOCK - 1) / K_THREAD_PER_BLOCK;
    dim3 block(numBlocks), threads(K_THREAD_PER_BLOCK);
    BackwardRelu_Cuda<<<block, threads>>>(_xG._dataGpu, yG->_dataGpu,
                                          _pX->_dataGpu, nJobs);
    assert(cudaGetLastError() == cudaSuccess);
  }
  return &_xG;
}

void ReluLayer::Pull() {
  _xRelu.PullFromGpu();
  _xG.PullFromGpu();
}

BatchNorm2dLayer::BatchNorm2dLayer(CudaNn* nn, int inDim)
    : CudaLayer(nn), _pX(nullptr), _accCount(0) {
  _meanAndVariance.ResetTensor(2, inDim);
  _meanAndVarianceAcc.ResetTensor(2, inDim);
  _meanAndVarianceAcc.SetZero();
  _meanAndVarianceG.ResetTensor(2, inDim);
  _w.ResetTensor(2, inDim);
  for (int i = 0; i < inDim; ++i) {
    _w._data[i * 2] = 1.0f;      // alpha
    _w._data[i * 2 + 1] = 0.0f;  // beta
  }
  _w.PushToGpu();
  _wG.ResetTensor(2, inDim);
  _wG_m.ResetTensor(2, inDim);
  _wG_m.SetZero();
  _wG_v.ResetTensor(2, inDim);
  _wG_v.SetZero();
}

template <int BLOCK_SIZE>
__global__ void ForwardBatchNorm2d_Train_Cuda(float* meanAndVariance,
                                              float* meanAndVarianceAcc,
                                              float* xHat, float* y,
                                              const float* w, const float* x,
                                              int nRow, int nCol) {
  int nChannel = gridDim.x;
  int nImages = blockDim.x;
  int ch = blockIdx.x;
  int image = threadIdx.x;

  int mDash = nImages * nRow * nCol;
  int imageStride = nChannel * nRow * nCol;
  int channelStride = nRow * nCol;
  int currChBaseIndex = ch * channelStride;
  int baseIndex = image * imageStride + currChBaseIndex;

  __shared__ float meanArr[BLOCK_SIZE];
  meanArr[image] = 0.0f;
  for (int i = 0; i < channelStride; ++i) {
    meanArr[image] += x[baseIndex + i];
  }
  __syncthreads();
  float mean = 0.0f;
  for (int i = 0; i < nImages; ++i) {
    mean += meanArr[i];
  }
  mean /= mDash;

  __shared__ float varianceArr[BLOCK_SIZE];
  varianceArr[image] = 0.0f;
  for (int i = 0; i < channelStride; ++i) {
    float a = x[baseIndex + i] - mean;
    varianceArr[image] += (a * a);
  }
  __syncthreads();
  float variance = 0.0f;
  for (int i = 0; i < nImages; ++i) {
    variance += varianceArr[i];
  }
  variance /= mDash;

  float alpha = w[ch * 2 + 0];
  float beta = w[ch * 2 + 1];
  float d = sqrtf(variance) + 1e-8f;
  for (int i = 0; i < channelStride; ++i) {
    float xDash = (x[baseIndex + i] - mean) / d;
    xHat[baseIndex + i] = xDash;
    y[baseIndex + i] = alpha * xDash + beta;
  }

  if (threadIdx.x == 0) {
    meanAndVariance[ch * 2 + 0] = mean;
    meanAndVariance[ch * 2 + 1] = variance;
    meanAndVarianceAcc[ch * 2 + 0] += mean;
    meanAndVarianceAcc[ch * 2 + 1] += variance;
  }
}

__global__ void ForwardBatchNorm2d_Cuda(float* y,
                                        const float* meanAndVarianceAcc,
                                        int accCount, const float* w,
                                        const float* x, int nRow, int nCol) {
  int nChannel = gridDim.x;
  int ch = blockIdx.x;
  int image = threadIdx.x;

  int imageStride = nChannel * nRow * nCol;
  int channelStride = nRow * nCol;
  int currChBaseIndex = ch * channelStride;
  int baseIndex = image * imageStride + currChBaseIndex;

  float alpha = w[ch * 2 + 0];
  float beta = w[ch * 2 + 1];

#if 1  // Note(dongwook): Deteministic network
  float mean = meanAndVarianceAcc[ch * 2 + 0] / accCount;
  float variance = meanAndVarianceAcc[ch * 2 + 1] / (accCount - 1);
  // float variance = (meanAndVarianceAcc[ch * 2 + 1] / (accCount));
#else  // Note(dongwook): Non-deteministic network for batch input.
  int nImages = blockDim.x;
  int mDash = nImages * nRow * nCol;
  __shared__ float meanArr[100];
  meanArr[image] = 0.0f;
  for (int i = 0; i < channelStride; ++i) {
    meanArr[image] += x[baseIndex + i];
  }
  __syncthreads();
  float mean = 0.0f;
  for (int i = 0; i < nImages; ++i) {
    mean += meanArr[i];
  }
  mean /= mDash;

  __shared__ float varianceArr[100];
  varianceArr[image] = 0.0f;
  for (int i = 0; i < channelStride; ++i) {
    float a = x[baseIndex + i] - mean;
    varianceArr[image] += (a * a);
  }
  __syncthreads();
  float variance = 0.0f;
  for (int i = 0; i < nImages; ++i) {
    variance += varianceArr[i];
  }
  variance /= mDash;
#endif
  float d = 1.0f / (sqrtf(variance) + 1e-8f);
  float a = alpha * d;
  float b = beta - (alpha * mean) * d;
  for (int i = 0; i < channelStride; ++i) {
    y[baseIndex + i] = a * x[baseIndex + i] + b;
  }
}

const CudaTensor* BatchNorm2dLayer::Forward(const CudaTensor* x) {
  assert(_w._d1 == x->_d2);
  _pX = x;
  _xHat.ResetTensor(x->_d0, x->_d1, x->_d2, x->_d3);
  _y.ResetTensor(x->_d0, x->_d1, x->_d2, x->_d3);

  if (_nn->IsTraining()) {
    assert(x->_d3 <= 256);
    ForwardBatchNorm2d_Train_Cuda<256><<<x->_d2, x->_d3>>>(
        _meanAndVariance._dataGpu, _meanAndVarianceAcc._dataGpu, _xHat._dataGpu,
        _y._dataGpu, _w._dataGpu, x->_dataGpu, x->_d1, x->_d0);
    assert(cudaGetLastError() == cudaSuccess);
    ++_accCount;
  } else {
    assert(x->_d3 <= 256);
    ForwardBatchNorm2d_Cuda<<<x->_d2, x->_d3>>>(
        _y._dataGpu, _meanAndVarianceAcc._dataGpu,
        _accCount > 1 ? _accCount : 2, _w._dataGpu, x->_dataGpu, x->_d1,
        x->_d0);
    assert(cudaGetLastError() == cudaSuccess);
  }

  return &_y;
}

template <int BLOCK_SIZE>
__global__ void BackwardBatchNorm2d_Cuda(
    float* wG, float* xG, float* meanAndVarianceG, float* meanAndVarianceAcc,
    const float* w, const float* x, const float* xHat,
    const float* meanAndVariance, const float* yG, int nRow, int nCol) {
  int nChannel = gridDim.x;
  int nImages = blockDim.x;
  int ch = blockIdx.x;
  int image = threadIdx.x;

  int mDash = nImages * nRow * nCol;
  int imageStride = nChannel * nRow * nCol;
  int channelStride = nRow * nCol;
  int currChBaseIndex = ch * channelStride;

  __shared__ float alphaArr[BLOCK_SIZE], betaArr[BLOCK_SIZE];
  alphaArr[image] = 0.0f;
  betaArr[image] = 0.0f;
  int baseIndex = image * imageStride + currChBaseIndex;
  for (int i = 0; i < channelStride; ++i) {
    float currYg = yG[baseIndex + i];
    alphaArr[image] += currYg * xHat[baseIndex + i];
    betaArr[image] += currYg;
    xG[baseIndex + i] = currYg * w[ch * 2 + 0];
  }
  __syncthreads();
  float alpha = 0.0f, beta = 0.0f;
  for (int i = 0; i < nImages; ++i) {
    alpha += alphaArr[i];
    beta += betaArr[i];
  }
  wG[ch * 2 + 0] = alpha;
  wG[ch * 2 + 1] = beta;

  __shared__ float varianceGarr[BLOCK_SIZE];
  varianceGarr[image] = 0.0f;
  float mean = meanAndVariance[ch * 2 + 0];
  float variance = meanAndVariance[ch * 2 + 1];
  float b = -0.5f * __powf(variance + 1e-8f, -1.5f);
  for (int i = 0; i < channelStride; ++i) {
    varianceGarr[image] += (xG[baseIndex + i] * (x[baseIndex + i] - mean) * b);
  }
  __syncthreads();
  float varianceG = 0.0f;
  for (int i = 0; i < nImages; ++i) {
    varianceG += varianceGarr[i];
  }
  meanAndVarianceG[ch * 2 + 1] = varianceG;

  __shared__ float meanGarr[BLOCK_SIZE];
  meanGarr[image] = 0.0f;
  float varianceSqrtInv = 1.0f / (sqrtf(variance) + 1e-8f);
  b = -varianceSqrtInv;
  for (int i = 0; i < channelStride; ++i) {
    meanGarr[image] +=
        xG[baseIndex + i] * b +
        varianceG * ((-2.0f * (x[baseIndex + i] - mean)) / mDash);
  }
  __syncthreads();
  float meanG = 0.0f;
  for (int i = 0; i < nImages; ++i) {
    meanG += meanGarr[i];
  }
  meanAndVarianceG[ch * 2 + 0] = meanG;

  b = meanG / mDash;
  for (int i = 0; i < channelStride; ++i) {
    float currXg = xG[baseIndex + i];
    xG[baseIndex + i] =
        currXg * varianceSqrtInv +
        varianceG * ((2.0f * (x[baseIndex + i] - mean)) / mDash) + b;
  }

  // Reset
  meanAndVarianceAcc[ch * 2 + 0] = 0.0f;
  meanAndVarianceAcc[ch * 2 + 1] = 0.0f;
}

const CudaTensor* BatchNorm2dLayer::Backward(const CudaTensor* yG,
                                             const int layerIndex) {
  // _pX.shape == _xG.shape == yG.shape == _y.shape
  _xG.ResetTensor(_pX->_d0, _pX->_d1, _pX->_d2, _pX->_d3);

  assert(_xG._d3 <= 256);
  BackwardBatchNorm2d_Cuda<256><<<_xG._d2, _xG._d3>>>(
      _wG._dataGpu, _xG._dataGpu, _meanAndVarianceG._dataGpu,
      _meanAndVarianceAcc._dataGpu, _w._dataGpu, _pX->_dataGpu, _xHat._dataGpu,
      _meanAndVariance._dataGpu, yG->_dataGpu, _xG._d1, _xG._d0);
  assert(cudaGetLastError() == cudaSuccess);
  _accCount = 0;
  return &_xG;
}

void BatchNorm2dLayer::UpdateWs(float learningRate, float beta1, float beta2,
                                float beta1t, float beta2t) {
  {
    int nJobs = _w._d1 * _w._d0;
    int numBlocks = (nJobs + K_THREAD_PER_BLOCK - 1) / K_THREAD_PER_BLOCK;
    dim3 block(numBlocks), threads(K_THREAD_PER_BLOCK);
    UpdateWs_Cuda<<<block, threads>>>(learningRate, beta1, beta2, beta1t,
                                      beta2t, _w._dataGpu, _wG._dataGpu,
                                      _wG_m._dataGpu, _wG_v._dataGpu, nJobs);
    assert(cudaGetLastError() == cudaSuccess);
  }
}

void BatchNorm2dLayer::SaveWeights(std::ofstream& file) {
  _w.WriteDataToFile(file);
}

void BatchNorm2dLayer::LoadWeights(FILE* file) {
  _w.LoadDataFromFile(file);
}

void BatchNorm2dLayer::Pull() {
  _meanAndVariance.PullFromGpu();
  _meanAndVarianceAcc.PullFromGpu();
  _meanAndVarianceG.PullFromGpu();
  _w.PullFromGpu();
  _wG.PullFromGpu();
  _wG_m.PullFromGpu();
  _wG_v.PullFromGpu();
  _xG.PullFromGpu();
  _xHat.PullFromGpu();
  _y.PullFromGpu();
}

__global__ void Dropout_Cuda(float* x, const float* inputX,
                             const float* dropoutMask, int nJobs) {
  int index = blockIdx.x * blockDim.x + threadIdx.x;
  if (index >= nJobs) return;

  x[index] = inputX[index] * dropoutMask[index];
}

const CudaTensor* DropoutLayer::Forward(const CudaTensor* x) {
  if (false == _nn->IsTraining()) {
    return x;
  }

  _crossCheck = 1;
  _dropoutMask.ResetTensor(x->_dataSize);
  _dropoutMask.SetDropoutMask(_dropoutRate);
  _xDropped.ResetTensor(x->_d0, x->_d1, x->_d2, x->_d3);

  int nJobs = _xDropped._dataSize;
  int numBlocks = (nJobs + K_THREAD_PER_BLOCK - 1) / K_THREAD_PER_BLOCK;
  dim3 block(numBlocks), threads(K_THREAD_PER_BLOCK);
  Dropout_Cuda<<<block, threads>>>(_xDropped._dataGpu, x->_dataGpu,
                                   _dropoutMask._dataGpu, nJobs);
  assert(cudaGetLastError() == cudaSuccess);
  return &_xDropped;
}

const CudaTensor* DropoutLayer::Backward(const CudaTensor* yG,
                                         const int layerIndex) {
  assert(_dropoutMask._dataSize == yG->_dataSize);
  if (1 != _crossCheck) return yG;
  _crossCheck = 0;

  if (layerIndex > 0) {
    _yGdropped.ResetTensor(yG->_d0, yG->_d1, yG->_d2, yG->_d3);
    int nJobs = _yGdropped._dataSize;
    int numBlocks = (nJobs + K_THREAD_PER_BLOCK - 1) / K_THREAD_PER_BLOCK;
    dim3 block(numBlocks), threads(K_THREAD_PER_BLOCK);
    Dropout_Cuda<<<block, threads>>>(_yGdropped._dataGpu, yG->_dataGpu,
                                     _dropoutMask._dataGpu, nJobs);
    assert(cudaGetLastError() == cudaSuccess);
  }
  return &_yGdropped;
}

void DropoutLayer::Pull() {
  _dropoutMask.PullFromGpu();
  _xDropped.PullFromGpu();
  _yGdropped.PullFromGpu();
}

__global__ void ForwardQuatNorm_Cuda(float* y, const float* x, int nQuats,
                                     int nJobs) {
  int index = blockIdx.x * blockDim.x + threadIdx.x;
  if (index >= nJobs) return;

  int batch = index / nQuats;
  int elem = index % nQuats;
  int baseIndex = batch * nQuats * 4 + elem * 4;
  float q[4] = {x[baseIndex], x[baseIndex + 1], x[baseIndex + 2],
                x[baseIndex + 3]};
  float l =
      sqrtf(q[0] * q[0] + q[1] * q[1] + q[2] * q[2] + q[3] * q[3]) + 1e-8f;
  y[baseIndex + 0] = q[0] / l;
  y[baseIndex + 1] = q[1] / l;
  y[baseIndex + 2] = q[2] / l;
  y[baseIndex + 3] = q[3] / l;
}

const CudaTensor* QuatNormLayer::Forward(const CudaTensor* x) {
  assert(x->_d0 % 4 == 0);
  _pX = x;
  _y.ResetTensor(x->_d0, x->_d1, x->_d2, x->_d3);

  int nJobs = _pX->_d0 * _pX->_d1 / 4;
  int nBlocks =
      (nJobs + K_SMALL_THREAD_PER_BLOCK - 1) / K_SMALL_THREAD_PER_BLOCK;
  dim3 block(nBlocks), threads(K_SMALL_THREAD_PER_BLOCK);
  ForwardQuatNorm_Cuda<<<block, threads>>>(_y._dataGpu, _pX->_dataGpu,
                                           _pX->_d1 / 4, nJobs);
  assert(cudaGetLastError() == cudaSuccess);

  return &_y;
}

__global__ void BackwardQuatNorm_Cuda(float* xG, const float* x,
                                      const float* yG, int nQuats, int nJobs) {
  int index = blockIdx.x * blockDim.x + threadIdx.x;
  if (index >= nJobs) return;

  int batch = index / nQuats;
  int elem = index % nQuats;
  int baseIndex = batch * nQuats * 4 + elem * 4;
  float q[4] = {x[baseIndex], x[baseIndex + 1], x[baseIndex + 2],
                x[baseIndex + 3]};
  float tYg[4] = {yG[baseIndex], yG[baseIndex + 1], yG[baseIndex + 2],
                  yG[baseIndex + 3]};
  float squaredSum = q[0] * q[0] + q[1] * q[1] + q[2] * q[2] + q[3] * q[3];
  float a = powf(squaredSum + 1e-8f, -1.5f);
  for (int i = 0; i < 4; ++i) {
    float b = squaredSum * tYg[i];
    float c = 0.0f;
    for (int j = 0; j < 4; ++j) {
      c += (q[i] * q[j] * tYg[j]);
    }
    xG[baseIndex + i] = a * (b - c);
  }
}

const CudaTensor* QuatNormLayer::Backward(const CudaTensor* yG,
                                          const int layerIndex) {
  assert(yG->_dataSize == _pX->_dataSize);
  _xG.ResetTensor(_pX->_d0, _pX->_d1, _pX->_d2, _pX->_d3);

  if (layerIndex > 0) {
    int nJobs = _xG._d0 * _xG._d1 / 4;
    int nBlocks =
        (nJobs + K_SMALL_THREAD_PER_BLOCK - 1) / K_SMALL_THREAD_PER_BLOCK;
    dim3 block(nBlocks), threads(K_SMALL_THREAD_PER_BLOCK);
    BackwardQuatNorm_Cuda<<<block, threads>>>(_xG._dataGpu, _pX->_dataGpu,
                                              yG->_dataGpu, _xG._d0 / 4, nJobs);
    assert(cudaGetLastError() == cudaSuccess);
  }

  return &_xG;
}

void QuatNormLayer::Pull() {
  _y.PullFromGpu();
  _xG.PullFromGpu();
}

__global__ void ForwardSoftmax_Cuda(float* softmax, const float* x, int nRow,
                                    int nCol) {
  int r = blockIdx.x * blockDim.x + threadIdx.x;
  if (nRow <= r) return;

  int baseIndex = r * nCol;
  float maxValue = x[baseIndex];
  for (int i = 1; i < nCol; ++i) {
    maxValue = max(maxValue, x[baseIndex + i]);
  }
  float sum = 0.0f;
  for (int i = 0; i < nCol; ++i) {
    sum += expf(x[baseIndex + i] - maxValue);
  }
  for (int i = 0; i < nCol; ++i) {
    softmax[baseIndex + i] = expf(x[baseIndex + i] - maxValue) / sum;
  }
}

const CudaTensor* SoftmaxLayer::Forward(const CudaTensor* x) {
  assert(1 == x->_d2 && 1 == x->_d3);
  _softmax.ResetTensor(x->_d0, x->_d1);

  int nBlocks = (x->_d1 + K_THREAD_PER_BLOCK - 1) / K_THREAD_PER_BLOCK;
  dim3 block(nBlocks), threads(K_THREAD_PER_BLOCK);
  ForwardSoftmax_Cuda<<<block, threads>>>(_softmax._dataGpu, x->_dataGpu,
                                          x->_d1, x->_d0);
  assert(cudaGetLastError() == cudaSuccess);
  return &_softmax;
}

__global__ void BackwardSoftmax_Cuda(float* lossG, const float* softmax,
                                     const float* yLabel, int nCol, int nJobs) {
  int index = blockIdx.x * blockDim.x + threadIdx.x;
  if (index >= nJobs) return;
  int r = index / nCol;
  int c = index % nCol;

  lossG[index] = softmax[index];
  if (yLabel[r] == c) lossG[index] -= 1.0f;
}

const CudaTensor* SoftmaxLayer::Backward(const CudaTensor* yG,
                                         const int layerIndex) {
  assert(yG->_d0 == yG->_dataSize && yG->_d0 == _softmax._d1);
  _lossG.ResetTensor(_softmax._d0, _softmax._d1);
  if (layerIndex > 0) {
    int nJobs = _lossG._d1 * _lossG._d0;
    int nBlocks = (nJobs + K_THREAD_PER_BLOCK - 1) / K_THREAD_PER_BLOCK;
    dim3 block(nBlocks), threads(K_THREAD_PER_BLOCK);
    BackwardSoftmax_Cuda<<<block, threads>>>(_lossG._dataGpu, _softmax._dataGpu,
                                             yG->_dataGpu, _lossG._d0, nJobs);
    assert(cudaGetLastError() == cudaSuccess);
  }
  return &_lossG;
}

void SoftmaxLayer::Pull() {
  _softmax.PullFromGpu();
  _lossG.PullFromGpu();
}

__global__ void BackwardSumOfSqures(float* yG, const float* y,
                                    const float* yLabel, int nJobs) {
  int index = blockIdx.x * blockDim.x + threadIdx.x;
  if (index >= nJobs) return;

  float diff = y[index] - yLabel[index];
  yG[index] = 2.0f * diff;
}

const CudaTensor* SumOfSquaresLayer::Forward(const CudaTensor* x) {
  _pY = x;
  return _pY;
}

const CudaTensor* SumOfSquaresLayer::Backward(const CudaTensor* yLabel,
                                              const int layerIndex) {
  assert(yLabel->_dataSize == _pY->_dataSize);
  _yG.ResetTensor(yLabel->_d0, yLabel->_d1, yLabel->_d2, yLabel->_d3);

  if (layerIndex > 0) {
    int nJobs = _yG._dataSize;
    int nBlocks = (nJobs + K_THREAD_PER_BLOCK - 1) / K_THREAD_PER_BLOCK;
    dim3 block(nBlocks), threads(K_THREAD_PER_BLOCK);
    BackwardSumOfSqures<<<block, threads>>>(_yG._dataGpu, _pY->_dataGpu,
                                            yLabel->_dataGpu, nJobs);
    assert(cudaGetLastError() == cudaSuccess);
  }
  return &_yG;
}

void SumOfSquaresLayer::Pull() { _yG.PullFromGpu(); }

///////////////////////////////////////////////////////////////////////
CudaNn::CudaNn() : _beta1t(kBeta1), _beta2t(kBeta2), _train(false) {}

CudaNn::~CudaNn() { InitializeCudaNn(""); }

bool CudaNn::InitializeCudaNn(const char* desc) {
  size_t numLayers = _layers.size();
  for (size_t i = 0; i < numLayers; ++i) {
    delete _layers[i];
  }
  _layers.clear();
  return true;
}

bool CudaNn::AddConv2d(int kernelSize, int nInChannel, int nOutChannel,
                       int stride, int padding) {
  _layers.push_back(new Conv2dLayer(this, kernelSize, nInChannel, nOutChannel,
                                    stride, padding));
  return true;
}

bool CudaNn::AddMaxPool() {
  _layers.push_back(new MaxPoolLayer(this));
  return true;
}

bool CudaNn::AddFc(int inDim, int outDim) {
  _layers.push_back(new FcLayer(this, inDim, outDim));
  return true;
}

bool CudaNn::AddRelu() {
  _layers.push_back(new ReluLayer(this));
  return true;
}

bool CudaNn::AddBatchNorm2d(int inDim) {
  _layers.push_back(new BatchNorm2dLayer(this, inDim));
  return true;
}

bool CudaNn::AddDropout(float dropoutRatio) {
  assert(dropoutRatio > 0.0 && dropoutRatio < 1.0);
  _layers.push_back(new DropoutLayer(this, dropoutRatio));
  return true;
}

bool CudaNn::AddSoftmax() {
  _layers.push_back(new SoftmaxLayer(this));
  return true;
}

bool CudaNn::AddSumOfSquares() {
  _layers.push_back(new SumOfSquaresLayer(this));
  return true;
}

bool CudaNn::AddQuatNorm() {
  _layers.push_back(new QuatNormLayer(this));
  return true;
}

const CudaTensor* CudaNn::Forward(const CudaTensor* x, bool train) {
  _train = train;
  const CudaTensor* y = nullptr;
  size_t numLayer = _layers.size();
  for (size_t i = 0; i < numLayer; ++i) {
    if (nullptr == x) return nullptr;

    // [afterdusk] Uncomment to set precision of layers
    // switch (i) {
    //   case 0:  // Conv2d(3, 1, 4, 1, 1)
    //   case 1:  // Relu
    //   case 2:  // MaxPool
    //     setenv("VF_SIGNIFICAND", "8", 1);
    //     setenv("VF_EXPONENT_MIN", "-132", 1);
    //     setenv("VF_EXPONENT_MAX", "128", 1);
    //     break;
    //   case 3:  // Conv2d(3, 4, 8, 1, 1)
    //   case 4:  // Relu
    //   case 5:  // MaxPool
    //     setenv("VF_SIGNIFICAND", "4", 1);
    //     setenv("VF_EXPONENT_MIN", "-128", 1);
    //     setenv("VF_EXPONENT_MAX", "128", 1);
    //     break;
    //   case 6:  // Conv2d(3, 8, 16, 1, 1)
    //   case 7:  // Relu
    //     setenv("VF_SIGNIFICAND", "3", 1);
    //     setenv("VF_EXPONENT_MIN", "-7", 1);
    //     setenv("VF_EXPONENT_MAX", "8", 1);
    //     break;
    //   case 8:   // Fc(7 * 7 * 16, 1000)
    //   case 9:   // Relu
    //   case 10:  // Fc(1000, 10)
    //   case 11:  // Softmax
    //     setenv("VF_SIGNIFICAND", "2", 1);
    //     setenv("VF_EXPONENT_MIN", "-6", 1);
    //     setenv("VF_EXPONENT_MAX", "8", 1);
    //     break;
    // }

    y = _layers[i]->Forward(x);
    x = y;
  }
  return y;
}

void CudaNn::Backward(const CudaTensor* yLabel) {
  const CudaTensor* y = yLabel;
  const CudaTensor* yGradient = nullptr;
  int numLayer = (int)_layers.size();
  for (int i = 0; i < numLayer; ++i) {
    int layerIndex = numLayer - i - 1;
    yGradient = _layers[layerIndex]->Backward(y, layerIndex);
    y = yGradient;
  }
}

void CudaNn::UpdateWs(float learningRate) {
  size_t numLayer = _layers.size();
  for (size_t i = 0; i < numLayer; ++i) {
    _layers[i]->UpdateWs(learningRate, kBeta1, kBeta2, _beta1t, _beta2t);
  }
  _beta1t *= kBeta1;
  _beta2t *= kBeta2;
}

void CudaNn::Pull() {
  size_t numLayer = _layers.size();
  for (size_t i = 0; i < numLayer; ++i) {
    _layers[i]->Pull();
  }
}

void CudaNn::SaveWeights(const char* weights_file) {
  std::ofstream file(weights_file);
  printf("Starting save\n");
  size_t numLayer = _layers.size();
  for (size_t i = 0; i < numLayer; ++i) {
    _layers[i]->SaveWeights(file);
  }
}

void CudaNn::LoadWeights(const char* weights_file) {
  FILE* file = fopen(weights_file, "r");
  printf("Starting load\n");
  size_t numLayer = _layers.size();
  for (size_t i = 0; i < numLayer; ++i) {
    _layers[i]->LoadWeights(file);
  }
}

}  // namespace ff