// Copyright 2021 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/vcvt.h"


$XINT8_T = {"QS8": "int8_t", "QU8": "uint8_t"}[DATATYPE]
$MIN_F32 = "__builtin_wasm_min_f32" if WASM else "math_min_f32"
$MAX_F32 = "__builtin_wasm_max_f32" if WASM else "math_max_f32"
$OUTPUT_MIN = {"QS8": -128, "QU8": 0}[DATATYPE]
$OUTPUT_MAX = {"QS8": 127, "QU8": 255}[DATATYPE]
void xnn_f32_${DATATYPE.lower()}_vcvt_ukernel__${"wasm" if WASM else "scalar"}_lrintf_u${BATCH_TILE}(
    size_t batch,
    const float* input,
    ${XINT8_T}* output,
    const struct xnn_f32_${DATATYPE.lower()}_cvt_params params[restrict XNN_MIN_ELEMENTS(1)])
{
  assert(batch != 0);
  assert(batch % sizeof(float) == 0);
  assert(input != NULL);
  assert(output != NULL);

  const float vscale = params->scalar.scale;
  const float voutput_min_less_zero_point = (float) ((int32_t) ${OUTPUT_MIN} - (int32_t) params->scalar.output_zero_point);
  const float voutput_max_less_zero_point = (float) ((int32_t) ${OUTPUT_MAX} - (int32_t) params->scalar.output_zero_point);
  const int32_t voutput_zero_point = params->scalar.output_zero_point;

  $if BATCH_TILE == 1:
    do {
      float vx = *input++;
      vx *= vscale;
      vx = ${MAX_F32}(vx, voutput_min_less_zero_point);
      vx = ${MIN_F32}(vx, voutput_max_less_zero_point);

      int32_t vy = (int32_t) lrintf(vx);
      vy += voutput_zero_point;

      *output++ = (${XINT8_T}) vy;

      batch -= sizeof(float);
    } while (batch != 0);
  $else:
    for (; batch >= ${BATCH_TILE} * sizeof(float); batch -= ${BATCH_TILE} * sizeof(float)) {
      $for N in range(BATCH_TILE):
        float vx${N} = input[${N}];
      input += ${BATCH_TILE};

      $for N in range(BATCH_TILE):
        vx${N} *= vscale;

      $for N in range(BATCH_TILE):
        vx${N} = ${MAX_F32}(vx${N}, voutput_min_less_zero_point);

      $for N in range(BATCH_TILE):
        vx${N} = ${MIN_F32}(vx${N}, voutput_max_less_zero_point);

      $for N in range(BATCH_TILE):
        int32_t vy${N} = (int32_t) lrintf(vx${N});

      $for N in range(BATCH_TILE):
        vy${N} += voutput_zero_point;

      $for N in range(BATCH_TILE):
        output[${N}] = (${XINT8_T}) vy${N};
      output += ${BATCH_TILE};
    }
    $if BATCH_TILE == 2:
      if XNN_UNLIKELY(batch != 0) {
        float vx = *input;
        vx *= vscale;
        vx = ${MAX_F32}(vx, voutput_min_less_zero_point);
        vx = ${MIN_F32}(vx, voutput_max_less_zero_point);

        int32_t vy = (int32_t) lrintf(vx);
        vy += voutput_zero_point;

        *output = (${XINT8_T}) vy;
      }
    $else:
      if XNN_UNLIKELY(batch != 0) {
        do {
          float vx = *input++;
          vx *= vscale;
          vx = ${MAX_F32}(vx, voutput_min_less_zero_point);
          vx = ${MIN_F32}(vx, voutput_max_less_zero_point);

          int32_t vy = (int32_t) lrintf(vx);
          vy += voutput_zero_point;

          *output++ = (${XINT8_T}) vy;

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