Commit 394f9f7df6 for aom
commit 394f9f7df624e4268fd8c894ead9f059d279e480
Author: Ihor Buhaiov <freedbrt@gmail.com>
Date: Thu Oct 15 11:48:36 2026 +0200
Add SSE2 and AVX2 palette centroid kernels
Add SSE2 and AVX2 versions of av1_calc_centroids_dim1(). They work the
same way as the Neon version: the leftover pixels in C first, then one
pass for each pair of centroids over 16 pixels at a time, a single pass
for the last centroid when k is odd, and the same empty cluster
fallback and rounding as C.
A byte compare on the indices gives the pixels of the centroid. The
mask is widened to 16 bits, the data is masked and added up in 32 bit
lanes with madd, and the mask itself is added up in 16 bit lanes for
the count. Neither can overflow for n up to 32768, and all sums are
exact integers, so the result equals C. The test now also checks the
SSE2 and AVX2 versions against C.
On a Ryzen 9 5900X with MSVC, the kernel is 2.2 to 6.1 times faster
with SSE2 and 2.5 to 7.8 times faster with AVX2 for 64 and 256 pixel
blocks. On 2560x1440 screen recordings the encoder takes 1.5% to 2.4%
less time.
Change-Id: I914c0a83b83bb4f58d0a73c263bc6e68791c6f92
diff --git a/av1/common/av1_rtcd_defs.pl b/av1/common/av1_rtcd_defs.pl
index 459ccb13e7..78bbf0327f 100644
--- a/av1/common/av1_rtcd_defs.pl
+++ b/av1/common/av1_rtcd_defs.pl
@@ -410,7 +410,7 @@ if (aom_config("CONFIG_AV1_ENCODER") eq "yes") {
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/;
+ specialize qw/av1_calc_centroids_dim1 sse2 avx2 neon/;
add_proto qw/void av1_calc_centroids_dim2/, "const int16_t *data, int16_t *centroids, const uint8_t *indices, int n, int k";
diff --git a/av1/encoder/x86/av1_k_means_avx2.c b/av1/encoder/x86/av1_k_means_avx2.c
index 98d8435478..61e6fa1661 100644
--- a/av1/encoder/x86/av1_k_means_avx2.c
+++ b/av1/encoder/x86/av1_k_means_avx2.c
@@ -8,10 +8,13 @@
* 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 <assert.h>
#include <immintrin.h> // AVX2
#include "config/av1_rtcd.h"
#include "aom_dsp/x86/synonyms.h"
+#include "aom_ports/mem.h"
+#include "av1/encoder/random.h"
static int64_t k_means_horizontal_sum_avx2(__m256i a) {
const __m128i low = _mm256_castsi256_si128(a);
@@ -130,3 +133,99 @@ void av1_calc_indices_dim2_avx2(const int16_t *data, const int16_t *centroids,
*total_dist = k_means_horizontal_sum_avx2(sum);
}
}
+
+// Adds the pixels of d picked by the 16 bit mask to the count and sum
+// accumulators.
+static inline void accumulate_cluster_avx2(__m256i mask, __m256i d,
+ __m256i *count_acc,
+ __m256i *sum_acc) {
+ const __m256i ones = _mm256_set1_epi16(1);
+ *count_acc = _mm256_sub_epi16(*count_acc, mask);
+ *sum_acc = _mm256_add_epi32(
+ *sum_acc, _mm256_madd_epi16(_mm256_and_si256(d, mask), ones));
+}
+
+static inline void add_cluster_avx2(__m256i count_acc, __m256i sum_acc,
+ int *count, int *sum) {
+ const __m256i ones = _mm256_set1_epi16(1);
+ // Add up the sum and the count together, they end up in lanes 0 and 2.
+ const __m256i count32 = _mm256_madd_epi16(count_acc, ones);
+ const __m256i t256 =
+ _mm256_add_epi32(_mm256_unpacklo_epi64(sum_acc, count32),
+ _mm256_unpackhi_epi64(sum_acc, count32));
+ __m128i t = _mm_add_epi32(_mm256_castsi256_si128(t256),
+ _mm256_extracti128_si256(t256, 1));
+ t = _mm_add_epi32(t, _mm_srli_epi64(t, 32));
+ *sum += _mm_cvtsi128_si32(t);
+ *count += _mm_cvtsi128_si32(_mm_srli_si128(t, 8));
+}
+
+// Adds the size and the pixel sum of cluster j to count and sum.
+static inline void count_and_sum_avx2(const int16_t *data,
+ const uint8_t *indices, int n16, int j,
+ int *count, int *sum) {
+ const __m128i index = _mm_set1_epi8(j);
+ __m256i count_acc = _mm256_setzero_si256();
+ __m256i sum_acc = _mm256_setzero_si256();
+ for (int i = 0; i < n16; i += 16) {
+ // All ones in the lanes that belong to cluster j. The sign extension
+ // keeps the 16 bit mask in the same order as the data.
+ const __m256i mask = _mm256_cvtepi8_epi16(
+ _mm_cmpeq_epi8(_mm_loadu_si128((const __m128i *)(indices + i)), index));
+ const __m256i d = _mm256_loadu_si256((const __m256i *)(data + i));
+ accumulate_cluster_avx2(mask, d, &count_acc, &sum_acc);
+ }
+ add_cluster_avx2(count_acc, sum_acc, &count[j], &sum[j]);
+}
+
+// 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_avx2(const int16_t *data,
+ const uint8_t *indices, int n16,
+ int j, int *count, int *sum) {
+ const __m128i index0 = _mm_set1_epi8(j);
+ const __m128i index1 = _mm_set1_epi8(j + 1);
+ __m256i count_acc[2] = { _mm256_setzero_si256(), _mm256_setzero_si256() };
+ __m256i sum_acc[2] = { _mm256_setzero_si256(), _mm256_setzero_si256() };
+ for (int i = 0; i < n16; i += 16) {
+ const __m128i idx = _mm_loadu_si128((const __m128i *)(indices + i));
+ const __m256i d = _mm256_loadu_si256((const __m256i *)(data + i));
+ // All ones in the lanes that belong to cluster j or j + 1.
+ const __m256i mask0 = _mm256_cvtepi8_epi16(_mm_cmpeq_epi8(idx, index0));
+ const __m256i mask1 = _mm256_cvtepi8_epi16(_mm_cmpeq_epi8(idx, index1));
+ accumulate_cluster_avx2(mask0, d, &count_acc[0], &sum_acc[0]);
+ accumulate_cluster_avx2(mask1, d, &count_acc[1], &sum_acc[1]);
+ }
+ add_cluster_avx2(count_acc[0], sum_acc[0], &count[j], &sum[j]);
+ add_cluster_avx2(count_acc[1], sum_acc[1], &count[j + 1], &sum[j + 1]);
+}
+
+void av1_calc_centroids_dim1_avx2(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_avx2(data, indices, n16, j, count, sum);
+ }
+ if (j < k) {
+ count_and_sum_avx2(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/x86/av1_k_means_sse2.c b/av1/encoder/x86/av1_k_means_sse2.c
index eb4c98e9f0..9063cf6b00 100644
--- a/av1/encoder/x86/av1_k_means_sse2.c
+++ b/av1/encoder/x86/av1_k_means_sse2.c
@@ -9,10 +9,13 @@
* PATENTS file, you can obtain it at www.aomedia.org/license/patent.
*/
+#include <assert.h>
#include <emmintrin.h> // SSE2
#include "config/av1_rtcd.h"
#include "aom_dsp/x86/synonyms.h"
+#include "aom_ports/mem.h"
+#include "av1/encoder/random.h"
static int64_t k_means_horizontal_sum_sse2(__m128i a) {
const __m128i sum1 = _mm_unpackhi_epi64(a, a);
@@ -122,3 +125,101 @@ void av1_calc_indices_dim2_sse2(const int16_t *data, const int16_t *centroids,
*total_dist = k_means_horizontal_sum_sse2(sum);
}
}
+
+// Adds the pixels of d0 and d1 picked by the byte mask to the count and sum
+// accumulators.
+static inline void accumulate_cluster_sse2(__m128i mask, __m128i d0, __m128i d1,
+ __m128i *count_acc,
+ __m128i *sum_acc) {
+ const __m128i ones = _mm_set1_epi16(1);
+ const __m128i mask0 = _mm_unpacklo_epi8(mask, mask);
+ const __m128i mask1 = _mm_unpackhi_epi8(mask, mask);
+ // Each int16 lane counts at most 2 * 32768 / 16 pixels.
+ *count_acc = _mm_sub_epi16(*count_acc, _mm_add_epi16(mask0, mask1));
+ *sum_acc = _mm_add_epi32(
+ *sum_acc, _mm_add_epi32(_mm_madd_epi16(_mm_and_si128(d0, mask0), ones),
+ _mm_madd_epi16(_mm_and_si128(d1, mask1), ones)));
+}
+
+static inline void add_cluster_sse2(__m128i count_acc, __m128i sum_acc,
+ int *count, int *sum) {
+ const __m128i ones = _mm_set1_epi16(1);
+ // Add up the sum and the count together, they end up in lanes 0 and 2.
+ const __m128i count32 = _mm_madd_epi16(count_acc, ones);
+ __m128i t = _mm_add_epi32(_mm_unpacklo_epi64(sum_acc, count32),
+ _mm_unpackhi_epi64(sum_acc, count32));
+ t = _mm_add_epi32(t, _mm_srli_epi64(t, 32));
+ *sum += _mm_cvtsi128_si32(t);
+ *count += _mm_cvtsi128_si32(_mm_srli_si128(t, 8));
+}
+
+// Adds the size and the pixel sum of cluster j to count and sum.
+static inline void count_and_sum_sse2(const int16_t *data,
+ const uint8_t *indices, int n16, int j,
+ int *count, int *sum) {
+ const __m128i index = _mm_set1_epi8(j);
+ __m128i count_acc = _mm_setzero_si128();
+ __m128i sum_acc = _mm_setzero_si128();
+ for (int i = 0; i < n16; i += 16) {
+ // All ones in the lanes that belong to cluster j.
+ const __m128i mask =
+ _mm_cmpeq_epi8(_mm_loadu_si128((const __m128i *)(indices + i)), index);
+ const __m128i d0 = _mm_loadu_si128((const __m128i *)(data + i));
+ const __m128i d1 = _mm_loadu_si128((const __m128i *)(data + i + 8));
+ accumulate_cluster_sse2(mask, d0, d1, &count_acc, &sum_acc);
+ }
+ add_cluster_sse2(count_acc, sum_acc, &count[j], &sum[j]);
+}
+
+// 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_sse2(const int16_t *data,
+ const uint8_t *indices, int n16,
+ int j, int *count, int *sum) {
+ const __m128i index0 = _mm_set1_epi8(j);
+ const __m128i index1 = _mm_set1_epi8(j + 1);
+ __m128i count_acc[2] = { _mm_setzero_si128(), _mm_setzero_si128() };
+ __m128i sum_acc[2] = { _mm_setzero_si128(), _mm_setzero_si128() };
+ for (int i = 0; i < n16; i += 16) {
+ const __m128i idx = _mm_loadu_si128((const __m128i *)(indices + i));
+ const __m128i d0 = _mm_loadu_si128((const __m128i *)(data + i));
+ const __m128i d1 = _mm_loadu_si128((const __m128i *)(data + i + 8));
+ // All ones in the lanes that belong to cluster j or j + 1.
+ accumulate_cluster_sse2(_mm_cmpeq_epi8(idx, index0), d0, d1, &count_acc[0],
+ &sum_acc[0]);
+ accumulate_cluster_sse2(_mm_cmpeq_epi8(idx, index1), d0, d1, &count_acc[1],
+ &sum_acc[1]);
+ }
+ add_cluster_sse2(count_acc[0], sum_acc[0], &count[j], &sum[j]);
+ add_cluster_sse2(count_acc[1], sum_acc[1], &count[j + 1], &sum[j + 1]);
+}
+
+void av1_calc_centroids_dim1_sse2(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_sse2(data, indices, n16, j, count, sum);
+ }
+ if (j < k) {
+ count_and_sum_sse2(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/test/av1_k_means_test.cc b/test/av1_k_means_test.cc
index 1903588fdc..5b6af7cd08 100644
--- a/test/av1_k_means_test.cc
+++ b/test/av1_k_means_test.cc
@@ -389,6 +389,8 @@ INSTANTIATE_TEST_SUITE_P(
SSE2, AV1KmeansTest2,
::testing::Combine(::testing::Values(&av1_calc_indices_dim2_sse2),
::testing::ValuesIn(kValidBlockSize)));
+INSTANTIATE_TEST_SUITE_P(SSE2, AV1KmeansCentroidsTest1,
+ ::testing::Values(&av1_calc_centroids_dim1_sse2));
#endif
#if HAVE_AVX2
@@ -400,6 +402,8 @@ INSTANTIATE_TEST_SUITE_P(
AVX2, AV1KmeansTest2,
::testing::Combine(::testing::Values(&av1_calc_indices_dim2_avx2),
::testing::ValuesIn(kValidBlockSize)));
+INSTANTIATE_TEST_SUITE_P(AVX2, AV1KmeansCentroidsTest1,
+ ::testing::Values(&av1_calc_centroids_dim1_avx2));
#endif
#if HAVE_NEON