// 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/vunary.h>

#include <fp16/bitcasts.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.prescale;
  const float valpha = params->scalar.alpha;
  const float vbeta = params->scalar.beta;

  const float vmagic_bias = 0x1.800000p19f;
  const float vlog2e = 0x1.715476p+0f;
  const uint32_t vindex_mask = UINT32_C(0xF);
  const float vsat_cutoff = -0x1.154246p+4f;
  const float vminus_ln2_hi = -0x1.62E400p-1f;
  const float vminus_ln2_lo = -0x1.7F7D1Cp-20f;
  const float vc3 = 0x1.55561Cp-3f;
  const float vc2 = 0x1.0001ECp-1f;
  const float vone = 1.0f;

  $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} = fp32_to_bits(vn${N}) << 19;
        const uint32_t vidx${N} = fp32_to_bits(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} = fp32_from_bits(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 = fp32_to_bits(vn) << 19;
      const uint32_t vidx = fp32_to_bits(vn) & vindex_mask;
      vn -= vmagic_bias;

      float vt = vn * vminus_ln2_hi + vz;
      float vs = fp32_from_bits(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 = fp32_to_bits(vn) << 19;
      const uint32_t vidx = fp32_to_bits(vn) & vindex_mask;
      vn -= vmagic_bias;

      float vt = vn * vminus_ln2_hi + vz;
      float vs = fp32_from_bits(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 = fp32_to_bits(vn) << 19;
        const uint32_t vidx = fp32_to_bits(vn) & vindex_mask;
        vn -= vmagic_bias;

        float vt = vn * vminus_ln2_hi + vz;
        float vs = fp32_from_bits(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);
    }
}