// Copyright 2022 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 DATATYPE in ["QS8", "QU8"]
$assert BATCH_TILE >= 1
$ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
#include <assert.h>

#include <xnnpack/math.h>
#include <xnnpack/vlrelu.h>


$XINT8_T = {"QS8": "int8_t", "QU8": "uint8_t"}[DATATYPE]
$OUTPUT_MIN = {"QS8": -128, "QU8": 0}[DATATYPE]
$OUTPUT_MAX = {"QS8": 127, "QU8": 255}[DATATYPE]
void xnn_${DATATYPE.lower()}_vlrelu_ukernel__scalar_select_x${BATCH_TILE}(
    size_t n,
    const ${XINT8_T}* x,
    ${XINT8_T}* y,
    const union xnn_${DATATYPE.lower()}_lrelu_params params[restrict XNN_MIN_ELEMENTS(1)])
{
  const int32_t vinput_zero_point = params->scalar_select.input_zero_point;
  const int32_t vpositive_multiplier = params->scalar_select.positive_multiplier;
  const int32_t vnegative_multiplier = params->scalar_select.negative_multiplier;
  const int32_t vbias = params->scalar_select.bias;
  $if BATCH_TILE == 1:
    do {
      int32_t vacc = (int32_t) *x++ - vinput_zero_point;
      const int32_t vmultiplier = XNN_UNPREDICTABLE(vacc >= 0) ? vpositive_multiplier : vnegative_multiplier;
      vacc = vbias + vacc * vmultiplier;

      int32_t vout = math_asr_s32(vacc, 8);
      vout = math_max_s32(vout, ${OUTPUT_MIN});
      vout = math_min_s32(vout, ${OUTPUT_MAX});
      *y++ = (${XINT8_T}) vout;

      n -= sizeof(${XINT8_T});
    } while (n != 0);
  $else:
    for (; n >= ${BATCH_TILE} * sizeof(${XINT8_T}); n -= ${BATCH_TILE} * sizeof(${XINT8_T})) {
      $for N in range(BATCH_TILE):
        int32_t vacc${ABC[N]} = (int32_t) x[${N}];
      x += ${BATCH_TILE};

      $for N in range(BATCH_TILE):
        vacc${ABC[N]} -= vinput_zero_point;

      $for N in range(BATCH_TILE):
        const int32_t vmultiplier${ABC[N]} = XNN_UNPREDICTABLE(vacc${ABC[N]} >= 0) ? vpositive_multiplier : vnegative_multiplier;

      $for N in range(BATCH_TILE):
        vacc${ABC[N]} = vbias + vacc${ABC[N]} * vmultiplier${ABC[N]};

      $for N in range(BATCH_TILE):
        int32_t vout${ABC[N]} = math_asr_s32(vacc${ABC[N]}, 8);

      $for N in range(BATCH_TILE):
        vout${ABC[N]} = math_max_s32(vout${ABC[N]}, ${OUTPUT_MIN});

      $for N in range(BATCH_TILE):
        vout${ABC[N]} = math_min_s32(vout${ABC[N]}, ${OUTPUT_MAX});

      $for N in range(BATCH_TILE):
        y[${N}] = (${XINT8_T}) vout${ABC[N]};
      y += ${BATCH_TILE};
    }
    if XNN_UNLIKELY(n != 0) {
      $if BATCH_TILE == 2:
        int32_t vacc = (int32_t) *x++ - vinput_zero_point;
        const int32_t vmultiplier = XNN_UNPREDICTABLE(vacc >= 0) ? vpositive_multiplier : vnegative_multiplier;
        vacc = vbias + vacc * vmultiplier;

        int32_t vout = math_asr_s32(vacc, 8);
        vout = math_max_s32(vout, ${OUTPUT_MIN});
        vout = math_min_s32(vout, ${OUTPUT_MAX});
        *y = (${XINT8_T}) vout;
      $else:
        do {
          int32_t vacc = (int32_t) *x++ - vinput_zero_point;
          const int32_t vmultiplier = XNN_UNPREDICTABLE(vacc >= 0) ? vpositive_multiplier : vnegative_multiplier;
          vacc = vbias + vacc * vmultiplier;

          int32_t vout = math_asr_s32(vacc, 8);
          vout = math_max_s32(vout, ${OUTPUT_MIN});
          vout = math_min_s32(vout, ${OUTPUT_MAX});
          *y++ = (${XINT8_T}) vout;

          n -= sizeof(${XINT8_T});
        } while (n != 0);
    }
}