av1/encoder/arm/av1_k_means_neon.c

/*
 * Copyright (c) 2023, Alliance for Open Media. All rights reserved.
 *
 * This source code is subject to the terms of the BSD 2 Clause License and
 * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
 * was not distributed with this source code in the LICENSE file, you can
 * obtain it at www.aomedia.org/license/software. If the Alliance for Open
 * Media Patent License 1.0 was not distributed with this source code in the
 * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
 */

#include <arm_neon.h>
#include <assert.h>

#include "aom_dsp/arm/sum_neon.h"
#include "aom_ports/mem.h"
#include "av1/encoder/random.h"
#include "config/aom_config.h"
#include "config/av1_rtcd.h"

static int32x4_t k_means_multiply_add_neon(const int16x8_t a) {
  const int32x4_t l = vmull_s16(vget_low_s16(a), vget_low_s16(a));
  const int32x4_t h = vmull_s16(vget_high_s16(a), vget_high_s16(a));
#if AOM_ARCH_AARCH64
  return vpaddq_s32(l, h);
#else
  const int32x2_t dl = vpadd_s32(vget_low_s32(l), vget_high_s32(l));
  const int32x2_t dh = vpadd_s32(vget_low_s32(h), vget_high_s32(h));
  return vcombine_s32(dl, dh);
#endif
}

void av1_calc_indices_dim1_neon(const int16_t *data, const int16_t *centroids,
                                uint8_t *indices, int64_t *total_dist, int n,
                                int k) {
  int64x2_t sum = vdupq_n_s64(0);
  int16x8_t cents[PALETTE_MAX_SIZE];
  for (int j = 0; j < k; ++j) {
    cents[j] = vdupq_n_s16(centroids[j]);
  }

  for (int i = 0; i < n; i += 8) {
    const int16x8_t in = vld1q_s16(data);
    uint16x8_t ind = vdupq_n_u16(0);
    // Compute the distance to the first centroid.
    int16x8_t dist_min = vabdq_s16(in, cents[0]);

    for (int j = 1; j < k; ++j) {
      // Compute the distance to the centroid.
      const int16x8_t dist = vabdq_s16(in, cents[j]);
      // Compare to the minimal one.
      const uint16x8_t cmp = vcgtq_s16(dist_min, dist);
      dist_min = vminq_s16(dist_min, dist);
      const uint16x8_t ind1 = vdupq_n_u16(j);
      ind = vbslq_u16(cmp, ind1, ind);
    }
    if (total_dist) {
      // Square, convert to 32 bit and add together.
      const int32x4_t l =
          vmull_s16(vget_low_s16(dist_min), vget_low_s16(dist_min));
      const int32x4_t sum32_tmp =
          vmlal_s16(l, vget_high_s16(dist_min), vget_high_s16(dist_min));
      // Pairwise sum, convert to 64 bit and add to sum.
      sum = vpadalq_s32(sum, sum32_tmp);
    }
    vst1_u8(indices, vmovn_u16(ind));
    indices += 8;
    data += 8;
  }
  if (total_dist) {
    *total_dist = horizontal_add_s64x2(sum);
  }
}

void av1_calc_indices_dim2_neon(const int16_t *data, const int16_t *centroids,
                                uint8_t *indices, int64_t *total_dist, int n,
                                int k) {
  int64x2_t sum = vdupq_n_s64(0);
  uint32x4_t ind[2];
  int16x8_t cents[PALETTE_MAX_SIZE];
  for (int j = 0; j < k; ++j) {
    const int16_t cx = centroids[2 * j], cy = centroids[2 * j + 1];
    const int16_t cxcy[8] = { cx, cy, cx, cy, cx, cy, cx, cy };
    cents[j] = vld1q_s16(cxcy);
  }

  for (int i = 0; i < n; i += 8) {
    for (int l = 0; l < 2; ++l) {
      const int16x8_t in = vld1q_s16(data);
      ind[l] = vdupq_n_u32(0);
      // Compute the distance to the first centroid.
      int16x8_t d1 = vsubq_s16(in, cents[0]);
      int32x4_t dist_min = k_means_multiply_add_neon(d1);

      for (int j = 1; j < k; ++j) {
        // Compute the distance to the centroid.
        d1 = vsubq_s16(in, cents[j]);
        const int32x4_t dist = k_means_multiply_add_neon(d1);
        // Compare to the minimal one.
        const uint32x4_t cmp = vcgtq_s32(dist_min, dist);
        dist_min = vminq_s32(dist_min, dist);
        const uint32x4_t ind1 = vdupq_n_u32(j);
        ind[l] = vbslq_u32(cmp, ind1, ind[l]);
      }
      if (total_dist) {
        // Pairwise sum, convert to 64 bit and add to sum.
        sum = vpadalq_s32(sum, dist_min);
      }
      data += 8;
    }
    // Cast to 8 bit and store.
    vst1_u8(indices,
            vmovn_u16(vcombine_u16(vmovn_u32(ind[0]), vmovn_u32(ind[1]))));
    indices += 8;
  }
  if (total_dist) {
    *total_dist = horizontal_add_s64x2(sum);
  }
}

// Adds the size and the pixel sum of cluster j to count and sum.
static inline void count_and_sum_neon(const int16_t *data,
                                      const uint8_t *indices, int n16, int j,
                                      int *count, int *sum) {
  const uint8x16_t index = vdupq_n_u8(j);
  int16x8_t count_acc = vdupq_n_s16(0);
  int32x4_t sum_acc[2] = { vdupq_n_s32(0), vdupq_n_s32(0) };
  for (int i = 0; i < n16; i += 16) {
    // All ones in the lanes that belong to cluster j.
    const uint8x16_t mask = vceqq_u8(vld1q_u8(indices + i), index);
    const uint8x16x2_t mask16 = vzipq_u8(mask, mask);
    const int16x8_t d0 = vld1q_s16(data + i);
    const int16x8_t d1 = vld1q_s16(data + i + 8);
    count_acc = vpadalq_s8(count_acc, vreinterpretq_s8_u8(mask));
    sum_acc[0] = vpadalq_s16(
        sum_acc[0], vandq_s16(d0, vreinterpretq_s16_u8(mask16.val[0])));
    sum_acc[1] = vpadalq_s16(
        sum_acc[1], vandq_s16(d1, vreinterpretq_s16_u8(mask16.val[1])));
  }
  count[j] -= horizontal_add_s16x8(count_acc);
  sum[j] += horizontal_add_s32x4(vaddq_s32(sum_acc[0], sum_acc[1]));
}

// Adds the size and the pixel sum of the clusters j and j + 1 to count and
// sum. Two clusters per pass share the loads.
static inline void count_and_sum_pair_neon(const int16_t *data,
                                           const uint8_t *indices, int n16,
                                           int j, int *count, int *sum) {
  const uint8x16_t index0 = vdupq_n_u8(j);
  const uint8x16_t index1 = vdupq_n_u8(j + 1);
  int16x8_t count_acc[2] = { vdupq_n_s16(0), vdupq_n_s16(0) };
  int32x4_t sum_acc[4] = { vdupq_n_s32(0), vdupq_n_s32(0), vdupq_n_s32(0),
                           vdupq_n_s32(0) };
  for (int i = 0; i < n16; i += 16) {
    const uint8x16_t idx = vld1q_u8(indices + i);
    const int16x8_t d0 = vld1q_s16(data + i);
    const int16x8_t d1 = vld1q_s16(data + i + 8);
    // All ones in the lanes that belong to cluster j or j + 1.
    const uint8x16_t mask0 = vceqq_u8(idx, index0);
    const uint8x16_t mask1 = vceqq_u8(idx, index1);
    const uint8x16x2_t mask0_16 = vzipq_u8(mask0, mask0);
    const uint8x16x2_t mask1_16 = vzipq_u8(mask1, mask1);
    count_acc[0] = vpadalq_s8(count_acc[0], vreinterpretq_s8_u8(mask0));
    count_acc[1] = vpadalq_s8(count_acc[1], vreinterpretq_s8_u8(mask1));
    sum_acc[0] = vpadalq_s16(
        sum_acc[0], vandq_s16(d0, vreinterpretq_s16_u8(mask0_16.val[0])));
    sum_acc[1] = vpadalq_s16(
        sum_acc[1], vandq_s16(d1, vreinterpretq_s16_u8(mask0_16.val[1])));
    sum_acc[2] = vpadalq_s16(
        sum_acc[2], vandq_s16(d0, vreinterpretq_s16_u8(mask1_16.val[0])));
    sum_acc[3] = vpadalq_s16(
        sum_acc[3], vandq_s16(d1, vreinterpretq_s16_u8(mask1_16.val[1])));
  }
  count[j] -= horizontal_add_s16x8(count_acc[0]);
  count[j + 1] -= horizontal_add_s16x8(count_acc[1]);
  sum[j] += horizontal_add_s32x4(vaddq_s32(sum_acc[0], sum_acc[1]));
  sum[j + 1] += horizontal_add_s32x4(vaddq_s32(sum_acc[2], sum_acc[3]));
}

void av1_calc_centroids_dim1_neon(const int16_t *data, int16_t *centroids,
                                  const uint8_t *indices, int n, int k) {
  int count[PALETTE_MAX_SIZE] = { 0 };
  int sum[PALETTE_MAX_SIZE] = { 0 };
  const int n16 = n & ~15;
  assert(n <= 32768);

  for (int i = n16; i < n; ++i) {
    ++count[indices[i]];
    sum[indices[i]] += data[i];
  }

  int j = 0;
  for (; j + 1 < k; j += 2) {
    count_and_sum_pair_neon(data, indices, n16, j, count, sum);
  }
  if (j < k) {
    count_and_sum_neon(data, indices, n16, j, count, sum);
  }

  unsigned int rand_state = (unsigned int)data[0];
  for (j = 0; j < k; ++j) {
    if (count[j] == 0) {
      centroids[j] = data[lcg_rand16(&rand_state) % n];
    } else {
      centroids[j] = DIVIDE_AND_ROUND(sum[j], count[j]);
    }
  }
}