// 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 BATCH_TILE >= 1
#include <assert.h>
#include <stddef.h>
#include <stdint.h>

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


void xnn_cs16_vsquareabs_ukernel__scalar_x${BATCH_TILE}(
    size_t batch,
    const int16_t* input,
    uint32_t* output)
{
  assert(batch != 0);
  assert(batch % (sizeof(int16_t) * 2) == 0);
  assert(input != NULL);
  assert(output != NULL);

  $if BATCH_TILE == 1:
    do {
      const int32_t vr = (int32_t) input[0];
      const int32_t vi = (int32_t) input[1];
      input += 2;

      uint32_t vacc = (uint32_t) (vr * vr);
      vacc += (uint32_t) (vi * vi);

      *output++ = vacc;
      batch -= sizeof(int16_t) * 2;
    } while (batch != 0);
  $else:
    for (; batch >= ${BATCH_TILE * 2} * sizeof(int16_t); batch -= ${BATCH_TILE * 2} * sizeof(int16_t)) {
      $for C in range(BATCH_TILE):
        const int32_t vr${C} = (int32_t) input[${C * 2}];
        const int32_t vi${C} = (int32_t) input[${C * 2 + 1}];
      input += ${BATCH_TILE} * 2;

      $for C in range(BATCH_TILE):
        uint32_t vacc${C} = (uint32_t) (vr${C} * vr${C});

      $for C in range(BATCH_TILE):
        vacc${C} += (uint32_t) (vi${C} * vi${C});

      $for C in range(BATCH_TILE):
        output[${C}] = vacc${C};
      output += ${BATCH_TILE};
    }
    $if BATCH_TILE == 2:
      if XNN_LIKELY(batch != 0) {
        assert(batch == 2 * sizeof(int16_t));

        const int32_t vr = (int32_t) input[0];
        const int32_t vi = (int32_t) input[1];

        uint32_t vacc = (uint32_t) (vr * vr);
        vacc += (uint32_t) (vi * vi);

        *output = vacc;
      }
    $else:
      if XNN_LIKELY(batch != 0) {
        do {
          const int32_t vr = (int32_t) input[0];
          const int32_t vi = (int32_t) input[1];
          input += 2;

          uint32_t vacc = (uint32_t) (vr * vr);
          vacc += (uint32_t) (vi * vi);

          *output++ = vacc;
          batch -= sizeof(int16_t) * 2;
        } while (batch != 0);
      }
}