diff options
Diffstat (limited to 'third_party/aom/av1/encoder/arm/neon/av1_k_means_neon.c')
-rw-r--r-- | third_party/aom/av1/encoder/arm/neon/av1_k_means_neon.c | 115 |
1 files changed, 115 insertions, 0 deletions
diff --git a/third_party/aom/av1/encoder/arm/neon/av1_k_means_neon.c b/third_party/aom/av1/encoder/arm/neon/av1_k_means_neon.c new file mode 100644 index 0000000000..d13cc65ae0 --- /dev/null +++ b/third_party/aom/av1/encoder/arm/neon/av1_k_means_neon.c @@ -0,0 +1,115 @@ +/* + * Copyright (c) 2023, Alliance for Open Media. All Rights Reserved. + * + * Use of this source code is governed by a BSD-style license + * that can be found in the LICENSE file in the root of the source + * tree. An additional intellectual property rights grant can be found + * in the file PATENTS. All contributing project authors may + * be found in the AUTHORS file in the root of the source tree. + */ + +#include <arm_neon.h> + +#include "aom_dsp/arm/sum_neon.h" +#include "config/aom_config.h" +#include "config/aom_dsp_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); + } +} |