// TODO(ataei): reduce the apparent redundancy of all the code below.
#include "caffe2/operators/pool_op.h"
#include "caffe2/utils/cpu_neon.h"

namespace caffe2 {

using std::max;
using std::min;

namespace {

template <typename T>
class AveragePool {
 public:
  static float initialize() {
    return 0.0;
  }

  static void process(
      const int x_col,
      const int y_col,
      ConstEigenMatrixMap<float>& x_mat,
      EigenMatrixMap<float>& y_mat) {
    y_mat.col(y_col) += x_mat.col(x_col);
  }

  static void process(const T& x_data, T& y_data) {
    y_data += x_data;
  }

  static void finalize(const int size, T& y_data) {
    y_data /= size;
  }

  static void
  finalize(const int size, const int col, EigenMatrixMap<float>& y_mat) {
    y_mat.col(col) /= size;
  }

  static bool runNeon() {
    return true;
  }
};

template <typename T>
class MaxPool {
 public:
  static float initialize() {
    return std::numeric_limits<float>::lowest();
  }

  static void process(
      const int x_col,
      const int y_col,
      ConstEigenMatrixMap<float>& x_mat,
      EigenMatrixMap<float>& y_mat) {
    y_mat.col(y_col) = y_mat.col(y_col).cwiseMax(x_mat.col(x_col));
  }

  static void process(const T& x_data, T& y_data) {
    if (x_data > y_data) {
      y_data = x_data;
    }
  }

  static void finalize(const int /*size*/, T& /*y_data*/) {}

  static void finalize(
      const int /*size*/,
      const int col,
      EigenMatrixMap<float>& /*y_mat*/) {}

  static bool runNeon() {
    return false;
  }
};

#ifdef __ARM_NEON__

bool isNeonEligible(int inputH, int inputW,
                    int outputH, int outputW,
                    int kH, int kW,
                    int strideH, int strideW,
                    int padT, int padL, int padB, int padR,
                    int dilationH, int dilationW,
                    const float* input,
                    float* output) {
  // Use this kernel only if:
  // Kernel width is 4x4
  // Kernel stride is 4x4
  // Padding is 0
  // Dilation is 1
  // Output width and height are even divisors of input width
  // Input width and height are divisible by 4 (should be implied by
  // all of the above, but just check again)
  // Input and output pointers are aligned by float32x4_t

  bool kernelOk = (kH == 4) && (kW == 4);
  bool strideOk = (strideH == 4) && (strideW == 4);
  bool padOk = (padT == 0) && (padL == 0) && (padB == 0) && (padR == 0);
  bool dilationOk = (dilationH == 1) && (dilationW == 1);

  bool outputOk = ((inputH % outputH) == 0) && ((inputW % outputW) == 0);
  bool inputOk = (inputW % 4 == 0) && (inputH % 4 == 0);
  bool alignOk = isPointerAligned(input, sizeof(float32x4_t)) &&
    isPointerAligned(output, sizeof(float32x4_t));

  return kernelOk && strideOk && padOk && dilationOk &&
    outputOk && inputOk && alignOk;
}

// Vectorizes 4x4p0s0 averge pooling for ARM NEON
void avgPoolNeon4x4p0s0Plane(int inputH, int inputW,
                             const float* input,
                             float* output) {
  constexpr int kKernelHeight = 4;
  constexpr int kKernelWidth = 4;
  constexpr float kDiv =
    (1.0f / ((float) kKernelHeight * (float) kKernelWidth));

  // Handle portion that can be unrolled by 4
  constexpr int kUnroll = 4;
  constexpr int kLoadSizeFloat = (sizeof(float32x4_t) / sizeof(float));
  constexpr int kLoadCols = kUnroll * kLoadSizeFloat;

  if (inputW % kLoadCols == 0) {
    //
    // Manually unroll by 4 (kUnroll)
    //

    for (int h = 0; h < inputH; h += kKernelHeight) {
      float* outputRow = output + (h / kKernelHeight) * (inputW / kKernelWidth);
      const float* curInput = input + h * inputW;

      for (int w = 0; w < inputW; w += kLoadCols) {
        float32x4_t out = {};

        {
          float32x4_t v0_0 = vld1q_f32_aligned(curInput + 0 * inputW);
          float32x4_t v0_1 = vld1q_f32_aligned(curInput + 1 * inputW);
          float32x4_t v0_2 = vld1q_f32_aligned(curInput + 2 * inputW);
          float32x4_t v0_3 = vld1q_f32_aligned(curInput + 3 * inputW);
          float v0 = horizontal_sum_f32(v0_0, v0_1, v0_2, v0_3);
          out = vsetq_lane_f32(v0, out, 0);
        }
        curInput += kLoadSizeFloat;

        {
          float32x4_t v0_0 = vld1q_f32_aligned(curInput + 0 * inputW);
          float32x4_t v0_1 = vld1q_f32_aligned(curInput + 1 * inputW);
          float32x4_t v0_2 = vld1q_f32_aligned(curInput + 2 * inputW);
          float32x4_t v0_3 = vld1q_f32_aligned(curInput + 3 * inputW);
          float v0 = horizontal_sum_f32(v0_0, v0_1, v0_2, v0_3);
          out = vsetq_lane_f32(v0, out, 1);
        }
        curInput += kLoadSizeFloat;

        {
          float32x4_t v0_0 = vld1q_f32_aligned(curInput + 0 * inputW);
          float32x4_t v0_1 = vld1q_f32_aligned(curInput + 1 * inputW);
          float32x4_t v0_2 = vld1q_f32_aligned(curInput + 2 * inputW);
          float32x4_t v0_3 = vld1q_f32_aligned(curInput + 3 * inputW);
          float v0 = horizontal_sum_f32(v0_0, v0_1, v0_2, v0_3);
          out = vsetq_lane_f32(v0, out, 2);
        }
        curInput += kLoadSizeFloat;

        {
          float32x4_t v0_0 = vld1q_f32_aligned(curInput + 0 * inputW);
          float32x4_t v0_1 = vld1q_f32_aligned(curInput + 1 * inputW);
          float32x4_t v0_2 = vld1q_f32_aligned(curInput + 2 * inputW);
          float32x4_t v0_3 = vld1q_f32_aligned(curInput + 3 * inputW);
          float v0 = horizontal_sum_f32(v0_0, v0_1, v0_2, v0_3);
          out = vsetq_lane_f32(v0, out, 3);
        }
        curInput += kLoadSizeFloat;

        out = vmulq_f32(out, vdupq_n_f32(kDiv));
        vst1q_f32_aligned(&outputRow[w / kKernelWidth], out);
      }
    }
  } else {
    //
    // Not unrolled
    //

    for (int h = 0; h < inputH; h += kKernelHeight) {
      const float* inputRow = input + h * inputW;
      float* outputRow = output + (h / kKernelHeight) * (inputW / kKernelWidth);

      for (int w = 0; w < inputW; w += kKernelWidth) {
        const float* curInput = inputRow + w;

        float32x4_t v0_0 = vld1q_f32_aligned(curInput + 0 * inputW);
        float32x4_t v0_1 = vld1q_f32_aligned(curInput + 1 * inputW);
        float32x4_t v0_2 = vld1q_f32_aligned(curInput + 2 * inputW);
        float32x4_t v0_3 = vld1q_f32_aligned(curInput + 3 * inputW);
        float v0 = horizontal_sum_f32(v0_0, v0_1, v0_2, v0_3) * kDiv;
        outputRow[w / kKernelWidth] = v0;
      }
    }
  }
}

void
runNeonAveragePool4x4p0s0NCHW(int N, int C, int inputH, int inputW,
                              const float* input,
                              float* output) {
  // We only have the 4x4p0s0 implementation at present, which is
  // checked at a higher level
  int outputH = inputH / 4;
  int outputW = inputW / 4;

  for (int n = 0; n < N; ++n) {
    for (int c = 0; c < C; ++c) {
      const float* curInput = input + (n * C + c) * inputH * inputW;
      float* curOutput = output + (n * C + c) * outputH * outputW;

      avgPoolNeon4x4p0s0Plane(inputH, inputW, curInput, curOutput);
    }
  }
}
#endif // __ARM_NEON__

}  // namespace

template <typename T, class Context, typename PoolType>
bool PoolOp<T, Context, PoolType>::RunOnDeviceWithOrderNCHW() {
  auto& X = Input(0);
  auto* Y = Output(0);
  ConvPoolOpBase<Context>::SetOutputSize(X, Y, X.dim32(1));

  const float* Xdata = X.template data<float>();
  float* Ydata = Y->template mutable_data<float>();
  math::Set<float, CPUContext>(
      Y->size(), PoolType::initialize(), Ydata, &context_);
  // The main loop
  int channels = X.dim32(1);
  int height = X.dim32(2);
  int width = kernel_.size() > 1 ? X.dim32(3) : 1;
  int depth = kernel_.size() > 2 ? X.dim32(4) : 1;
  int pooled_height = Y->dim32(2);
  int pooled_width = kernel_.size() > 1 ? Y->dim32(3) : 1;
  int pooled_depth = kernel_.size() > 2 ? Y->dim32(4) : 1;

#ifdef __ARM_NEON__
  // We specialize certain variants on ARM for vectorization
  if (PoolType::runNeon() && isNeonEligible(
                                 X.dim32(2),
                                 X.dim32(3),
                                 Y->dim32(2),
                                 Y->dim32(3),
                                 kernel_h(),
                                 kernel_w(),
                                 stride_h(),
                                 stride_w(),
                                 pad_t(),
                                 pad_l(),
                                 pad_b(),
                                 pad_r(),
                                 dilation_h(),
                                 dilation_w(),
                                 Xdata,
                                 Ydata)) {
    runNeonAveragePool4x4p0s0NCHW(
        X.dim32(0), X.dim32(1), X.dim32(2), X.dim32(3), Xdata, Ydata);
    return true;
  }
#endif // __ARM_NEON__
  switch (kernel_.size()) {
    case 1:
      for (int n = 0; n < X.dim32(0); ++n) {
        for (int c = 0; c < channels; ++c) {
          for (int ph = 0; ph < pooled_height; ++ph) {
            int hstart = ph * stride_h() - pad_t();
            int hend = min(hstart + kernel_h(), height);
            hstart = max(hstart, 0);
            for (int h = hstart; h < hend; ++h) {
              PoolType::process(Xdata[h], Ydata[ph]);
            }
            PoolType::finalize(hend - hstart, Ydata[ph]);
          }
          // Do offset.
          Xdata += height;
          Ydata += pooled_height;
        }
      }
      break;
    case 2:
      for (int n = 0; n < X.dim32(0); ++n) {
        for (int c = 0; c < channels; ++c) {
          for (int ph = 0; ph < pooled_height; ++ph) {
            int hstart = ph * stride_h() - pad_t();
            int hend = min(hstart + kernel_h(), height);
            hstart = max(hstart, 0);
            for (int pw = 0; pw < pooled_width; ++pw) {
              int wstart = pw * stride_w() - pad_l();
              int wend = min(wstart + kernel_w(), width);
              wstart = max(wstart, 0);
              const int pool_index = ph * pooled_width + pw;
              for (int h = hstart; h < hend; ++h) {
                for (int w = wstart; w < wend; ++w) {
                  const int input_index = h * width + w;
                  PoolType::process(Xdata[input_index], Ydata[pool_index]);
                }
              }
              PoolType::finalize(
                  (hend - hstart) * (wend - wstart), Ydata[pool_index]);
            }
          }
          // Do offset.
          Xdata += height * width;
          Ydata += pooled_height * pooled_width;
        }
      }
      break;
    case 3:
      for (int n = 0; n < X.dim32(0); ++n) {
        for (int c = 0; c < channels; ++c) {
          for (int ph = 0; ph < pooled_height; ++ph) {
            int hstart = ph * stride_h() - pad_t();
            int hend = min(hstart + kernel_h(), height);
            hstart = max(hstart, 0);
            for (int pw = 0; pw < pooled_width; ++pw) {
              int wstart = pw * stride_w() - pad_l();
              int wend = min(wstart + kernel_w(), width);
              wstart = max(wstart, 0);
              for (int pd = 0; pd < pooled_depth; ++pd) {
                int dstart = pd * stride_[2] - pads_[2];
                int dend = min(dstart + kernel_[2], depth);
                dstart = max(dstart, 0);
                const int pool_index =
                    ph * pooled_width * pooled_depth + pw * pooled_depth + pd;
                for (int h = hstart; h < hend; ++h) {
                  for (int w = wstart; w < wend; ++w) {
                    for (int d = dstart; d < dend; ++d) {
                      const int input_index = h * width * depth + w * depth + d;
                      PoolType::process(Xdata[input_index], Ydata[pool_index]);
                    }
                  }
                }
                PoolType::finalize(
                    (hend - hstart) * (wend - wstart) * (dend - dstart),
                    Ydata[pool_index]);
              }
            }
          }
          // Do offset.
          Xdata += height * width * depth;
          Ydata += pooled_height * pooled_width * pooled_depth;
        }
      }
      break;
    default:
      CAFFE_THROW("Unsupported pooling size : ", kernel_.size());
      return false;
  }
  return true;
}

template <typename T, class Context, typename PoolType>
bool PoolOp<T, Context, PoolType>::RunOnDeviceWithOrderNHWC() {
  auto& X = Input(0);
  auto* Y = Output(0);
  int height = X.dim32(1);
  int width = kernel_.size() > 1 ? X.dim32(2) : 1;
  int depth = kernel_.size() > 2 ? X.dim32(3) : 1;
  int channels = X.dim32(X.ndim() - 1);
  ConvPoolOpBase<Context>::SetOutputSize(X, Y, channels);

  EigenMatrixMap<float> Ymat(
      Y->template mutable_data<float>(), channels, Y->size() / channels);
  ConstEigenMatrixMap<float> Xmat(
      X.template data<float>(), channels, X.size() / channels);
  int pooled_height = Y->dim32(1);
  int pooled_width = kernel_.size() > 1 ? Y->dim32(2) : 1;
  int pooled_depth = kernel_.size() > 2 ? Y->dim32(3) : 1;
  // The main loop
  switch (kernel_.size()) {
    case 1:
      for (int n = 0; n < X.dim32(0); ++n) {
        for (int ph = 0; ph < pooled_height; ++ph) {
          int hstart = ph * stride_h() - pad_t();
          int hend = min(hstart + kernel_h(), height);
          hstart = max(hstart, 0);
          const int y_col = n * pooled_height + ph;
          Ymat.col(y_col).setConstant(PoolType::initialize());
          for (int h = hstart; h < hend; ++h) {
            const int x_col = n * height + h;
            PoolType::process(x_col, y_col, Xmat, Ymat);
          }
          PoolType::finalize((hend - hstart), y_col, Ymat);
        }
      }
      break;
    case 2:
      for (int n = 0; n < X.dim32(0); ++n) {
        for (int ph = 0; ph < pooled_height; ++ph) {
          int hstart = ph * stride_h() - pad_t();
          int hend = min(hstart + kernel_h(), height);
          hstart = max(hstart, 0);
          for (int pw = 0; pw < pooled_width; ++pw) {
            int wstart = pw * stride_w() - pad_l();
            int wend = min(wstart + kernel_w(), width);
            wstart = max(wstart, 0);
            const int y_col = (n * pooled_height + ph) * pooled_width + pw;
            Ymat.col(y_col).setConstant(PoolType::initialize());
            for (int h = hstart; h < hend; ++h) {
              for (int w = wstart; w < wend; ++w) {
                const int x_col = (n * height + h) * width + w;
                PoolType::process(x_col, y_col, Xmat, Ymat);
              }
            }
            PoolType::finalize((hend - hstart) * (wend - wstart), y_col, Ymat);
          }
        }
      }
      break;
    case 3:
      for (int n = 0; n < X.dim32(0); ++n) {
        for (int ph = 0; ph < pooled_height; ++ph) {
          int hstart = ph * stride_h() - pad_t();
          int hend = min(hstart + kernel_h(), height);
          hstart = max(hstart, 0);
          for (int pw = 0; pw < pooled_width; ++pw) {
            int wstart = pw * stride_w() - pad_l();
            int wend = min(wstart + kernel_w(), width);
            wstart = max(wstart, 0);
            for (int pd = 0; pd < pooled_depth; ++pd) {
              int dstart = pd * stride_[2] - pads_[2];
              int dend = min(dstart + kernel_[2], depth);
              dstart = max(dstart, 0);
              const int y_col = ((n * pooled_height + ph) * pooled_width + pw) *
                      pooled_depth +
                  pd;
              Ymat.col(y_col).setConstant(PoolType::initialize());
              for (int h = hstart; h < hend; ++h) {
                for (int w = wstart; w < wend; ++w) {
                  for (int d = dstart; d < dend; ++d) {
                    const int x_col =
                        ((n * height + h) * width + w) * depth + d;
                    PoolType::process(x_col, y_col, Xmat, Ymat);
                  }
                }
              }
              PoolType::finalize(
                  (hend - hstart) * (wend - wstart) * (dend - dstart),
                  y_col,
                  Ymat);
            }
          }
        }
    }
    break;
    default:
      CAFFE_THROW("Unsupported pooling size : ", kernel_.size());
      return false;
  }
  return true;
}

namespace {
REGISTER_CPU_OPERATOR(
    AveragePool,
    PoolOp<float, CPUContext, AveragePool<float>>);

OPERATOR_SCHEMA(AveragePool)
  .NumInputs(1)
  .NumOutputs(1)
  .TensorInferenceFunction(ConvPoolOpBase<CPUContext>::TensorInferenceForPool)
  .SetDoc(R"DOC(
AveragePool consumes an input blob X and applies average pooling across the
the blob according to kernel sizes, stride sizes, and pad lengths defined by the
ConvPoolOpBase operator. Average pooling consisting of averaging all values of a
subset of the input tensor according to the kernel size and downsampling the
data into the output blob Y for further processing.
  )DOC")
  .Input(0, "X", "Input data tensor from the previous operator; dimensions "
  "depend on whether the NCHW or NHWC operators are being used. For example, "
  "in the former, the input has size (N x C x H x W), where N is the batch "
  "size, C is the number of channels, and H and W are the height and the width "
  "of the data. The corresponding permutation of dimensions is used in the "
  "latter case. ")
  .Output(0, "Y", "Output data tensor from average pooling across the input "
  "tensor. Dimensions will vary based on various kernel, stride, and pad "
  "sizes.");

REGISTER_CPU_OPERATOR(MaxPool, PoolOp<float, CPUContext, MaxPool<float>>);

OPERATOR_SCHEMA(MaxPool)
  .NumInputs(1)
  .NumOutputs(1)
  .TensorInferenceFunction(ConvPoolOpBase<CPUContext>::TensorInferenceForPool)
  .SetDoc(R"DOC(
MaxPool consumes an input blob X and applies max pooling across the
the blob according to kernel sizes, stride sizes, and pad lengths defined by the
ConvPoolOpBase operator. Max pooling consisting of taking the maximumvalue of a
subset of the input tensor according to the kernel size and downsampling the
data into the output blob Y for further processing.
  )DOC")
  .Input(0, "X", "Input data tensor from the previous operator; dimensions "
  "depend on whether the NCHW or NHWC operators are being used. For example, "
  "in the former, the input has size (N x C x H x W), where N is the batch "
  "size, C is the number of channels, and H and W are the height and the width "
  "of the data. The corresponding permutation of dimensions is used in the "
  "latter case. ")
  .Output(0, "Y", "Output data tensor from max pooling across the input "
  "tensor. Dimensions will vary based on various kernel, stride, and pad "
  "sizes.");

}  // namespace
}  // namespace caffe2
