// conv_transpose_op_impl.h is the templated implementation of the
// conv_transpose_op.h file.
#ifndef CAFFE2_OPERATORS_CONV_TRANSPOSE_MOBILE_OP_IMPL_H_
#define CAFFE2_OPERATORS_CONV_TRANSPOSE_MOBILE_OP_IMPL_H_

#include "caffe2/core/common.h"

#ifdef C10_MOBILE

#include "caffe2/core/logging.h"
#include "caffe2/operators/conv_op_shared.h"
#include "caffe2/operators/conv_transpose_op_mobile.h"
#include "caffe2/utils/cpu_neon.h"
#include "caffe2/utils/eigen_utils.h"
#include "caffe2/utils/fixed_divisor.h"
#include "caffe2/utils/math.h"
#include "caffe2/utils/math/utils.h"

C10_DECLARE_bool(caffe2_force_shared_col_buffer);

namespace caffe2 {

template <typename T, typename Context>
void runTileContiguous(
    int tileId,
    int N,
    int M,
    int H,
    int W,
    int outputH,
    int outputW,
    int C,
    int kernelH,
    int kernelW,
    int strideH,
    int strideW,
    int padT,
    const T* filterData,
    const T* Xdata,
    T* colBufferData,
    T* Ydata,
    Context* context) {
  // The tile size is exactly the length of a single row
  int tileSize = W;

  auto kernelDataSize = C * kernelH * kernelW;
  auto currentTileStart = tileSize * tileId;

  // gemm tile
  math::GemmEx<T, Context>(
      CblasTrans,
      CblasNoTrans,
      kernelDataSize,
      tileSize,
      M,
      1,
      filterData,
      kernelDataSize,
      Xdata + currentTileStart,
      H * W,
      0,
      colBufferData,
      tileSize,
      context);

  // col2im tile
  // We assume that there is no padding in the columns (padL and padR
  // == 0).
  // FIXME: it is actually possible for us to handle padding, figure
  // out how to adjust the bounds

  // We write into Y in a de-interleaved fashion; in other words,
  // every column (mod strideW) == 0 together in one block,
  // every column (mod strideW) == 1 in another,
  // ... and so on.
  int colBlockSize = (W + kernelW / strideW);
  int numColBlocks = strideW;

  for (int c = 0; c < kernelDataSize; ++c) {
    int w_offset = c % kernelW;
    int h_offset = (c / kernelW) % kernelH;
    int c_im = c / kernelH / kernelW;

    // Each row is a separate tile that we handle. First determine the
    // row into which we are writing the output.
    // We can properly handle padding for the rows.
    int rowY = tileId * strideH - padT + h_offset;

    // If this row is out of bounds, then skip it
    if (!math::utils::IsAGeZeroAndALtB(rowY, outputH)) {
      continue;
    }

    // FIXME: we don't actually handle a dynamic padL > 0
    constexpr int kPadL = 0;
    int colOffsetStart = -kPadL + w_offset;
    int colBlockY = colOffsetStart % strideW;

    // However, within a block we may not start writing at offset
    // 0. The offset at which we begin writing is determined by
    // colOffsetStart
    int colWithinBlockOffsetY = colOffsetStart / strideW;

    // So, this is where we begin reading/writing in Y
    int colY = colBlockY * colBlockSize + colWithinBlockOffsetY;

    // This is the complete offset into Y from the start
    // Each row has strideW blocks of size colBlockSize
    int offsetY = rowY * colBlockSize * numColBlocks + colY;

    T* colBufferPointer = colBufferData + c * tileSize;
    T* yPointer =
        Ydata + c_im * outputH * (colBlockSize * numColBlocks) + offsetY;

    int b = 0;
#if defined(__ARM_NEON__) || defined(__ARM_NEON)
    // We vectorize the loop within the row
    {
      constexpr int kUnroll = (sizeof(float32x4_t) / sizeof(float)) * 4;
      int limit = (tileSize / kUnroll) * kUnroll;

      for (; b < limit; b += kUnroll) {
        float32x4_t cb0 = vld1q_f32(colBufferPointer + 0);
        float32x4_t cb1 = vld1q_f32(colBufferPointer + 4);
        float32x4_t cb2 = vld1q_f32(colBufferPointer + 8);
        float32x4_t cb3 = vld1q_f32(colBufferPointer + 12);

        float32x4_t y0 = vld1q_f32(yPointer + 0);
        float32x4_t y1 = vld1q_f32(yPointer + 4);
        float32x4_t y2 = vld1q_f32(yPointer + 8);
        float32x4_t y3 = vld1q_f32(yPointer + 12);

        y0 = vaddq_f32(y0, cb0);
        y1 = vaddq_f32(y1, cb1);
        y2 = vaddq_f32(y2, cb2);
        y3 = vaddq_f32(y3, cb3);

        vst1q_f32(yPointer + 0, y0);
        vst1q_f32(yPointer + 4, y1);
        vst1q_f32(yPointer + 8, y2);
        vst1q_f32(yPointer + 12, y3);

        colBufferPointer += kUnroll;
        yPointer += kUnroll;
      }
    }

    {
      constexpr int kUnroll = (sizeof(float32x4_t) / sizeof(float));
      int limit = (tileSize / kUnroll) * kUnroll;

      for (; b < limit; b += kUnroll) {
        float32x4_t cb0 = vld1q_f32(colBufferPointer);
        float32x4_t y0 = vld1q_f32(yPointer);

        y0 = vaddq_f32(y0, cb0);

        vst1q_f32(yPointer, y0);

        colBufferPointer += kUnroll;
        yPointer += kUnroll;
      }
    }
#endif

    // Handle un-vectorizable epilogue
    for (; b < tileSize; ++b) {
      *yPointer += *colBufferPointer;
      ++yPointer;
      ++colBufferPointer;
    }
  }
}

template <typename T, int N>
struct StoreInterleaved {};

template <>
struct StoreInterleaved<float, 1> {
#if defined(__ARM_NEON__) || defined(__ARM_NEON)
  inline static void store(float* p, float32x4_t v[1]) {
    vst1q_f32(p, v[0]);
  }
#endif

  inline static void store(float* p, float v[1]) {
    p[0] = v[0];
  }
};

template <>
struct StoreInterleaved<float, 2> {
#if defined(__ARM_NEON__) || defined(__ARM_NEON)
  inline static void store(float* p, float32x4_t v[2]) {
    float32x4x2_t x = {{v[0], v[1]}};
    vst2q_f32(p, x);
  }
#endif

  inline static void store(float* p, float v[2]) {
    p[0] = v[0];
    p[1] = v[1];
  }
};

template <>
struct StoreInterleaved<float, 3> {
#if defined(__ARM_NEON__) || defined(__ARM_NEON)
  inline static void store(float* p, float32x4_t v[3]) {
    float32x4x3_t x = {{v[0], v[1], v[2]}};
    vst3q_f32(p, x);
  }
#endif

  inline static void store(float* p, float v[3]) {
    p[0] = v[0];
    p[1] = v[1];
    p[2] = v[2];
  }
};

template <>
struct StoreInterleaved<float, 4> {
#if defined(__ARM_NEON__) || defined(__ARM_NEON)
  inline static void store(float* p, float32x4_t v[4]) {
    float32x4x4_t x = {{v[0], v[1], v[2], v[3]}};
    vst4q_f32(p, x);
  }
#endif

  inline static void store(float* p, float v[4]) {
    p[0] = v[0];
    p[1] = v[1];
    p[2] = v[2];
    p[3] = v[3];
  }
};

template <int kStrideW>
void reinterleaveRows(
    const float* src,
    const float* bias,
    int c,
    int h,
    float* dst,
    int outputC,
    int outputH,
    int outputW,
    int inputW,
    int kernelW,
    int strideW,
    int adjH) {
  // Each row in src is of the form:
  // [w mod strideW == 0 elements]...[w mod strideW == strideW - 1
  // elements]
  // We need to re-interleave the values and write them in the output
  int colBlockSize = inputW + kernelW / kStrideW;
  int noAdjOutputW = (inputW - 1) * kStrideW + kernelW;

  int point = c * outputH + h;
  src += point * colBlockSize * kStrideW;
  dst += point * outputW;

  float b = bias ? bias[c] : 0;
#if defined(__ARM_NEON__) || defined(__ARM_NEON)
  float32x4_t biasV = vdupq_n_f32(b);
#endif

  int w = 0;
#if defined(__ARM_NEON__) || defined(__ARM_NEON)
  constexpr int kUnroll = (sizeof(float32x4_t) / sizeof(float)) * 2;
  int limit = ((inputW - 1) / kUnroll) * kUnroll;

  for (; w < limit; w += kUnroll) {
    // We need to interleave in terms of kStrideW units
    float32x4_t v0[kStrideW];
    float32x4_t v1[kStrideW];

    for (int i = 0; i < kStrideW; ++i) {
      v0[i] = vld1q_f32(src + i * colBlockSize);
      v1[i] = vld1q_f32(src + i * colBlockSize + 4);
    }

    // add per-channel bias
    for (int i = 0; i < kStrideW; ++i) {
      v0[i] = vaddq_f32(v0[i], biasV);
      v1[i] = vaddq_f32(v1[i], biasV);
    }

    // Write interleaved into the output
    StoreInterleaved<float, kStrideW>::store(dst + 0 * kStrideW, v0);
    StoreInterleaved<float, kStrideW>::store(dst + 4 * kStrideW, v1);

    src += kUnroll;
    dst += kUnroll * kStrideW;
  }
#endif

  // Handle non-vectorizable remainder
  for (; w < inputW - 1; ++w) {
    float v[kStrideW];

    for (int i = 0; i < kStrideW; ++i) {
      v[i] = src[i * colBlockSize];
    }

    // add per-channel bias
    for (int i = 0; i < kStrideW; ++i) {
      v[i] += b;
    }

    // Write interleaved into the output
    StoreInterleaved<float, kStrideW>::store(dst, v);

    src += 1;
    dst += kStrideW;
  }

  // We have handled 0 .. (inputW - 1) * stride inclusive so far.
  // Handle the remainder
  int outputPoint = (inputW - 1) * kStrideW;
  int block = 0;

  // Output width may include adjustment into which we don't
  // write; ignore it
  while (outputPoint < noAdjOutputW) {
    float v = src[block * colBlockSize];
    dst[0] = v + b;
    ++outputPoint;
    dst += 1;

    ++block;
    if (block >= kStrideW) {
      block = 0;
      src += 1;
    }
  }

  // Remainder of the buffer comprised of just the `adj` must have
  // bias added
  for (; outputPoint < outputW; ++outputPoint) {
    dst[0] = b;
    dst += 1;
  }
}

template <int N, typename T, typename Context>
void reinterleaveMultithreaded(
    const T* y0,
    const T* bias_data,
    T* y,
    int outputC,
    int outputH,
    int outputW,
    int inputW,
    int kernelW,
    int strideW,
    int adjH,
    ThreadPool* pool) {
  // # channels times height
  size_t totalTiles = (size_t)outputC * outputH;
  FixedDivisor<int> divOutputH(outputH);

#define REINTERLEAVE(N)  \
  do {                   \
    reinterleaveRows<N>( \
        y0,              \
        bias_data,       \
        c,               \
        h,               \
        y,               \
        outputC,         \
        outputH,         \
        outputW,         \
        inputW,          \
        kernelW,         \
        strideW,         \
        adjH);           \
  } while (false)

  std::function<void(int, size_t)> fnReinterleave = [&](int threadId,
                                                        size_t tileId) {
    int h;
    int c;
    divOutputH.DivMod((int)tileId, &c, &h);

    REINTERLEAVE(N);
  };

#undef REINTERLEAVE

  pool->run(fnReinterleave, totalTiles);
}

#if defined(__ARM_NEON__) || defined(__ARM_NEON)
template <int N>
struct SumMultiple {
  static void sumInto(float* acc, float** toSum, size_t size);
};

template <>
struct SumMultiple<1> {
  static void sumInto(float* acc, float** toSum, size_t size) {
    constexpr int kUnroll = (sizeof(float32x4_t) / sizeof(float));
    int limit = (size / kUnroll) * kUnroll;

    auto toSum0 = toSum[0];

    size_t i = 0;
    for (; i < limit; i += kUnroll) {
      float32x4_t v0 = vld1q_f32_aligned(acc + i);
      float32x4_t v1 = vld1q_f32_aligned(toSum0 + i);

      v0 = vaddq_f32(v0, v1);

      vst1q_f32_aligned(acc + i, v0);
    }

    for (; i < size; ++i) {
      float v0 = acc[i];
      float v1 = toSum0[i];

      v0 += v1;

      acc[i] = v0;
    }
  }
};

template <>
struct SumMultiple<2> {
  static void sumInto(float* acc, float** toSum, size_t size) {
    constexpr int kUnroll = (sizeof(float32x4_t) / sizeof(float));
    int limit = (size / kUnroll) * kUnroll;

    auto toSum0 = toSum[0];
    auto toSum1 = toSum[1];

    size_t i = 0;
    for (; i < limit; i += kUnroll) {
      float32x4_t v0 = vld1q_f32_aligned(acc + i);
      float32x4_t v1 = vld1q_f32_aligned(toSum0 + i);
      float32x4_t v2 = vld1q_f32_aligned(toSum1 + i);

      v0 = vaddq_f32(v0, v1);
      v0 = vaddq_f32(v0, v2);

      vst1q_f32_aligned(acc + i, v0);
    }

    for (; i < size; ++i) {
      float v0 = acc[i];
      float v1 = toSum0[i];
      float v2 = toSum1[i];

      v0 += v1;
      v0 += v2;

      acc[i] = v0;
    }
  }
};

template <>
struct SumMultiple<3> {
  static void sumInto(float* acc, float** toSum, size_t size) {
    constexpr int kUnroll = (sizeof(float32x4_t) / sizeof(float));
    int limit = (size / kUnroll) * kUnroll;

    auto toSum0 = toSum[0];
    auto toSum1 = toSum[1];
    auto toSum2 = toSum[2];

    size_t i = 0;
    for (; i < limit; i += kUnroll) {
      float32x4_t v0 = vld1q_f32_aligned(acc + i);
      float32x4_t v1 = vld1q_f32_aligned(toSum0 + i);
      float32x4_t v2 = vld1q_f32_aligned(toSum1 + i);
      float32x4_t v3 = vld1q_f32_aligned(toSum2 + i);

      v0 = vaddq_f32(v0, v1);
      v2 = vaddq_f32(v2, v3);
      v0 = vaddq_f32(v0, v2);

      vst1q_f32_aligned(acc + i, v0);
    }

    for (; i < size; ++i) {
      float v0 = acc[i];
      float v1 = toSum0[i];
      float v2 = toSum1[i];
      float v3 = toSum2[i];

      v0 += v1;
      v2 += v3;
      v0 += v2;

      acc[i] = v0;
    }
  }
};
#endif

// Performs acc[i] += sum_j toSum_j[i] pointwise
void sumInto(float* acc, std::vector<float*>& toSum, size_t size) {
#if defined(__ARM_NEON__) || defined(__ARM_NEON)
  if (toSum.size() == 1) {
    SumMultiple<1>::sumInto(acc, toSum.data(), size);
    return;
  } else if (toSum.size() == 2) {
    SumMultiple<2>::sumInto(acc, toSum.data(), size);
    return;
  } else if (toSum.size() == 3) {
    SumMultiple<3>::sumInto(acc, toSum.data(), size);
    return;
  }
#endif

  // Otherwise, use fallback implementation
  EigenVectorArrayMap<float> accT(acc, size);

  for (auto p : toSum) {
    accT += ConstEigenVectorArrayMap<float>(p, size);
  }
}

template <typename T, class Context>
bool ConvTransposeMobileOp<T, Context>::RunOnDeviceWithOrderNCHW() {
  const Tensor& X = Input(INPUT);
  auto& filter = Input(FILTER);
  const int N = X.dim32(0), M = X.dim32(1), H = X.dim32(2), W = X.dim32(3);
  CAFFE_ENFORCE(filter.ndim() == 4, "filter must be 4D tensor");
  CAFFE_ENFORCE(
      filter.dim32(0) == M,
      "filter number must be equal to input channel number");
  const int C = filter.dim32(1);
  CAFFE_ENFORCE(
      filter.dim32(2) == this->kernel_h(),
      "filter height must be equal to kernel height");
  CAFFE_ENFORCE(
      filter.dim32(3) == this->kernel_w(),
      "filter width must be equal to kernel width");
  if (InputSize() == 3) {
    auto& bias = Input(BIAS);
    CAFFE_ENFORCE(bias.ndim() == 1, "bias must be 1D tensor");
    CAFFE_ENFORCE(
        bias.dim32(0) == C,
        "bias dimension must be equal to output channel number");
  }

  if (X.numel() == 0) {
    VLOG(2) << "Number of elements is 0 in ConvTrasposeOp";
    return true;
  }

  auto sizes = ConvTransposeUnpoolBase<Context>::GetOutputSize(X, C);
  Tensor* Y = Output(0, sizes, at::dtype<T>());

  const int outputH = Y->dim32(2);
  const int outputW = Y->dim32(3);
  const int outputPlaneSize = outputH * outputW;
  const int outputBatchElementSize = Y->dim32(1) * outputPlaneSize;

  auto Xdata = X.template data<T>();
  auto Ydata = Y->template mutable_data<T>();

  auto pool = ws_->GetThreadPool();
  auto numThreads = pool->getNumThreads();

  // Initialize per-thread buffers for output
  // The main thread will write directly into the output Y, we just
  // need buffers for the worker threads
  size_t colBlockSize = W + this->kernel_w() / this->stride_w();
  size_t threadYBufferSize = C * outputH * colBlockSize * this->stride_w();
  // Require 16 byte alignment, so 4-element alignment as these are floats.
  size_t threadYBufferSizeAligned =
      ((C * outputH * colBlockSize * this->stride_w() + 3) / 4) * 4;
  size_t threadColBufferSize = C * this->kernel_h() * this->kernel_w() * W;

  // Work around GCC 4.9 bug when this is declared inside the inner lambda.
  auto runLocalTile = [&](TensorCPU* threadBuffer,
                          int threadId,
                          size_t tileId) {
    auto localYData = threadBuffer->template mutable_data<T>() +
        threadId * threadYBufferSizeAligned;

    auto localColBufferData = threadBuffer->template mutable_data<T>() +
        numThreads * threadYBufferSizeAligned + threadId * threadColBufferSize;

    runTileContiguous<T, Context>(
        tileId,
        N,
        M,
        H,
        W,
        outputH,
        outputW,
        C,
        this->kernel_h(),
        this->kernel_w(),
        this->stride_h(),
        this->stride_w(),
        this->pad_t(),
        filter.template data<T>(),
        Xdata,
        localColBufferData,
        localYData,
        &context_);
  };

  auto f = [&](Tensor* threadBuffer) {
    threadBuffer->Resize(
        numThreads * threadYBufferSizeAligned +
        numThreads * threadColBufferSize);
    // Group together thread buffers for accumulation
    std::vector<T*> toSum(numThreads - 1);
    for (int i = 1; i < numThreads; ++i) {
      toSum[i - 1] = threadBuffer->template mutable_data<T>() +
          i * threadYBufferSizeAligned;
    }

    for (auto image_id = 0; image_id < N; ++image_id) {
      // Each time through, we have to reset all per-thread output
      // buffers, since the output buffer is only per-batch element
      // The column buffers are overwritten by the matrix multiplication
      // each time, so we need not clear them out each round
      math::Set<T, Context>(
          numThreads * threadYBufferSizeAligned,
          0,
          threadBuffer->template mutable_data<T>(),
          &context_);

      // Run tiled gemm and col2im in our threadpool; all of these tiles
      // are guaranteed to be full tiles
      // Each tile handles a single row of the input
      pool->run(
          [&](int threadId, int tileId) {
            runLocalTile(threadBuffer, threadId, tileId);
          },
          H);

      // We need to accumulate the per-thread results into the output
      // Y; the first worker thread (main thread) already produced its
      // results in Y
      sumInto(
          threadBuffer->template mutable_data<T>(), toSum, threadYBufferSize);

// y0 now contains the final output, but it is in deinterleaved
// form. We have to re-interleave it to produce the final form in Y
// This operation also handles adding the per-channel bias.
#define REINTERLEAVE(N)                                              \
  do {                                                               \
    reinterleaveMultithreaded<N, T, Context>(                        \
        threadBuffer->template mutable_data<T>(),                    \
        InputSize() == 3 ? Input(BIAS).template data<T>() : nullptr, \
        Ydata,                                                       \
        Y->dim32(1),                                                 \
        Y->dim32(2),                                                 \
        Y->dim32(3),                                                 \
        W,                                                           \
        this->kernel_w(),                                            \
        this->stride_w(),                                            \
        this->adj_h(),                                               \
        pool);                                                       \
  } while (false)

      if (this->stride_w() == 1) {
        REINTERLEAVE(1);
      } else if (this->stride_w() == 2) {
        REINTERLEAVE(2);
      } else if (this->stride_w() == 3) {
        REINTERLEAVE(3);
      } else if (this->stride_w() == 4) {
        REINTERLEAVE(4);
      }

#undef REINTERLEAVE

      Xdata += M * H * W;
      Ydata += Y->size() / Y->dim32(0);
    }
  };
  if (FLAGS_caffe2_force_shared_col_buffer || shared_buffer_) {
    runWithSharedBuffer<Context>(ws_, f);
  } else {
    f(&threadBuffer_);
  }

  return true;
}

template <typename T, class Context>
bool ConvTransposeMobileOp<T, Context>::RunOnDeviceWithOrderNHWC() {
  CAFFE_THROW("Not implemented.");
}

} // namespace caffe2

#endif // C10_MOBILE

#endif // CAFFE2_OPERATORS_CONV_TRANSPOSE_MOBILE_OP_IMPL_H_