Commit 15aa8c01 authored by Alex Converse's avatar Alex Converse
Browse files

palette: Templatize k_means routines

Drops k_means from 15% of profile to 4.2%.

BUG=aomedia:670

Change-Id: I9c60f024abde9112eec8c32ead482f885ed6e57a
parent 6ae0054c
...@@ -75,6 +75,7 @@ AV1_CX_SRCS-yes += encoder/treewriter.h ...@@ -75,6 +75,7 @@ AV1_CX_SRCS-yes += encoder/treewriter.h
AV1_CX_SRCS-yes += encoder/mcomp.c AV1_CX_SRCS-yes += encoder/mcomp.c
AV1_CX_SRCS-yes += encoder/encoder.c AV1_CX_SRCS-yes += encoder/encoder.c
ifeq ($(CONFIG_PALETTE),yes) ifeq ($(CONFIG_PALETTE),yes)
AV1_CX_SRCS-yes += encoder/k_means_template.h
AV1_CX_SRCS-yes += encoder/palette.h AV1_CX_SRCS-yes += encoder/palette.h
AV1_CX_SRCS-yes += encoder/palette.c AV1_CX_SRCS-yes += encoder/palette.c
endif endif
......
/*
* 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 <string.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)
static float RENAME(calc_dist)(const float *p1, const float *p2) {
float dist = 0;
int i;
for (i = 0; i < AV1_K_MEANS_DIM; ++i) {
const float diff = p1[i] - p2[i];
dist += diff * diff;
}
return dist;
}
void RENAME(av1_calc_indices)(const float *data, const float *centroids,
uint8_t *indices, int n, int k) {
int i, j;
for (i = 0; i < n; ++i) {
float min_dist = RENAME(calc_dist)(data + i * AV1_K_MEANS_DIM, centroids);
indices[i] = 0;
for (j = 1; j < k; ++j) {
const float 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;
}
}
}
}
static void RENAME(calc_centroids)(const float *data, float *centroids,
const uint8_t *indices, int n, int k) {
int i, j, index;
int count[PALETTE_MAX_SIZE];
unsigned int rand_state = (unsigned int)data[0];
assert(n <= 32768);
memset(count, 0, sizeof(count[0]) * k);
memset(centroids, 0, sizeof(centroids[0]) * k * AV1_K_MEANS_DIM);
for (i = 0; i < n; ++i) {
index = indices[i];
assert(index < k);
++count[index];
for (j = 0; j < AV1_K_MEANS_DIM; ++j) {
centroids[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 {
const float norm = 1.0f / count[i];
for (j = 0; j < AV1_K_MEANS_DIM; ++j)
centroids[i * AV1_K_MEANS_DIM + j] *= norm;
}
}
// Round to nearest integers.
for (i = 0; i < k * AV1_K_MEANS_DIM; ++i) {
centroids[i] = roundf(centroids[i]);
}
}
static float RENAME(calc_total_dist)(const float *data, const float *centroids,
const uint8_t *indices, int n, int k) {
float dist = 0;
int i;
(void)k;
for (i = 0; i < n; ++i)
dist += RENAME(calc_dist)(data + i * AV1_K_MEANS_DIM,
centroids + indices[i] * AV1_K_MEANS_DIM);
return dist;
}
void RENAME(av1_k_means)(const float *data, float *centroids, uint8_t *indices,
int n, int k, int max_itr) {
int i;
float this_dist;
float pre_centroids[2 * PALETTE_MAX_SIZE];
uint8_t pre_indices[MAX_SB_SQUARE];
RENAME(av1_calc_indices)(data, centroids, indices, n, k);
this_dist = RENAME(calc_total_dist)(data, centroids, indices, n, k);
for (i = 0; i < max_itr; ++i) {
const float pre_dist = this_dist;
memcpy(pre_centroids, centroids,
sizeof(pre_centroids[0]) * k * AV1_K_MEANS_DIM);
memcpy(pre_indices, indices, sizeof(pre_indices[0]) * n);
RENAME(calc_centroids)(data, centroids, indices, n, k);
RENAME(av1_calc_indices)(data, centroids, indices, n, k);
this_dist = RENAME(calc_total_dist)(data, centroids, indices, n, k);
if (this_dist > pre_dist) {
memcpy(centroids, pre_centroids,
sizeof(pre_centroids[0]) * k * AV1_K_MEANS_DIM);
memcpy(indices, pre_indices, sizeof(pre_indices[0]) * n);
break;
}
if (!memcmp(centroids, pre_centroids,
sizeof(pre_centroids[0]) * k * AV1_K_MEANS_DIM))
break;
}
}
#undef RENAME_
#undef RENAME
...@@ -16,109 +16,12 @@ ...@@ -16,109 +16,12 @@
#include "av1/encoder/palette.h" #include "av1/encoder/palette.h"
#include "av1/encoder/random.h" #include "av1/encoder/random.h"
static float calc_dist(const float *p1, const float *p2, int dim) { #define AV1_K_MEANS_DIM 1
float dist = 0; #include "av1/encoder/k_means_template.h"
int i; #undef AV1_K_MEANS_DIM
for (i = 0; i < dim; ++i) { #define AV1_K_MEANS_DIM 2
const float diff = p1[i] - p2[i]; #include "av1/encoder/k_means_template.h"
dist += diff * diff; #undef AV1_K_MEANS_DIM
}
return dist;
}
void av1_calc_indices(const float *data, const float *centroids,
uint8_t *indices, int n, int k, int dim) {
int i, j;
for (i = 0; i < n; ++i) {
float min_dist = calc_dist(data + i * dim, centroids, dim);
indices[i] = 0;
for (j = 1; j < k; ++j) {
const float this_dist =
calc_dist(data + i * dim, centroids + j * dim, dim);
if (this_dist < min_dist) {
min_dist = this_dist;
indices[i] = j;
}
}
}
}
static void calc_centroids(const float *data, float *centroids,
const uint8_t *indices, int n, int k, int dim) {
int i, j, index;
int count[PALETTE_MAX_SIZE];
unsigned int rand_state = (unsigned int)data[0];
assert(n <= 32768);
memset(count, 0, sizeof(count[0]) * k);
memset(centroids, 0, sizeof(centroids[0]) * k * dim);
for (i = 0; i < n; ++i) {
index = indices[i];
assert(index < k);
++count[index];
for (j = 0; j < dim; ++j) {
centroids[index * dim + j] += data[i * dim + j];
}
}
for (i = 0; i < k; ++i) {
if (count[i] == 0) {
memcpy(centroids + i * dim, data + (lcg_rand16(&rand_state) % n) * dim,
sizeof(centroids[0]) * dim);
} else {
const float norm = 1.0f / count[i];
for (j = 0; j < dim; ++j) centroids[i * dim + j] *= norm;
}
}
// Round to nearest integers.
for (i = 0; i < k * dim; ++i) {
centroids[i] = roundf(centroids[i]);
}
}
static float calc_total_dist(const float *data, const float *centroids,
const uint8_t *indices, int n, int k, int dim) {
float dist = 0;
int i;
(void)k;
for (i = 0; i < n; ++i)
dist += calc_dist(data + i * dim, centroids + indices[i] * dim, dim);
return dist;
}
void av1_k_means(const float *data, float *centroids, uint8_t *indices, int n,
int k, int dim, int max_itr) {
int i;
float this_dist;
float pre_centroids[2 * PALETTE_MAX_SIZE];
uint8_t pre_indices[MAX_SB_SQUARE];
av1_calc_indices(data, centroids, indices, n, k, dim);
this_dist = calc_total_dist(data, centroids, indices, n, k, dim);
for (i = 0; i < max_itr; ++i) {
const float pre_dist = this_dist;
memcpy(pre_centroids, centroids, sizeof(pre_centroids[0]) * k * dim);
memcpy(pre_indices, indices, sizeof(pre_indices[0]) * n);
calc_centroids(data, centroids, indices, n, k, dim);
av1_calc_indices(data, centroids, indices, n, k, dim);
this_dist = calc_total_dist(data, centroids, indices, n, k, dim);
if (this_dist > pre_dist) {
memcpy(centroids, pre_centroids, sizeof(pre_centroids[0]) * k * dim);
memcpy(indices, pre_indices, sizeof(pre_indices[0]) * n);
break;
}
if (!memcmp(centroids, pre_centroids, sizeof(pre_centroids[0]) * k * dim))
break;
}
}
static int float_comparer(const void *a, const void *b) { static int float_comparer(const void *a, const void *b) {
const float fa = *(const float *)a; const float fa = *(const float *)a;
......
...@@ -18,17 +18,49 @@ ...@@ -18,17 +18,49 @@
extern "C" { extern "C" {
#endif #endif
#define AV1_K_MEANS_RENAME(func, dim) func##_dim##dim
void AV1_K_MEANS_RENAME(av1_calc_indices, 1)(const float *data,
const float *centroids,
uint8_t *indices, int n, int k);
void AV1_K_MEANS_RENAME(av1_calc_indices, 2)(const float *data,
const float *centroids,
uint8_t *indices, int n, int k);
void AV1_K_MEANS_RENAME(av1_k_means, 1)(const float *data, float *centroids,
uint8_t *indices, int n, int k,
int max_itr);
void AV1_K_MEANS_RENAME(av1_k_means, 2)(const float *data, float *centroids,
uint8_t *indices, int n, int k,
int max_itr);
// Given 'n' 'data' points and 'k' 'centroids' each of dimension 'dim', // Given 'n' 'data' points and 'k' 'centroids' each of dimension 'dim',
// calculate the centroid 'indices' for the data points. // calculate the centroid 'indices' for the data points.
void av1_calc_indices(const float *data, const float *centroids, static INLINE void av1_calc_indices(const float *data, const float *centroids,
uint8_t *indices, int n, int k, int dim); uint8_t *indices, int n, int k, int dim) {
if (dim == 1) {
AV1_K_MEANS_RENAME(av1_calc_indices, 1)(data, centroids, indices, n, k);
} else if (dim == 2) {
AV1_K_MEANS_RENAME(av1_calc_indices, 2)(data, centroids, indices, n, k);
} else {
assert(0 && "Untemplated k means dimension");
}
}
// Given 'n' 'data' points and an initial guess of 'k' 'centroids' each of // Given 'n' 'data' points and an initial guess of 'k' 'centroids' each of
// dimension 'dim', runs up to 'max_itr' iterations of k-means algorithm to get // dimension 'dim', runs up to 'max_itr' iterations of k-means algorithm to get
// updated 'centroids' and the centroid 'indices' for elements in 'data'. // updated 'centroids' and the centroid 'indices' for elements in 'data'.
// Note: the output centroids are rounded off to nearest integers. // Note: the output centroids are rounded off to nearest integers.
void av1_k_means(const float *data, float *centroids, uint8_t *indices, int n, static INLINE void av1_k_means(const float *data, float *centroids,
int k, int dim, int max_itr); uint8_t *indices, int n, int k, int dim,
int max_itr) {
if (dim == 1) {
AV1_K_MEANS_RENAME(av1_k_means, 1)(data, centroids, indices, n, k, max_itr);
} else if (dim == 2) {
AV1_K_MEANS_RENAME(av1_k_means, 2)(data, centroids, indices, n, k, max_itr);
} else {
assert(0 && "Untemplated k means dimension");
}
}
// Given a list of centroids, returns the unique number of centroids 'k', and // Given a list of centroids, returns the unique number of centroids 'k', and
// puts these unique centroids in first 'k' indices of 'centroids' array. // puts these unique centroids in first 'k' indices of 'centroids' array.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment