// 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 CHANNEL_TILE % 8 == 0
$assert CHANNEL_TILE >= 8
$assert ROW_TILE >= 1
$ABC = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
#include <assert.h>

#include <immintrin.h>

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


void xnn_f16_prelu_ukernel__f16c_${ROW_TILE}x${CHANNEL_TILE}(
    size_t rows,
    size_t channels,
    const void* restrict input,
    size_t input_stride,
    const void* restrict weights,
    void* restrict output,
    size_t output_stride) XNN_OOB_READS
{
  assert(rows != 0);
  assert(channels != 0);
  assert(channels % sizeof(uint16_t) == 0);

  const uint16_t* i0 = (const uint16_t*) input;
  uint16_t* o0 = (uint16_t*) output;
  $for M in range(1, ROW_TILE):
    const uint16_t* i${M} = (const uint16_t*) ((uintptr_t) i${M-1} + input_stride);
    uint16_t* o${M} = (uint16_t*) ((uintptr_t) o${M-1} + output_stride);

  const size_t input_increment = input_stride * ${ROW_TILE} - channels;
  const size_t output_increment = output_stride * ${ROW_TILE} - channels;

  do {
    $for M in range(1, ROW_TILE):
      $if M % 2 == 0:
        if XNN_UNPREDICTABLE(rows <= ${M}) {
          i${M} = i${M-1};
          o${M} = o${M-1};
        }
      $else:
        if XNN_UNPREDICTABLE(rows < ${M+1}) {
          i${M} = i${M-1};
          o${M} = o${M-1};
        }

    const uint16_t* w = (const uint16_t*) weights;
    size_t c = channels;
    $if CHANNEL_TILE > 8:
      for (; c >= ${CHANNEL_TILE} * sizeof(uint16_t); c -= ${CHANNEL_TILE} * sizeof(uint16_t)) {
        const __m256 vw${ABC[0:8]} = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) w));
        $for C in range(8, CHANNEL_TILE, 8):
          const __m256 vw${ABC[C:C+8]} = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) (w + ${C})));
        w += ${CHANNEL_TILE};

        $for M in range(ROW_TILE):
          const __m256 vi${M}x0${ABC[0:8]} = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) i${M}));
          $for C in range(8, CHANNEL_TILE, 8):
            const __m256 vi${M}x0${ABC[C:C+8]} = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) (i${M} + ${C})));
          i${M} += ${CHANNEL_TILE};

        $for M in range(ROW_TILE):
          $for C in range(0, CHANNEL_TILE, 8):
            __m256 vacc${M}x0${ABC[C:C+8]} = _mm256_mul_ps(vi${M}x0${ABC[C:C+8]}, vw${ABC[C:C+8]});

        $for M in range(ROW_TILE):
          $for C in range(0, CHANNEL_TILE, 8):
            vacc${M}x0${ABC[C:C+8]} = _mm256_blendv_ps(vi${M}x0${ABC[C:C+8]}, vacc${M}x0${ABC[C:C+8]}, vi${M}x0${ABC[C:C+8]});

        $for M in range(ROW_TILE):
          _mm_storeu_si128((__m128i*) o${M}, _mm256_cvtps_ph(vacc${M}x0${ABC[C:C+8]}, _MM_FROUND_NO_EXC));
          $for C in range(0, CHANNEL_TILE, 8):
            _mm_storeu_si128((__m128i*) (o${M} + ${C}), _mm256_cvtps_ph(vacc${M}x0${ABC[C:C+8]}, _MM_FROUND_NO_EXC));
          o${M} += ${CHANNEL_TILE};
      }
    for (; c >= 8 * sizeof(uint16_t); c -= 8 * sizeof(uint16_t)) {
      const __m256 vw01234567 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) w));
      w += 8;

      $for M in range(ROW_TILE):
        const __m256 vi${M}x01234567 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) i${M}));
        i${M} += 8;

      $for M in range(ROW_TILE):
        __m256 vacc${M}x01234567 = _mm256_mul_ps(vi${M}x01234567, vw01234567);

      $for M in range(ROW_TILE):
        vacc${M}x01234567 = _mm256_blendv_ps(vi${M}x01234567, vacc${M}x01234567, vi${M}x01234567);

      $for M in range(ROW_TILE):
        _mm_storeu_si128((__m128i*) o${M}, _mm256_cvtps_ph(vacc${M}x01234567, _MM_FROUND_NO_EXC));
        o${M} += 8;
    }
    if XNN_UNLIKELY(c != 0) {
      const __m256 vw01234567 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) w));

      $for M in range(ROW_TILE):
        const __m256 vi${M}x01234567 = _mm256_cvtph_ps(_mm_loadu_si128((const __m128i*) i${M}));
        i${M} = (const uint16_t*) ((uintptr_t) i${M} + c);

      $for M in range(ROW_TILE):
        __m256 vacc${M}x01234567 = _mm256_mul_ps(vi${M}x01234567, vw01234567);

      $for M in range(ROW_TILE):
        vacc${M}x01234567 = _mm256_blendv_ps(vi${M}x01234567, vacc${M}x01234567, vi${M}x01234567);

      $for M in range(ROW_TILE):
        __m128i vh${M}x01234567 = _mm256_cvtps_ph(vacc${M}x01234567, _MM_FROUND_NO_EXC);
      if (c & (4 * sizeof(uint16_t))) {
        $for M in range(ROW_TILE):
          _mm_storel_epi64((__m128i*) o${M}, vh${M}x01234567);

        $for M in range(ROW_TILE):
          vh${M}x01234567 = _mm_unpackhi_epi64(vh${M}x01234567, vh${M}x01234567);

        $for M in range(ROW_TILE):
          o${M} += 4;
      }
      if (c & (2 * sizeof(uint16_t))) {
        $for M in range(ROW_TILE):
          *((uint32_t*) o${M}) = (uint32_t) _mm_cvtsi128_si32(vh${M}x01234567);

        $for M in range(ROW_TILE):
          vh${M}x01234567 = _mm_srli_epi64(vh${M}x01234567, 32);

        $for M in range(ROW_TILE):
          o${M} += 2;
      }
      if (c & (1 * sizeof(uint16_t))) {
        $for M in range(ROW_TILE):
          *o${M} = (uint16_t) _mm_extract_epi16(vh${M}x01234567, 0);

        $for M in range(ROW_TILE):
          o${M} += 1;
      }
    }
    $for M in range(ROW_TILE):
      i${M} = (const uint16_t*) ((uintptr_t) i${M} + input_increment);
      o${M} = (uint16_t*) ((uintptr_t) o${M} + output_increment);
    rows = doz(rows, ${ROW_TILE});
  } while (rows != 0);
}