<commit_before>#include "caffe2/operators/prelu_op.h"

#include "caffe2/utils/cpu_neon.h"
#include "caffe2/utils/math.h"

namespace caffe2 {

#ifdef __ARM_NEON__
namespace {

void runNeonPrelu(float* out, const float* in, int size, float w) {
  float32x4_t vZero = vdupq_n_f32(0.0f);
  float32x4_t vW = vdupq_n_f32(w);

  constexpr int kVecSizeInFloat = sizeof(float32x4_t) / sizeof(float);

  if (size < kVecSizeInFloat) {
    for (int i = 0; i < size; ++i) {
      float v = in[i];
      out[i] = v > 0 ? v : v * w;
    }

    return;
  }

  // We want to load aligned from the input, but assume the output is unaligned
  int prologue =
    kVecSizeInFloat -
    // remainder in floats
    (((uintptr_t) in) % (sizeof(float32x4_t))) / sizeof(float);

  int i = 0;

  // Prologue loop
  for (; i < prologue; ++i) {
    float v = in[i];
    out[i] = v > 0 ? v : v * w;
  }

  // The loop is manually unrolled by 6; seems to be the limit for
  // armv7 to avoid register spills
  constexpr int kUnroll = 6;
  constexpr int kFloatsPerLoop = kUnroll * kVecSizeInFloat;

  int remainder = size - prologue;
  int vectorizable = prologue + (remainder / kFloatsPerLoop) * kFloatsPerLoop;

  for (; i < vectorizable; i += kFloatsPerLoop) {
    float32x4_t v0 = vld1q_f32_aligned(in + i + 0);
    float32x4_t v1 = vld1q_f32_aligned(in + i + 4);
    float32x4_t v2 = vld1q_f32_aligned(in + i + 8);
    float32x4_t v3 = vld1q_f32_aligned(in + i + 12);
    float32x4_t v4 = vld1q_f32_aligned(in + i + 16);
    float32x4_t v5 = vld1q_f32_aligned(in + i + 20);

    uint32x4_t gz0 = vcgtq_f32(v0, vZero);
    uint32x4_t gz1 = vcgtq_f32(v1, vZero);
    uint32x4_t gz2 = vcgtq_f32(v2, vZero);
    uint32x4_t gz3 = vcgtq_f32(v3, vZero);
    uint32x4_t gz4 = vcgtq_f32(v4, vZero);
    uint32x4_t gz5 = vcgtq_f32(v5, vZero);

    float32x4_t v0neg = vmulq_f32(v0, vW);
    float32x4_t v1neg = vmulq_f32(v1, vW);
    float32x4_t v2neg = vmulq_f32(v2, vW);
    float32x4_t v3neg = vmulq_f32(v3, vW);
    float32x4_t v4neg = vmulq_f32(v4, vW);
    float32x4_t v5neg = vmulq_f32(v5, vW);

    // v0 > 0 ? v0 : v0 * w
    v0 = vbslq_f32(gz0, v0, v0neg);
    v1 = vbslq_f32(gz1, v1, v1neg);
    v2 = vbslq_f32(gz2, v2, v2neg);
    v3 = vbslq_f32(gz3, v3, v3neg);
    v4 = vbslq_f32(gz4, v4, v4neg);
    v5 = vbslq_f32(gz5, v5, v5neg);

    vst1q_f32(out + i + 0, v0);
    vst1q_f32(out + i + 4, v1);
    vst1q_f32(out + i + 8, v2);
    vst1q_f32(out + i + 12, v3);
    vst1q_f32(out + i + 16, v4);
    vst1q_f32(out + i + 20, v5);
  }

  for (; i < size; ++i) {
    float v = in[i];
    out[i] = v > 0 ? v : v * w;
  }
}

}
#endif // __ARM_NEON__

template <>
bool PReluOp<float, CPUContext>::RunOnDevice() {
  const auto& X = Input(0);
  const auto& W = Input(1);
  auto* Y = Output(0);
  Y->ResizeLike(X);
  const auto* Xdata = X.template data<float>();
  const auto* Wdata = W.template data<float>();
  auto* Ydata = Y->template mutable_data<float>();

  const auto C = order_ == StorageOrder::NCHW ? X.dim(1) : X.dim(X.ndim() - 1);
  const auto C_shared = (W.size() == 1);

  if (!C_shared) {
    CAFFE_ENFORCE_EQ(C, W.size());
  }

  if (C_shared) {
#ifdef __ARM_NEON__
    // The function is completely pointwise
    runNeonPrelu(Ydata, Xdata, X.size(), Wdata[0]);
#else
    ConstEigenVectorMap<float> Xvec(Xdata, X.size());
    EigenVectorMap<float> Yvec(Ydata, Y->size());
    Yvec = Xvec.cwiseMax(0.f) + Xvec.cwiseMin(0.f) * Wdata[0];
    return true;
#endif // __ARM_NEON__
  }

  // non-shared case.
  switch (order_) {
    case StorageOrder::NCHW: {
      const auto N = X.dim(0);
      const auto dim = X.size_from_dim(2);

#ifdef __ARM_NEON__
      // Pointwise for each channel
      for (int n = 0; n < N; ++n) {
        for (int c = 0; c < C; ++c) {
          runNeonPrelu(Ydata + (n * C + c) * dim,
                       Xdata + (n * C + c) * dim,
                       dim, Wdata[c]);
        }
      }
#else
      int nc = 0;
      for (int n = 0; n < N; ++n) {<commit_msg>Fix ARM_NEON codepath for non-shared scale case.<commit_after>#include "caffe2/operators/prelu_op.h"

#include "caffe2/utils/cpu_neon.h"
#include "caffe2/utils/math.h"

namespace caffe2 {

#ifdef __ARM_NEON__
namespace {

void runNeonPrelu(float* out, const float* in, int size, float w) {
  float32x4_t vZero = vdupq_n_f32(0.0f);
  float32x4_t vW = vdupq_n_f32(w);

  constexpr int kVecSizeInFloat = sizeof(float32x4_t) / sizeof(float);

  if (size < kVecSizeInFloat) {
    for (int i = 0; i < size; ++i) {
      float v = in[i];
      out[i] = v > 0 ? v : v * w;
    }

    return;
  }

  // We want to load aligned from the input, but assume the output is unaligned
  int prologue =
    kVecSizeInFloat -
    // remainder in floats
    (((uintptr_t) in) % (sizeof(float32x4_t))) / sizeof(float);

  int i = 0;

  // Prologue loop
  for (; i < prologue; ++i) {
    float v = in[i];
    out[i] = v > 0 ? v : v * w;
  }

  // The loop is manually unrolled by 6; seems to be the limit for
  // armv7 to avoid register spills
  constexpr int kUnroll = 6;
  constexpr int kFloatsPerLoop = kUnroll * kVecSizeInFloat;

  int remainder = size - prologue;
  int vectorizable = prologue + (remainder / kFloatsPerLoop) * kFloatsPerLoop;

  for (; i < vectorizable; i += kFloatsPerLoop) {
    float32x4_t v0 = vld1q_f32_aligned(in + i + 0);
    float32x4_t v1 = vld1q_f32_aligned(in + i + 4);
    float32x4_t v2 = vld1q_f32_aligned(in + i + 8);
    float32x4_t v3 = vld1q_f32_aligned(in + i + 12);
    float32x4_t v4 = vld1q_f32_aligned(in + i + 16);
    float32x4_t v5 = vld1q_f32_aligned(in + i + 20);

    uint32x4_t gz0 = vcgtq_f32(v0, vZero);
    uint32x4_t gz1 = vcgtq_f32(v1, vZero);
    uint32x4_t gz2 = vcgtq_f32(v2, vZero);
    uint32x4_t gz3 = vcgtq_f32(v3, vZero);
    uint32x4_t gz4 = vcgtq_f32(v4, vZero);
    uint32x4_t gz5 = vcgtq_f32(v5, vZero);

    float32x4_t v0neg = vmulq_f32(v0, vW);
    float32x4_t v1neg = vmulq_f32(v1, vW);
    float32x4_t v2neg = vmulq_f32(v2, vW);
    float32x4_t v3neg = vmulq_f32(v3, vW);
    float32x4_t v4neg = vmulq_f32(v4, vW);
    float32x4_t v5neg = vmulq_f32(v5, vW);

    // v0 > 0 ? v0 : v0 * w
    v0 = vbslq_f32(gz0, v0, v0neg);
    v1 = vbslq_f32(gz1, v1, v1neg);
    v2 = vbslq_f32(gz2, v2, v2neg);
    v3 = vbslq_f32(gz3, v3, v3neg);
    v4 = vbslq_f32(gz4, v4, v4neg);
    v5 = vbslq_f32(gz5, v5, v5neg);

    vst1q_f32(out + i + 0, v0);
    vst1q_f32(out + i + 4, v1);
    vst1q_f32(out + i + 8, v2);
    vst1q_f32(out + i + 12, v3);
    vst1q_f32(out + i + 16, v4);
    vst1q_f32(out + i + 20, v5);
  }

  for (; i < size; ++i) {
    float v = in[i];
    out[i] = v > 0 ? v : v * w;
  }
}

}
#endif // __ARM_NEON__

template <>
bool PReluOp<float, CPUContext>::RunOnDevice() {
  const auto& X = Input(0);
  const auto& W = Input(1);
  auto* Y = Output(0);
  Y->ResizeLike(X);
  const auto* Xdata = X.template data<float>();
  const auto* Wdata = W.template data<float>();
  auto* Ydata = Y->template mutable_data<float>();

  const auto C = order_ == StorageOrder::NCHW ? X.dim(1) : X.dim(X.ndim() - 1);
  const auto C_shared = (W.size() == 1);

  if (!C_shared) {
    CAFFE_ENFORCE_EQ(C, W.size());
  }

  if (C_shared) {
#ifdef __ARM_NEON__
    // The function is completely pointwise
    runNeonPrelu(Ydata, Xdata, X.size(), Wdata[0]);
#else
    ConstEigenVectorMap<float> Xvec(Xdata, X.size());
    EigenVectorMap<float> Yvec(Ydata, Y->size());
    Yvec = Xvec.cwiseMax(0.f) + Xvec.cwiseMin(0.f) * Wdata[0];
#endif // __ARM_NEON__
    return true;
  }

  // non-shared case.
  switch (order_) {
    case StorageOrder::NCHW: {
      const auto N = X.dim(0);
      const auto dim = X.size_from_dim(2);

#ifdef __ARM_NEON__
      // Pointwise for each channel
      for (int n = 0; n < N; ++n) {
        for (int c = 0; c < C; ++c) {
          runNeonPrelu(Ydata + (n * C + c) * dim,
                       Xdata + (n * C + c) * dim,
                       dim, Wdata[c]);
        }
      }
#else
      int nc = 0;
      for (int n = 0; n < N; ++n) {
{'fraction_non_alphanumeric': 0.14520376175548588, 'fraction_numerical': 0.06369905956112852, 'mean_word_length': 2.746359793330202, 'pattern_counts': {'":': 0, '<': 33, '<?xml version=': 0, '>': 31, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1526062', 'n_tokens_mistral': 3837, 'n_tokens_neox': 3617, 'n_words': 834}