// Copyright 2020 Google LLC
//
// This source code is licensed under the BSD-style license found in the
// LICENSE file in the root directory of this source tree.

$assert BATCH_TILE >= 1
#include <assert.h>
#include <math.h>

#include <xnnpack/common.h>
#include <xnnpack/math.h>
#include <xnnpack/vunary.h>


extern XNN_INTERNAL const uint32_t xnn_table_exp2minus_k_over_16[16];

void xnn_f32_velu_ukernel__${"wasm" if WASM else "scalar"}_rr2_lut16_p3_x${BATCH_TILE}(
    size_t n,
    const float* x,
    float* y,
    const union xnn_f32_elu_params params[restrict XNN_MIN_ELEMENTS(1)])
{
  assert(n % sizeof(float) == 0);

  const float vprescale = params->scalar_rr2_lut16_p3.prescale;
  const float valpha = params->scalar_rr2_lut16_p3.alpha;
  const float vbeta = params->scalar_rr2_lut16_p3.beta;
  const float vmagic_bias = params->scalar_rr2_lut16_p3.magic_bias;
  const float vlog2e = params->scalar_rr2_lut16_p3.log2e;
  const uint32_t vindex_mask = UINT32_C(0xF);
  const float vsat_cutoff = params->scalar_rr2_lut16_p3.sat_cutoff;
  const float vminus_ln2_hi = params->scalar_rr2_lut16_p3.minus_ln2_hi;
  const float vminus_ln2_lo = params->scalar_rr2_lut16_p3.minus_ln2_lo;
  const float vc3 = params->scalar_rr2_lut16_p3.c3;
  const float vc2 = params->scalar_rr2_lut16_p3.c2;
  const float vone = params->scalar_rr2_lut16_p3.one;

  $if BATCH_TILE > 1:
    for (; n >= ${BATCH_TILE} * sizeof(float); n -= ${BATCH_TILE} * sizeof(float)) {
      $for N in range(BATCH_TILE):
        float vx${N} = x[${N}];
      x += ${BATCH_TILE};

      $for N in range(BATCH_TILE):
        $if WASM:
          const float vz${N} = __builtin_wasm_min_f32(__builtin_wasm_max_f32(vx${N} * vprescale, vsat_cutoff), 0.0f);
        $else:
          const float vz${N} = vx${N} * vprescale;

      $for N in range(BATCH_TILE):
        float vn${N} = vz${N} * vlog2e + vmagic_bias;

      $for N in range(BATCH_TILE):
        const uint32_t ven${N} = float_as_uint32(vn${N}) << 19;
        const uint32_t vidx${N} = float_as_uint32(vn${N}) & vindex_mask;
        vn${N} -= vmagic_bias;

      $for N in range(BATCH_TILE):
        float vt${N} = vn${N} * vminus_ln2_hi + vz${N};
        float vs${N} = uint32_as_float(xnn_table_exp2minus_k_over_16[vidx${N}] + ven${N});

      $for N in range(BATCH_TILE):
        vt${N} = vn${N} * vminus_ln2_lo + vt${N};
        $if not WASM:
          if XNN_UNPREDICTABLE(vz${N} <= vsat_cutoff) {
            vs${N} = 0.0f;
            vt${N} = 0.0f;
          }

      $for N in range(BATCH_TILE):
        float vp${N} = vc3 * vt${N} + vc2;

      $for N in range(BATCH_TILE):
        vp${N} *= vt${N};

      $for N in range(BATCH_TILE):
        vt${N} *= vs${N};
        vs${N} -= vone;

      $for N in range(BATCH_TILE):
        vp${N} = vp${N} * vt${N} + vt${N};

      $for N in range(BATCH_TILE):
        const float ve${N} = (vp${N} + vs${N}) * valpha;
        $if WASM:
          float vy${N} = __builtin_wasm_max_f32(vx${N} * vbeta, 0.0f);
        $else:
          float vy${N} = vx${N} * vbeta;

      $if WASM:
        $for N in range(BATCH_TILE):
          vy${N} += __builtin_wasm_min_f32(ve${N}, 0.0f);
      $else:
        $for N in range(BATCH_TILE):
          if XNN_UNPREDICTABLE(vx${N} < 0.0f) {
            vy${N} = ve${N};
          }

      $for N in range(BATCH_TILE):
        y[${N}] = vy${N};
      y += ${BATCH_TILE};
    }
  $if BATCH_TILE == 1:
    do {
      float vx = *x++;

      $if WASM:
        const float vz = __builtin_wasm_min_f32(__builtin_wasm_max_f32(vx * vprescale, vsat_cutoff), 0.0f);
      $else:
        const float vz = vx * vprescale;

      float vn = vz * vlog2e + vmagic_bias;
      const uint32_t ven = float_as_uint32(vn) << 19;
      const uint32_t vidx = float_as_uint32(vn) & vindex_mask;
      vn -= vmagic_bias;

      float vt = vn * vminus_ln2_hi + vz;
      float vs = uint32_as_float(xnn_table_exp2minus_k_over_16[vidx] + ven);

      vt = vn * vminus_ln2_lo + vt;
      $if not WASM:
        if XNN_UNPREDICTABLE(vz <= vsat_cutoff) {
          vs = 0.0f;
          vt = 0.0f;
        }

      float vp = vc3 * vt + vc2;
      vp *= vt;

      vt *= vs;
      vs -= vone;
      vp = vp * vt + vt;
      const float ve = (vp + vs) * valpha;

      $if WASM:
        float vy = __builtin_wasm_max_f32(vx * vbeta, 0.0f);
        vy += __builtin_wasm_min_f32(ve, 0.0f);
      $else:
        float vy = vx * vbeta;
        if XNN_UNPREDICTABLE(vx < 0.0f) {
          vy = ve;
        }

      *y++ = vy;

      n -= sizeof(float);
    } while (n != 0);
  $elif BATCH_TILE == 2:
    if XNN_UNLIKELY(n != 0) {
      float vx = *x;

      $if WASM:
        const float vz = __builtin_wasm_min_f32(__builtin_wasm_max_f32(vx * vprescale, vsat_cutoff), 0.0f);
      $else:
        const float vz = vx * vprescale;

      float vn = vz * vlog2e + vmagic_bias;
      const uint32_t ven = float_as_uint32(vn) << 19;
      const uint32_t vidx = float_as_uint32(vn) & vindex_mask;
      vn -= vmagic_bias;

      float vt = vn * vminus_ln2_hi + vz;
      float vs = uint32_as_float(xnn_table_exp2minus_k_over_16[vidx] + ven);

      vt = vn * vminus_ln2_lo + vt;
      $if not WASM:
        if XNN_UNPREDICTABLE(vz <= vsat_cutoff) {
          vs = 0.0f;
          vt = 0.0f;
        }

      float vp = vc3 * vt + vc2;
      vp *= vt;

      vt *= vs;
      vs -= vone;
      vp = vp * vt + vt;
      const float ve = (vp + vs) * valpha;

      $if WASM:
        float vy = __builtin_wasm_max_f32(vx * vbeta, 0.0f);
        vy += __builtin_wasm_min_f32(ve, 0.0f);
      $else:
        float vy = vx * vbeta;
        if XNN_UNPREDICTABLE(vx < 0.0f) {
          vy = ve;
        }

      *y = vy;
    }
  $else:
    if XNN_UNLIKELY(n != 0) {
      do {
        float vx = *x++;

        $if WASM:
          const float vz = __builtin_wasm_min_f32(__builtin_wasm_max_f32(vx * vprescale, vsat_cutoff), 0.0f);
        $else:
          const float vz = vx * vprescale;

        float vn = vz * vlog2e + vmagic_bias;
        const uint32_t ven = float_as_uint32(vn) << 19;
        const uint32_t vidx = float_as_uint32(vn) & vindex_mask;
        vn -= vmagic_bias;

        float vt = vn * vminus_ln2_hi + vz;
        float vs = uint32_as_float(xnn_table_exp2minus_k_over_16[vidx] + ven);

        vt = vn * vminus_ln2_lo + vt;
        $if not WASM:
          if XNN_UNPREDICTABLE(vz <= vsat_cutoff) {
            vs = 0.0f;
            vt = 0.0f;
          }

        float vp = vc3 * vt + vc2;
        vp *= vt;

        vt *= vs;
        vs -= vone;
        vp = vp * vt + vt;
        const float ve = (vp + vs) * valpha;

        $if WASM:
          float vy = __builtin_wasm_max_f32(vx * vbeta, 0.0f);
          vy += __builtin_wasm_min_f32(ve, 0.0f);
        $else:
          float vy = vx * vbeta;
          if XNN_UNPREDICTABLE(vx < 0.0f) {
            vy = ve;
          }

        *y++ = vy;

        n -= sizeof(float);
      } while (n != 0);
    }
}