xref: /aosp_15_r20/external/libaom/av1/encoder/k_means_template.h (revision 77c1e3ccc04c968bd2bc212e87364f250e820521)
1 /*
2  * Copyright (c) 2016, Alliance for Open Media. All rights reserved.
3  *
4  * This source code is subject to the terms of the BSD 2 Clause License and
5  * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
6  * was not distributed with this source code in the LICENSE file, you can
7  * obtain it at www.aomedia.org/license/software. If the Alliance for Open
8  * Media Patent License 1.0 was not distributed with this source code in the
9  * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
10  */
11 
12 #include <assert.h>
13 #include <stdint.h>
14 #include <stdlib.h>
15 #include <string.h>
16 
17 #include "av1/common/blockd.h"
18 #include "av1/encoder/palette.h"
19 #include "av1/encoder/random.h"
20 
21 #ifndef AV1_K_MEANS_DIM
22 #error "This template requires AV1_K_MEANS_DIM to be defined"
23 #endif
24 
25 #define RENAME_(x, y) AV1_K_MEANS_RENAME(x, y)
26 #define RENAME(x) RENAME_(x, AV1_K_MEANS_DIM)
27 #define K_MEANS_RENAME_C(x, y) x##_dim##y##_c
28 #define RENAME_C_(x, y) K_MEANS_RENAME_C(x, y)
29 #define RENAME_C(x) RENAME_C_(x, AV1_K_MEANS_DIM)
30 
31 // Though we want to compute the smallest L2 norm, in 1 dimension,
32 // it is equivalent to find the smallest L1 norm and then square it.
33 // This is preferrable for speed, especially on the SIMD side.
RENAME(calc_dist)34 static int RENAME(calc_dist)(const int16_t *p1, const int16_t *p2) {
35 #if AV1_K_MEANS_DIM == 1
36   return abs(p1[0] - p2[0]);
37 #else
38   int dist = 0;
39   for (int i = 0; i < AV1_K_MEANS_DIM; ++i) {
40     const int diff = p1[i] - p2[i];
41     dist += diff * diff;
42   }
43   return dist;
44 #endif
45 }
46 
RENAME_C(av1_calc_indices)47 void RENAME_C(av1_calc_indices)(const int16_t *data, const int16_t *centroids,
48                                 uint8_t *indices, int64_t *dist, int n, int k) {
49   if (dist) {
50     *dist = 0;
51   }
52   for (int i = 0; i < n; ++i) {
53     int min_dist = RENAME(calc_dist)(data + i * AV1_K_MEANS_DIM, centroids);
54     indices[i] = 0;
55     for (int j = 1; j < k; ++j) {
56       const int this_dist = RENAME(calc_dist)(data + i * AV1_K_MEANS_DIM,
57                                               centroids + j * AV1_K_MEANS_DIM);
58       if (this_dist < min_dist) {
59         min_dist = this_dist;
60         indices[i] = j;
61       }
62     }
63     if (dist) {
64 #if AV1_K_MEANS_DIM == 1
65       *dist += min_dist * min_dist;
66 #else
67       *dist += min_dist;
68 #endif
69     }
70   }
71 }
72 
RENAME(calc_centroids)73 static void RENAME(calc_centroids)(const int16_t *data, int16_t *centroids,
74                                    const uint8_t *indices, int n, int k) {
75   int i, j;
76   int count[PALETTE_MAX_SIZE] = { 0 };
77   int centroids_sum[AV1_K_MEANS_DIM * PALETTE_MAX_SIZE];
78   unsigned int rand_state = (unsigned int)data[0];
79   assert(n <= 32768);
80   memset(centroids_sum, 0, sizeof(centroids_sum[0]) * k * AV1_K_MEANS_DIM);
81 
82   for (i = 0; i < n; ++i) {
83     const int index = indices[i];
84     assert(index < k);
85     ++count[index];
86     for (j = 0; j < AV1_K_MEANS_DIM; ++j) {
87       centroids_sum[index * AV1_K_MEANS_DIM + j] +=
88           data[i * AV1_K_MEANS_DIM + j];
89     }
90   }
91 
92   for (i = 0; i < k; ++i) {
93     if (count[i] == 0) {
94       memcpy(centroids + i * AV1_K_MEANS_DIM,
95              data + (lcg_rand16(&rand_state) % n) * AV1_K_MEANS_DIM,
96              sizeof(centroids[0]) * AV1_K_MEANS_DIM);
97     } else {
98       for (j = 0; j < AV1_K_MEANS_DIM; ++j) {
99         centroids[i * AV1_K_MEANS_DIM + j] =
100             DIVIDE_AND_ROUND(centroids_sum[i * AV1_K_MEANS_DIM + j], count[i]);
101       }
102     }
103   }
104 }
105 
RENAME(av1_k_means)106 void RENAME(av1_k_means)(const int16_t *data, int16_t *centroids,
107                          uint8_t *indices, int n, int k, int max_itr) {
108   int16_t centroids_tmp[AV1_K_MEANS_DIM * PALETTE_MAX_SIZE];
109   uint8_t indices_tmp[MAX_PALETTE_BLOCK_WIDTH * MAX_PALETTE_BLOCK_HEIGHT];
110   int16_t *meta_centroids[2] = { centroids, centroids_tmp };
111   uint8_t *meta_indices[2] = { indices, indices_tmp };
112   int i, l = 0, prev_l, best_l = 0;
113   int64_t this_dist;
114 
115   assert(n <= MAX_PALETTE_BLOCK_WIDTH * MAX_PALETTE_BLOCK_HEIGHT);
116 
117 #if AV1_K_MEANS_DIM == 1
118   av1_calc_indices_dim1(data, centroids, indices, &this_dist, n, k);
119 #else
120   av1_calc_indices_dim2(data, centroids, indices, &this_dist, n, k);
121 #endif
122 
123   for (i = 0; i < max_itr; ++i) {
124     const int64_t prev_dist = this_dist;
125     prev_l = l;
126     l = (l == 1) ? 0 : 1;
127 
128     RENAME(calc_centroids)(data, meta_centroids[l], meta_indices[prev_l], n, k);
129     if (!memcmp(meta_centroids[l], meta_centroids[prev_l],
130                 sizeof(centroids[0]) * k * AV1_K_MEANS_DIM)) {
131       break;
132     }
133 #if AV1_K_MEANS_DIM == 1
134     av1_calc_indices_dim1(data, meta_centroids[l], meta_indices[l], &this_dist,
135                           n, k);
136 #else
137     av1_calc_indices_dim2(data, meta_centroids[l], meta_indices[l], &this_dist,
138                           n, k);
139 #endif
140 
141     if (this_dist > prev_dist) {
142       best_l = prev_l;
143       break;
144     }
145   }
146   if (i == max_itr) best_l = l;
147   if (best_l != 0) {
148     memcpy(centroids, meta_centroids[1],
149            sizeof(centroids[0]) * k * AV1_K_MEANS_DIM);
150     memcpy(indices, meta_indices[1], sizeof(indices[0]) * n);
151   }
152 }
153 #undef RENAME_
154 #undef RENAME
155 #undef K_MEANS_RENAME_C
156 #undef RENAME_C_
157 #undef RENAME_C
158