// 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>


void xnn_f32_f16_vcvt_ukernel__scalar_fabsf_x${BATCH_TILE}(
    size_t n,
    const float* input,
    void* output,
    const union xnn_f32_f16_cvt_params params[restrict XNN_MIN_ELEMENTS(1)])
{
  assert(n != 0);
  assert(n % sizeof(float) == 0);
  assert(input != NULL);
  assert(output != NULL);

  const float vscale_to_inf = params->scalar_fabsf.scale_to_inf;
  const uint32_t vexp_bias = params->scalar_fabsf.exp_bias;
  const float vscale_to_zero = params->scalar_fabsf.scale_to_zero;
  const uint32_t vexpw_max = params->scalar_fabsf.expw_max;
  const uint32_t vbias_min = params->scalar_fabsf.bias_min;
  const uint16_t vexph_mask = params->scalar_fabsf.exph_mask;
  const uint16_t vmanth_mask = params->scalar_fabsf.manth_mask;
  const uint16_t vnanh = params->scalar_fabsf.nanh;

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

      $for N in range(BATCH_TILE):
        const float vabsx${N} = fabsf(vx${N});
      $for N in range(BATCH_TILE):
        uint32_t vsignw${N} = float_as_uint32(vx${N});

      $for N in range(BATCH_TILE):
        const uint32_t vnonsignw${N} = float_as_uint32(vabsx${N});
      $for N in range(BATCH_TILE):
        float vf${N} = vabsx${N} * vscale_to_inf;

      $for N in range(BATCH_TILE):
        uint32_t vbias${N} = vnonsignw${N} + vexp_bias;
      $for N in range(BATCH_TILE):
        vsignw${N} ^= vnonsignw${N};

      $for N in range(BATCH_TILE):
        vf${N} *= vscale_to_zero;
      $for N in range(BATCH_TILE):
        vbias${N} &= vexpw_max;

      $for N in range(BATCH_TILE):
        vbias${N} = math_max_u32(vbias${N}, vbias_min);

      $for N in range(BATCH_TILE):
        vf${N} += uint32_as_float(vbias${N});

      $for N in range(BATCH_TILE):
        const uint32_t vbits${N} = float_as_uint32(vf${N});

      $for N in range(BATCH_TILE):
        const uint16_t vexph${N} = (uint16_t) (vbits${N} >> 13) & vexph_mask;
      $for N in range(BATCH_TILE):
        const uint16_t vmanth${N} = (uint16_t) vbits${N} & vmanth_mask;
      $for N in range(BATCH_TILE):
        const uint16_t vsignh${N} = (uint16_t) (vsignw${N} >> 16);

      $for N in range(BATCH_TILE):
        uint16_t vh${N} = vexph${N} + vmanth${N};
      $for N in range(BATCH_TILE):
        if XNN_UNPREDICTABLE(vnonsignw${N} > vexpw_max) {
          vh${N} = vnanh;
        }
      $for N in range(BATCH_TILE):
        vh${N} |= vsignh${N};

      $for N in range(BATCH_TILE):
        o[${N}] = vh${N};
      o += ${BATCH_TILE};
    }
  $if BATCH_TILE == 1:
    do {
      const float vx = *input++;

      const float vabsx = fabsf(vx);
      uint32_t vsignw = float_as_uint32(vx);

      const uint32_t vnonsignw = float_as_uint32(vabsx);
      float vf = vabsx * vscale_to_inf;

      uint32_t vbias = vnonsignw + vexp_bias;
      vsignw ^= vnonsignw;

      vf *= vscale_to_zero;
      vbias &= vexpw_max;

      vbias = math_max_u32(vbias, vbias_min);

      vf += uint32_as_float(vbias);

      const uint32_t vbits = float_as_uint32(vf);

      const uint16_t vexph = (uint16_t) (vbits >> 13) & vexph_mask;
      const uint16_t vmanth = (uint16_t) vbits & vmanth_mask;
      const uint16_t vsignh = (uint16_t) (vsignw >> 16);

      uint16_t vh = vexph + vmanth;
      if XNN_UNPREDICTABLE(vnonsignw > vexpw_max) {
        vh = vnanh;
      }
      vh |= vsignh;

      *o++ = vh;

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

      const float vabsx = fabsf(vx);
      uint32_t vsignw = float_as_uint32(vx);

      const uint32_t vnonsignw = float_as_uint32(vabsx);
      float vf = vabsx * vscale_to_inf;

      uint32_t vbias = vnonsignw + vexp_bias;
      vsignw ^= vnonsignw;

      vf *= vscale_to_zero;
      vbias &= vexpw_max;

      vbias = math_max_u32(vbias, vbias_min);

      vf += uint32_as_float(vbias);

      const uint32_t vbits = float_as_uint32(vf);

      const uint16_t vexph = (uint16_t) (vbits >> 13) & vexph_mask;
      const uint16_t vmanth = (uint16_t) vbits & vmanth_mask;
      const uint16_t vsignh = (uint16_t) (vsignw >> 16);

      uint16_t vh = vexph + vmanth;
      if XNN_UNPREDICTABLE(vnonsignw > vexpw_max) {
        vh = vnanh;
      }
      vh |= vsignh;

      *o = vh;
    }
  $else:
    if XNN_UNLIKELY(n != 0) {
      do {
        const float vx = *input++;

        const float vabsx = fabsf(vx);
        uint32_t vsignw = float_as_uint32(vx);

        const uint32_t vnonsignw = float_as_uint32(vabsx);
        float vf = vabsx * vscale_to_inf;

        uint32_t vbias = vnonsignw + vexp_bias;
        vsignw ^= vnonsignw;

        vf *= vscale_to_zero;
        vbias &= vexpw_max;

        vbias = math_max_u32(vbias, vbias_min);

        vf += uint32_as_float(vbias);

        const uint32_t vbits = float_as_uint32(vf);

        const uint16_t vexph = (uint16_t) (vbits >> 13) & vexph_mask;
        const uint16_t vmanth = (uint16_t) vbits & vmanth_mask;
        const uint16_t vsignh = (uint16_t) (vsignw >> 16);

        uint16_t vh = vexph + vmanth;
        if XNN_UNPREDICTABLE(vnonsignw > vexpw_max) {
          vh = vnanh;
        }
        vh |= vsignh;

        *o++ = vh;

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