Commit 5cc71c7f89 for aom
commit 5cc71c7f89ec9b85a50fec006885a23027f9d95d
Author: Ihor Buhaiov <freedbrt@gmail.com>
Date: Thu Oct 15 19:37:24 2026 +0200
Add a Neon version of the k-means centroid step
The centroid step of av1_k_means() was a static C function. Its loop
updates one counter and one sum per pixel, and when neighboring pixels
pick the same color each update waits on the store of the one before.
Move it into RTCD as av1_calc_centroids_dim1() and dim2(), the same
way as av1_calc_indices, and add a Neon version for dim1 that counts
and sums two centroids per pass, 16 pixels per step. The empty cluster
fallback and the rounding are the same as in C, so the output does not
change. Add a test that checks the Neon version against C.
On real palette inputs the Neon version is 3.9 to 4.8 times faster. On
2560x1440 screen recordings the encoder uses up to 5.5% fewer cycles.
Change-Id: I3b5e8e2c2a4e9358875377ab3297aa2f5a6c63fa
diff --git a/av1/common/av1_rtcd_defs.pl b/av1/common/av1_rtcd_defs.pl
index 560d33c08d..459ccb13e7 100644
--- a/av1/common/av1_rtcd_defs.pl
+++ b/av1/common/av1_rtcd_defs.pl
@@ -409,6 +409,11 @@ if (aom_config("CONFIG_AV1_ENCODER") eq "yes") {
add_proto qw/void av1_calc_indices_dim2/, "const int16_t *data, const int16_t *centroids, uint8_t *indices, int64_t *total_dist, int n, int k";
specialize qw/av1_calc_indices_dim2 sse2 avx2 neon/;
+ add_proto qw/void av1_calc_centroids_dim1/, "const int16_t *data, int16_t *centroids, const uint8_t *indices, int n, int k";
+ specialize qw/av1_calc_centroids_dim1 neon/;
+
+ add_proto qw/void av1_calc_centroids_dim2/, "const int16_t *data, int16_t *centroids, const uint8_t *indices, int n, int k";
+
# ENCODEMB INVOKE
if (aom_config("CONFIG_AV1_HIGHBITDEPTH") eq "yes") {
add_proto qw/int64_t av1_highbd_block_error/, "const tran_low_t *coeff, const tran_low_t *dqcoeff, intptr_t block_size, int64_t *ssz, int bd";
diff --git a/av1/encoder/arm/av1_k_means_neon.c b/av1/encoder/arm/av1_k_means_neon.c
index eab9587b4d..21bc878e4e 100644
--- a/av1/encoder/arm/av1_k_means_neon.c
+++ b/av1/encoder/arm/av1_k_means_neon.c
@@ -10,8 +10,11 @@
*/
#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"
@@ -114,3 +117,92 @@ void av1_calc_indices_dim2_neon(const int16_t *data, const int16_t *centroids,
*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]);
+ }
+ }
+}
diff --git a/av1/encoder/k_means_template.h b/av1/encoder/k_means_template.h
index 4199159c89..ce74f93a4c 100644
--- a/av1/encoder/k_means_template.h
+++ b/av1/encoder/k_means_template.h
@@ -70,8 +70,8 @@ void RENAME_C(av1_calc_indices)(const int16_t *data, const int16_t *centroids,
}
}
-static void RENAME(calc_centroids)(const int16_t *data, int16_t *centroids,
- const uint8_t *indices, int n, int k) {
+void RENAME_C(av1_calc_centroids)(const int16_t *data, int16_t *centroids,
+ const uint8_t *indices, int n, int k) {
int i, j;
int count[PALETTE_MAX_SIZE] = { 0 };
int centroids_sum[AV1_K_MEANS_DIM * PALETTE_MAX_SIZE];
@@ -125,7 +125,13 @@ void RENAME(av1_k_means)(const int16_t *data, int16_t *centroids,
prev_l = l;
l = (l == 1) ? 0 : 1;
- RENAME(calc_centroids)(data, meta_centroids[l], meta_indices[prev_l], n, k);
+#if AV1_K_MEANS_DIM == 1
+ av1_calc_centroids_dim1(data, meta_centroids[l], meta_indices[prev_l], n,
+ k);
+#else
+ av1_calc_centroids_dim2(data, meta_centroids[l], meta_indices[prev_l], n,
+ k);
+#endif
if (!memcmp(meta_centroids[l], meta_centroids[prev_l],
sizeof(centroids[0]) * k * AV1_K_MEANS_DIM)) {
break;
diff --git a/test/av1_k_means_test.cc b/test/av1_k_means_test.cc
index 58dcb016e3..1903588fdc 100644
--- a/test/av1_k_means_test.cc
+++ b/test/av1_k_means_test.cc
@@ -251,6 +251,127 @@ TEST_P(AV1KmeansTest2, DISABLED_Speed) {
RunSpeedTest(GET_PARAM(0), GET_PARAM(1), 8);
}
+using av1_calc_centroids_func = void (*)(const int16_t *data,
+ int16_t *centroids,
+ const uint8_t *indices, int n, int k);
+
+class AV1KmeansCentroidsTest1
+ : public ::testing::TestWithParam<av1_calc_centroids_func> {
+ public:
+ ~AV1KmeansCentroidsTest1() override;
+ void SetUp() override;
+
+ protected:
+ enum DataType { kData8Bit, kData12Bit, kDataInt16 };
+ enum IndexType { kIndexRandom, kIndexRuns, kIndexFewUsed, kIndexSame };
+
+ void FillData(DataType data_type);
+ void FillIndices(IndexType index_type, int n, int k);
+ void RunCheckOutput(int n, int k);
+ void RunSpeedTest(int n, int k);
+
+ libaom_test::ACMRandom rnd_;
+ int16_t data_[4096];
+ uint8_t indices_[4096];
+ int16_t centroids1_[PALETTE_MAX_SIZE];
+ int16_t centroids2_[PALETTE_MAX_SIZE];
+};
+GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(AV1KmeansCentroidsTest1);
+
+AV1KmeansCentroidsTest1::~AV1KmeansCentroidsTest1() = default;
+
+void AV1KmeansCentroidsTest1::SetUp() {
+ rnd_.Reset(libaom_test::ACMRandom::DeterministicSeed());
+}
+
+void AV1KmeansCentroidsTest1::FillData(DataType data_type) {
+ for (int i = 0; i < 4096; ++i) {
+ switch (data_type) {
+ case kData8Bit: data_[i] = rnd_.Rand8(); break;
+ case kData12Bit: data_[i] = rnd_.Rand16() & 0xfff; break;
+ default: data_[i] = rnd_.Rand16Signed(); break;
+ }
+ }
+}
+
+void AV1KmeansCentroidsTest1::FillIndices(IndexType index_type, int n, int k) {
+ int i = 0;
+ while (i < n) {
+ switch (index_type) {
+ case kIndexRandom: indices_[i++] = rnd_(k); break;
+ case kIndexRuns: {
+ const int index = rnd_(k);
+ for (int run = 1 + rnd_(32); run > 0 && i < n; --run) {
+ indices_[i++] = index;
+ }
+ break;
+ }
+ // Leaves some centroids without pixels.
+ case kIndexFewUsed: indices_[i++] = rnd_((k + 1) / 2); break;
+ default: indices_[i++] = k - 1; break;
+ }
+ }
+}
+
+void AV1KmeansCentroidsTest1::RunCheckOutput(int n, int k) {
+ const av1_calc_centroids_func test_impl = GetParam();
+ for (int data_type = kData8Bit; data_type <= kDataInt16; ++data_type) {
+ for (int index_type = kIndexRandom; index_type <= kIndexSame;
+ ++index_type) {
+ FillData(static_cast<DataType>(data_type));
+ FillIndices(static_cast<IndexType>(index_type), n, k);
+ av1_calc_centroids_dim1_c(data_, centroids1_, indices_, n, k);
+ API_REGISTER_STATE_CHECK(test_impl(data_, centroids2_, indices_, n, k));
+ for (int j = 0; j < k; ++j) {
+ ASSERT_EQ(centroids1_[j], centroids2_[j])
+ << "n " << n << " k " << k << " data " << data_type << " index "
+ << index_type << " centroid " << j;
+ }
+ }
+ }
+}
+
+void AV1KmeansCentroidsTest1::RunSpeedTest(int n, int k) {
+ const av1_calc_centroids_func test_impl = GetParam();
+ const int num_loops = 1000000000 / n;
+ FillData(kData8Bit);
+ FillIndices(kIndexRuns, n, k);
+
+ av1_calc_centroids_func funcs[2] = { av1_calc_centroids_dim1_c, test_impl };
+ double elapsed_time[2] = { 0 };
+ for (int i = 0; i < 2; ++i) {
+ aom_usec_timer timer;
+ aom_usec_timer_start(&timer);
+ for (int j = 0; j < num_loops; ++j) {
+ funcs[i](data_, centroids1_, indices_, n, k);
+ }
+ aom_usec_timer_mark(&timer);
+ double time = static_cast<double>(aom_usec_timer_elapsed(&timer));
+ elapsed_time[i] = 1000.0 * time / num_loops;
+ }
+ printf("av1_calc_centroids_dim1 n=%d centroids=%d: %7.2f/%7.2fns", n, k,
+ elapsed_time[0], elapsed_time[1]);
+ printf("(%3.2f)\n", elapsed_time[0] / elapsed_time[1]);
+}
+
+TEST_P(AV1KmeansCentroidsTest1, CheckOutput) {
+ // Block sizes 8x8 to 64x64, plus sizes that are not a multiple of 16.
+ const int kSizes[] = { 64, 128, 256, 512, 1024, 2048, 4096, 1, 15, 33, 4095 };
+ for (int n : kSizes) {
+ for (int k = PALETTE_MIN_SIZE; k <= PALETTE_MAX_SIZE; ++k) {
+ RunCheckOutput(n, k);
+ }
+ }
+}
+
+TEST_P(AV1KmeansCentroidsTest1, DISABLED_Speed) {
+ for (int n : { 64, 256, 1024, 4096 }) {
+ for (int k = PALETTE_MIN_SIZE; k <= PALETTE_MAX_SIZE; ++k) {
+ RunSpeedTest(n, k);
+ }
+ }
+}
+
#if HAVE_SSE2 || HAVE_AVX2 || HAVE_NEON
const BLOCK_SIZE kValidBlockSize[] = { BLOCK_8X8, BLOCK_8X16, BLOCK_8X32,
BLOCK_16X8, BLOCK_16X16, BLOCK_16X32,
@@ -290,6 +411,8 @@ INSTANTIATE_TEST_SUITE_P(
NEON, AV1KmeansTest2,
::testing::Combine(::testing::Values(&av1_calc_indices_dim2_neon),
::testing::ValuesIn(kValidBlockSize)));
+INSTANTIATE_TEST_SUITE_P(NEON, AV1KmeansCentroidsTest1,
+ ::testing::Values(&av1_calc_centroids_dim1_neon));
#endif
} // namespace AV1Kmeans