
$assert BATCH_TILE % 32 == 0
$assert BATCH_TILE >= 32
$SIMD_TILE = BATCH_TILE // 32
$assert OP in ["ADD", "MAX", "MIN", "MUL", "SUB", "SQRDIFF"]
#include <assert.h>

#include "xnnpack/simd/f32-hvx.h"

#include "xnnpack/math.h"
#include "xnnpack/vbinary.h"

$_HEXAGON_OP_HVX = {
$  "ADD": "xnn_add_f32",
$  "MAX": "xnn_max_f32",
$  "MIN": "xnn_min_f32", 
$  "MUL": "xnn_mul_f32", 
$  "SUB": "xnn_sub_f32",
$  "SQRDIFF": "xnn_sub_f32",
$}[OP]
void xnn_f32_v${OP.lower()}_ukernel__hvx_u${BATCH_TILE}(
    size_t batch,
    const float* input_a,
    const float* input_b,
    float* output,
    const struct xnn_f32_default_params params[restrict XNN_MIN_ELEMENTS(1)]) XNN_OOB_READS
{
  assert(batch != 0);
  assert(batch % sizeof(float) == 0);
  assert(input_a != NULL);
  assert(input_b != NULL);
  assert(output != NULL);

  $if BATCH_TILE > 32:
    for (; batch >= ${BATCH_TILE} * sizeof(float); batch -= ${BATCH_TILE} * sizeof(float)) {
      HVX_Vector va0 = xnn_loadu_f32(input_a);
      $for N in range(32, BATCH_TILE, 32):
        HVX_Vector va${int(N/32)} = xnn_loadu_f32(input_a + ${N});
      HVX_Vector vb0 = xnn_loadu_f32(input_b);
      $for N in range(32, BATCH_TILE, 32):
        HVX_Vector vb${int(N/32)} = xnn_loadu_f32(input_b + ${N});
      input_a += ${BATCH_TILE};
      input_b += ${BATCH_TILE};

      $for N in range(SIMD_TILE):
        HVX_Vector vacc${N} = ${_HEXAGON_OP_HVX}(va${N}, vb${N});

      $if OP == "SQRDIFF":
        $for N in range(SIMD_TILE):
          vacc${N} = xnn_mul_f32(vacc${N}, vacc${N});

      xnn_storeu_f32(output, vacc0);
      $for N in range(32, BATCH_TILE, 32):
        xnn_storeu_f32(output + ${N}, vacc${int(N/32)});
      output += ${BATCH_TILE};
    }
  for (; batch >= 32 * sizeof(float); batch -= 32 * sizeof(float)) {
    HVX_Vector va = xnn_loadu_f32(input_a);
    HVX_Vector vb = xnn_loadu_f32(input_b);
    input_a += 32;
    input_b += 32;

    HVX_Vector vacc = ${_HEXAGON_OP_HVX}(va, vb);
    $if OP == "SQRDIFF":
      vacc = xnn_mul_f32(vacc, vacc);

    xnn_storeu_f32(output, vacc);
    output += 32;
  }
  if XNN_UNLIKELY(batch != 0) {
     HVX_Vector va = xnn_loadu_f32(input_a);
     HVX_Vector vb = xnn_loadu_f32(input_b);

     HVX_Vector vacc = ${_HEXAGON_OP_HVX}(va, vb);
     $if OP == "SQRDIFF":
       vacc = xnn_mul_f32(vacc, vacc);

     Q6_V_vstu_variable(output, batch, vacc);
  }
}
