1 files changed, 151 insertions, 0 deletions
diff --git a/third_party/aom/av1/encoder/k_means_template.h b/third_party/aom/av1/encoder/k_means_template.h
new file mode 100644
index 0000000000..4be2038a6f
--- /dev/null
+++ b/third_party/aom/av1/encoder/k_means_template.h
@@ -0,0 +1,151 @@
+/*
+ * Copyright (c) 2016, 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 <assert.h>
+#include <stdint.h>
+#include <stdlib.h>
+#include <string.h>
+
+#include "av1/common/blockd.h"
+#include "av1/encoder/palette.h"
+#include "av1/encoder/random.h"
+
+#ifndef AV1_K_MEANS_DIM
+#error "This template requires AV1_K_MEANS_DIM to be defined"
+#endif
+
+#define RENAME_(x, y) AV1_K_MEANS_RENAME(x, y)
+#define RENAME(x) RENAME_(x, AV1_K_MEANS_DIM)
+
+// Though we want to compute the smallest L2 norm, in 1 dimension,
+// it is equivalent to find the smallest L1 norm and then square it.
+// This is preferrable for speed, especially on the SIMD side.
+static int RENAME(calc_dist)(const int16_t *p1, const int16_t *p2) {
+#if AV1_K_MEANS_DIM == 1
+  return abs(p1[0] - p2[0]);
+#else
+  int dist = 0;
+  for (int i = 0; i < AV1_K_MEANS_DIM; ++i) {
+    const int diff = p1[i] - p2[i];
+    dist += diff * diff;
+  }
+  return dist;
+#endif
+}
+
+void RENAME(av1_calc_indices)(const int16_t *data, const int16_t *centroids,
+                              uint8_t *indices, int64_t *dist, int n, int k) {
+  if (dist) {
+    *dist = 0;
+  }
+  for (int i = 0; i < n; ++i) {
+    int min_dist = RENAME(calc_dist)(data + i * AV1_K_MEANS_DIM, centroids);
+    indices[i] = 0;
+    for (int j = 1; j < k; ++j) {
+      const int this_dist = RENAME(calc_dist)(data + i * AV1_K_MEANS_DIM,
+                                              centroids + j * AV1_K_MEANS_DIM);
+      if (this_dist < min_dist) {
+        min_dist = this_dist;
+        indices[i] = j;
+      }
+    }
+    if (dist) {
+#if AV1_K_MEANS_DIM == 1
+      *dist += min_dist * min_dist;
+#else
+      *dist += min_dist;
+#endif
+    }
+  }
+}
+
+static void RENAME(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];
+  unsigned int rand_state = (unsigned int)data[0];
+  assert(n <= 32768);
+  memset(centroids_sum, 0, sizeof(centroids_sum[0]) * k * AV1_K_MEANS_DIM);
+
+  for (i = 0; i < n; ++i) {
+    const int index = indices[i];
+    assert(index < k);
+    ++count[index];
+    for (j = 0; j < AV1_K_MEANS_DIM; ++j) {
+      centroids_sum[index * AV1_K_MEANS_DIM + j] +=
+          data[i * AV1_K_MEANS_DIM + j];
+    }
+  }
+
+  for (i = 0; i < k; ++i) {
+    if (count[i] == 0) {
+      memcpy(centroids + i * AV1_K_MEANS_DIM,
+             data + (lcg_rand16(&rand_state) % n) * AV1_K_MEANS_DIM,
+             sizeof(centroids[0]) * AV1_K_MEANS_DIM);
+    } else {
+      for (j = 0; j < AV1_K_MEANS_DIM; ++j) {
+        centroids[i * AV1_K_MEANS_DIM + j] =
+            DIVIDE_AND_ROUND(centroids_sum[i * AV1_K_MEANS_DIM + j], count[i]);
+      }
+    }
+  }
+}
+
+void RENAME(av1_k_means)(const int16_t *data, int16_t *centroids,
+                         uint8_t *indices, int n, int k, int max_itr) {
+  int16_t centroids_tmp[AV1_K_MEANS_DIM * PALETTE_MAX_SIZE];
+  uint8_t indices_tmp[MAX_PALETTE_BLOCK_WIDTH * MAX_PALETTE_BLOCK_HEIGHT];
+  int16_t *meta_centroids[2] = { centroids, centroids_tmp };
+  uint8_t *meta_indices[2] = { indices, indices_tmp };
+  int i, l = 0, prev_l, best_l = 0;
+  int64_t this_dist;
+
+  assert(n <= MAX_PALETTE_BLOCK_WIDTH * MAX_PALETTE_BLOCK_HEIGHT);
+
+#if AV1_K_MEANS_DIM == 1
+  av1_calc_indices_dim1(data, centroids, indices, &this_dist, n, k);
+#else
+  av1_calc_indices_dim2(data, centroids, indices, &this_dist, n, k);
+#endif
+
+  for (i = 0; i < max_itr; ++i) {
+    const int64_t prev_dist = this_dist;
+    prev_l = l;
+    l = (l == 1) ? 0 : 1;
+
+    RENAME(calc_centroids)(data, meta_centroids[l], meta_indices[prev_l], n, k);
+    if (!memcmp(meta_centroids[l], meta_centroids[prev_l],
+                sizeof(centroids[0]) * k * AV1_K_MEANS_DIM)) {
+      break;
+    }
+#if AV1_K_MEANS_DIM == 1
+    av1_calc_indices_dim1(data, meta_centroids[l], meta_indices[l], &this_dist,
+                          n, k);
+#else
+    av1_calc_indices_dim2(data, meta_centroids[l], meta_indices[l], &this_dist,
+                          n, k);
+#endif
+
+    if (this_dist > prev_dist) {
+      best_l = prev_l;
+      break;
+    }
+  }
+  if (i == max_itr) best_l = l;
+  if (best_l != 0) {
+    memcpy(centroids, meta_centroids[1],
+           sizeof(centroids[0]) * k * AV1_K_MEANS_DIM);
+    memcpy(indices, meta_indices[1], sizeof(indices[0]) * n);
+  }
+}
+#undef RENAME_
+#undef RENAME