diff --git a/av1/common/av1_rtcd_defs.pl b/av1/common/av1_rtcd_defs.pl
index b5baaae..ba1dcbb 100644
--- a/av1/common/av1_rtcd_defs.pl
+++ b/av1/common/av1_rtcd_defs.pl
@@ -412,10 +412,10 @@
 
   add_proto qw/void av1_quantize_b/, "const tran_low_t *coeff_ptr, intptr_t n_coeffs, const int16_t *zbin_ptr, const int16_t *round_ptr, const int16_t *quant_ptr, const int16_t *quant_shift_ptr, tran_low_t *qcoeff_ptr, tran_low_t *dqcoeff_ptr, const int16_t *dequant_ptr, uint16_t *eob_ptr, const int16_t *scan, const int16_t *iscan, const qm_val_t * qm_ptr, const qm_val_t * iqm_ptr, int log_scale";
 
-  add_proto qw/void av1_calc_indices_dim1/, "const int *data, const int *centroids, uint8_t *indices, int64_t *total_dist, int n, int k";
+  add_proto qw/void av1_calc_indices_dim1/, "const int16_t *data, const int16_t *centroids, uint8_t *indices, int64_t *total_dist, int n, int k";
   specialize qw/av1_calc_indices_dim1 sse2 avx2/;
 
-  add_proto qw/void av1_calc_indices_dim2/, "const int *data, const int *centroids, uint8_t *indices, int64_t *total_dist, int n, int k";
+  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/;
 
   # ENCODEMB INVOKE
diff --git a/av1/encoder/block.h b/av1/encoder/block.h
index 3dec881..4185798 100644
--- a/av1/encoder/block.h
+++ b/av1/encoder/block.h
@@ -326,7 +326,7 @@
   //! The best color map found.
   uint8_t best_palette_color_map[MAX_PALETTE_SQUARE];
   //! A temporary buffer used for k-means clustering.
-  int kmeans_data_buf[2 * MAX_PALETTE_SQUARE];
+  int16_t kmeans_data_buf[2 * MAX_PALETTE_SQUARE];
 } PALETTE_BUFFER;
 
 /*! \brief Contains buffers used by av1_compound_type_rd()
diff --git a/av1/encoder/k_means_template.h b/av1/encoder/k_means_template.h
index 0cf72b7..31ffdcf 100644
--- a/av1/encoder/k_means_template.h
+++ b/av1/encoder/k_means_template.h
@@ -28,7 +28,7 @@
 // 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 int *p1, const int *p2) {
+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
@@ -41,7 +41,7 @@
 #endif
 }
 
-void RENAME(av1_calc_indices)(const int *data, const int *centroids,
+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;
@@ -67,20 +67,22 @@
   }
 }
 
-static void RENAME(calc_centroids)(const int *data, int *centroids,
+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, 0, sizeof(centroids[0]) * k * AV1_K_MEANS_DIM);
+  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[index * AV1_K_MEANS_DIM + j] += data[i * AV1_K_MEANS_DIM + j];
+      centroids_sum[index * AV1_K_MEANS_DIM + j] +=
+          data[i * AV1_K_MEANS_DIM + j];
     }
   }
 
@@ -92,17 +94,17 @@
     } else {
       for (j = 0; j < AV1_K_MEANS_DIM; ++j) {
         centroids[i * AV1_K_MEANS_DIM + j] =
-            DIVIDE_AND_ROUND(centroids[i * AV1_K_MEANS_DIM + j], count[i]);
+            DIVIDE_AND_ROUND(centroids_sum[i * AV1_K_MEANS_DIM + j], count[i]);
       }
     }
   }
 }
 
-void RENAME(av1_k_means)(const int *data, int *centroids, uint8_t *indices,
-                         int n, int k, int max_itr) {
-  int centroids_tmp[AV1_K_MEANS_DIM * PALETTE_MAX_SIZE];
+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];
-  int *meta_centroids[2] = { centroids, centroids_tmp };
+  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;
diff --git a/av1/encoder/palette.c b/av1/encoder/palette.c
index 4375175..9c3d407 100644
--- a/av1/encoder/palette.c
+++ b/av1/encoder/palette.c
@@ -31,14 +31,14 @@
 #include "av1/encoder/k_means_template.h"
 #undef AV1_K_MEANS_DIM
 
-static int int_comparer(const void *a, const void *b) {
-  return (*(int *)a - *(int *)b);
+static int int16_comparer(const void *a, const void *b) {
+  return (*(int16_t *)a - *(int16_t *)b);
 }
 
-int av1_remove_duplicates(int *centroids, int num_centroids) {
+int av1_remove_duplicates(int16_t *centroids, int num_centroids) {
   int num_unique;  // number of unique centroids
   int i;
-  qsort(centroids, num_centroids, sizeof(*centroids), int_comparer);
+  qsort(centroids, num_centroids, sizeof(*centroids), int16_comparer);
   // Remove duplicates.
   num_unique = 1;
   for (i = 1; i < num_centroids; ++i) {
@@ -189,14 +189,14 @@
 // TODO(huisu): Try other schemes to improve compression.
 static AOM_INLINE void optimize_palette_colors(uint16_t *color_cache,
                                                int n_cache, int n_colors,
-                                               int stride, int *centroids,
+                                               int stride, int16_t *centroids,
                                                int bit_depth) {
   if (n_cache <= 0) return;
   for (int i = 0; i < n_colors * stride; i += stride) {
-    int min_diff = abs(centroids[i] - (int)color_cache[0]);
+    int min_diff = abs((int)centroids[i] - (int)color_cache[0]);
     int idx = 0;
     for (int j = 1; j < n_cache; ++j) {
-      const int this_diff = abs(centroids[i] - color_cache[j]);
+      const int this_diff = abs((int)centroids[i] - (int)color_cache[j]);
       if (this_diff < min_diff) {
         min_diff = this_diff;
         idx = j;
@@ -216,8 +216,8 @@
  */
 static AOM_INLINE void palette_rd_y(
     const AV1_COMP *const cpi, MACROBLOCK *x, MB_MODE_INFO *mbmi,
-    BLOCK_SIZE bsize, int dc_mode_cost, const int *data, int *centroids, int n,
-    uint16_t *color_cache, int n_cache, bool do_header_rd_based_gating,
+    BLOCK_SIZE bsize, int dc_mode_cost, const int16_t *data, int16_t *centroids,
+    int n, uint16_t *color_cache, int n_cache, bool do_header_rd_based_gating,
     MB_MODE_INFO *best_mbmi, uint8_t *best_palette_color_map, int64_t *best_rd,
     int *rate, int *rate_tokenonly, int64_t *distortion, uint8_t *skippable,
     int *beat_best_rd, PICK_MODE_CONTEXT *ctx, uint8_t *blk_skip,
@@ -324,14 +324,14 @@
 // returns the best number of colors found.
 static AOM_INLINE int perform_top_color_palette_search(
     const AV1_COMP *const cpi, MACROBLOCK *x, MB_MODE_INFO *mbmi,
-    BLOCK_SIZE bsize, int dc_mode_cost, const int *data, int *top_colors,
-    int start_n, int end_n, int step_size, bool do_header_rd_based_gating,
-    int *last_n_searched, uint16_t *color_cache, int n_cache,
-    MB_MODE_INFO *best_mbmi, uint8_t *best_palette_color_map, int64_t *best_rd,
-    int *rate, int *rate_tokenonly, int64_t *distortion, uint8_t *skippable,
-    int *beat_best_rd, PICK_MODE_CONTEXT *ctx, uint8_t *best_blk_skip,
-    uint8_t *tx_type_map, int discount_color_cost) {
-  int centroids[PALETTE_MAX_SIZE];
+    BLOCK_SIZE bsize, int dc_mode_cost, const int16_t *data,
+    int16_t *top_colors, int start_n, int end_n, int step_size,
+    bool do_header_rd_based_gating, int *last_n_searched, uint16_t *color_cache,
+    int n_cache, MB_MODE_INFO *best_mbmi, uint8_t *best_palette_color_map,
+    int64_t *best_rd, int *rate, int *rate_tokenonly, int64_t *distortion,
+    uint8_t *skippable, int *beat_best_rd, PICK_MODE_CONTEXT *ctx,
+    uint8_t *best_blk_skip, uint8_t *tx_type_map, int discount_color_cost) {
+  int16_t centroids[PALETTE_MAX_SIZE];
   int n = start_n;
   int top_color_winner = end_n;
   /* clang-format off */
@@ -371,7 +371,7 @@
 // returns the best number of colors found.
 static AOM_INLINE int perform_k_means_palette_search(
     const AV1_COMP *const cpi, MACROBLOCK *x, MB_MODE_INFO *mbmi,
-    BLOCK_SIZE bsize, int dc_mode_cost, const int *data, int lower_bound,
+    BLOCK_SIZE bsize, int dc_mode_cost, const int16_t *data, int lower_bound,
     int upper_bound, int start_n, int end_n, int step_size,
     bool do_header_rd_based_gating, int *last_n_searched, uint16_t *color_cache,
     int n_cache, MB_MODE_INFO *best_mbmi, uint8_t *best_palette_color_map,
@@ -379,7 +379,7 @@
     uint8_t *skippable, int *beat_best_rd, PICK_MODE_CONTEXT *ctx,
     uint8_t *best_blk_skip, uint8_t *tx_type_map, uint8_t *color_map,
     int data_points, int discount_color_cost) {
-  int centroids[PALETTE_MAX_SIZE];
+  int16_t centroids[PALETTE_MAX_SIZE];
   const int max_itr = 50;
   int n = start_n;
   int top_color_winner = end_n;
@@ -435,16 +435,19 @@
   *step_size = AOMMAX(1, *max_n - *min_n);
 }
 
-static AOM_INLINE void fill_data_and_get_bounds(
-    const uint8_t *src, const int src_stride, const int rows, const int cols,
-    const int is_high_bitdepth, int *data, int *lower_bound, int *upper_bound) {
+static AOM_INLINE void fill_data_and_get_bounds(const uint8_t *src,
+                                                const int src_stride,
+                                                const int rows, const int cols,
+                                                const int is_high_bitdepth,
+                                                int16_t *data, int *lower_bound,
+                                                int *upper_bound) {
   if (is_high_bitdepth) {
     const uint16_t *src_ptr = CONVERT_TO_SHORTPTR(src);
     *lower_bound = *upper_bound = src_ptr[0];
     for (int r = 0; r < rows; ++r) {
       for (int c = 0; c < cols; ++c) {
         const int val = src_ptr[c];
-        data[c] = val;
+        data[c] = (int16_t)val;
         *lower_bound = AOMMIN(*lower_bound, val);
         *upper_bound = AOMMAX(*upper_bound, val);
       }
@@ -459,7 +462,7 @@
   for (int r = 0; r < rows; ++r) {
     for (int c = 0; c < cols; ++c) {
       const int val = src[c];
-      data[c] = val;
+      data[c] = (int16_t)val;
       *lower_bound = AOMMIN(*lower_bound, val);
       *upper_bound = AOMMAX(*upper_bound, val);
     }
@@ -487,7 +490,7 @@
 }
 
 static void find_top_colors(const int *const count_buf, int bit_depth,
-                            int n_colors, int *top_colors) {
+                            int n_colors, int16_t *top_colors) {
   // Top color array, serving as a priority queue if more than n_colors are
   // found.
   struct ColorCount top_color_counts[PALETTE_MAX_SIZE] = { { 0 } };
@@ -562,8 +565,8 @@
 
   uint8_t *const color_map = xd->plane[0].color_index_map;
   if (colors_threshold > 1 && colors_threshold <= 64) {
-    int *const data = x->palette_buffer->kmeans_data_buf;
-    int centroids[PALETTE_MAX_SIZE];
+    int16_t *const data = x->palette_buffer->kmeans_data_buf;
+    int16_t centroids[PALETTE_MAX_SIZE];
     int lower_bound, upper_bound;
     fill_data_and_get_bounds(src, src_stride, rows, cols, is_hbd, data,
                              &lower_bound, &upper_bound);
@@ -575,7 +578,7 @@
     const int n_cache = av1_get_palette_cache(xd, 0, color_cache);
 
     // Find the dominant colors, stored in top_colors[].
-    int top_colors[PALETTE_MAX_SIZE] = { 0 };
+    int16_t top_colors[PALETTE_MAX_SIZE] = { 0 };
     find_top_colors(count_buf, bit_depth, AOMMIN(colors, PALETTE_MAX_SIZE),
                     top_colors);
 
@@ -791,8 +794,8 @@
     const int max_itr = 50;
     int lb_u, ub_u, val_u;
     int lb_v, ub_v, val_v;
-    int *const data = x->palette_buffer->kmeans_data_buf;
-    int centroids[2 * PALETTE_MAX_SIZE];
+    int16_t *const data = x->palette_buffer->kmeans_data_buf;
+    int16_t centroids[2 * PALETTE_MAX_SIZE];
 
     uint16_t *src_u16 = CONVERT_TO_SHORTPTR(src_u);
     uint16_t *src_v16 = CONVERT_TO_SHORTPTR(src_v);
@@ -917,8 +920,8 @@
   int src_stride = x->plane[1].src.stride;
   const uint8_t *const src_u = x->plane[1].src.buf;
   const uint8_t *const src_v = x->plane[2].src.buf;
-  int *const data = x->palette_buffer->kmeans_data_buf;
-  int centroids[2 * PALETTE_MAX_SIZE];
+  int16_t *const data = x->palette_buffer->kmeans_data_buf;
+  int16_t centroids[2 * PALETTE_MAX_SIZE];
   uint8_t *const color_map = xd->plane[1].color_index_map;
   int r, c;
   const uint16_t *const src_u16 = CONVERT_TO_SHORTPTR(src_u);
diff --git a/av1/encoder/palette.h b/av1/encoder/palette.h
index 6f33f44..7da863a 100644
--- a/av1/encoder/palette.h
+++ b/av1/encoder/palette.h
@@ -28,10 +28,10 @@
 /*!\cond */
 #define AV1_K_MEANS_RENAME(func, dim) func##_dim##dim##_c
 
-void AV1_K_MEANS_RENAME(av1_k_means, 1)(const int *data, int *centroids,
+void AV1_K_MEANS_RENAME(av1_k_means, 1)(const int16_t *data, int16_t *centroids,
                                         uint8_t *indices, int n, int k,
                                         int max_itr);
-void AV1_K_MEANS_RENAME(av1_k_means, 2)(const int *data, int *centroids,
+void AV1_K_MEANS_RENAME(av1_k_means, 2)(const int16_t *data, int16_t *centroids,
                                         uint8_t *indices, int n, int k,
                                         int max_itr);
 /*!\endcond */
@@ -51,8 +51,9 @@
  *
  * \remark Returns nothing, but saves each data's cluster index in \a indices.
  */
-static INLINE void av1_calc_indices(const int *data, const int *centroids,
-                                    uint8_t *indices, int n, int k, int dim) {
+static INLINE void av1_calc_indices(const int16_t *data,
+                                    const int16_t *centroids, uint8_t *indices,
+                                    int n, int k, int dim) {
   assert(n > 0);
   assert(k > 0);
   if (dim == 1) {
@@ -84,7 +85,7 @@
  *
  * \attention The output centroids are rounded off to nearest integers.
  */
-static INLINE void av1_k_means(const int *data, int *centroids,
+static INLINE void av1_k_means(const int16_t *data, int16_t *centroids,
                                uint8_t *indices, int n, int k, int dim,
                                int max_itr) {
   assert(n > 0);
@@ -110,7 +111,7 @@
  * \attention The centroids should be rounded to integers before calling this
  * method.
  */
-int av1_remove_duplicates(int *centroids, int num_centroids);
+int av1_remove_duplicates(int16_t *centroids, int num_centroids);
 
 /*!\brief Checks what colors are in the color cache.
  *
diff --git a/av1/encoder/x86/av1_k_means_avx2.c b/av1/encoder/x86/av1_k_means_avx2.c
index 854b90c..a2db222 100644
--- a/av1/encoder/x86/av1_k_means_avx2.c
+++ b/av1/encoder/x86/av1_k_means_avx2.c
@@ -23,7 +23,55 @@
   return res;
 }
 
-void av1_calc_indices_dim1_avx2(const int *data, const int *centroids,
+void av1_calc_indices_dim1_avx2(const int16_t *data, const int16_t *centroids,
+                                uint8_t *indices, int64_t *total_dist, int n,
+                                int k) {
+  __m256i dist[PALETTE_MAX_SIZE];
+  const __m256i v_zero = _mm256_setzero_si256();
+  __m256i sum = _mm256_setzero_si256();
+
+  for (int i = 0; i < n; i += 16) {
+    __m256i ind = _mm256_loadu_si256((__m256i *)data);
+    for (int j = 0; j < k; j++) {
+      __m256i cent = _mm256_set1_epi16(centroids[j]);
+      __m256i d1 = _mm256_sub_epi16(ind, cent);
+      dist[j] = _mm256_abs_epi16(d1);
+    }
+
+    ind = _mm256_setzero_si256();
+    for (int j = 1; j < k; j++) {
+      __m256i cmp = _mm256_cmpgt_epi16(dist[0], dist[j]);
+      dist[0] = _mm256_min_epi16(dist[0], dist[j]);
+      __m256i ind1 = _mm256_set1_epi16(j);
+      ind = _mm256_or_si256(_mm256_andnot_si256(cmp, ind),
+                            _mm256_and_si256(cmp, ind1));
+    }
+
+    __m256i p1 = _mm256_packus_epi16(ind, v_zero);
+    __m256i px = _mm256_permute4x64_epi64(p1, 0x58);
+    __m128i d1 = _mm256_extracti128_si256(px, 0);
+
+    _mm_storeu_si128((__m128i *)indices, d1);
+
+    if (total_dist) {
+      // Square, convert to 32 bit and add together.
+      dist[0] = _mm256_madd_epi16(dist[0], dist[0]);
+      // Convert to 64 bit and add to sum.
+      const __m256i dist1 = _mm256_unpacklo_epi32(dist[0], v_zero);
+      const __m256i dist2 = _mm256_unpackhi_epi32(dist[0], v_zero);
+      sum = _mm256_add_epi64(sum, dist1);
+      sum = _mm256_add_epi64(sum, dist2);
+    }
+
+    indices += 16;
+    data += 16;
+  }
+  if (total_dist) {
+    *total_dist = k_means_horizontal_sum_avx2(sum);
+  }
+}
+
+void av1_calc_indices_dim2_avx2(const int16_t *data, const int16_t *centroids,
                                 uint8_t *indices, int64_t *total_dist, int n,
                                 int k) {
   __m256i dist[PALETTE_MAX_SIZE];
@@ -33,77 +81,18 @@
   for (int i = 0; i < n; i += 8) {
     __m256i ind = _mm256_loadu_si256((__m256i *)data);
     for (int j = 0; j < k; j++) {
-      __m256i cent = _mm256_set1_epi32(centroids[j]);
-      __m256i d1 = _mm256_sub_epi32(ind, cent);
-      dist[j] = _mm256_abs_epi32(d1);
+      const int16_t cx = centroids[2 * j], cy = centroids[2 * j + 1];
+      const __m256i cent = _mm256_set_epi16(cy, cx, cy, cx, cy, cx, cy, cx, cy,
+                                            cx, cy, cx, cy, cx, cy, cx);
+      const __m256i d1 = _mm256_sub_epi16(ind, cent);
+      dist[j] = _mm256_madd_epi16(d1, d1);
     }
 
     ind = _mm256_setzero_si256();
     for (int j = 1; j < k; j++) {
       __m256i cmp = _mm256_cmpgt_epi32(dist[0], dist[j]);
-      __m256i dist1 = _mm256_andnot_si256(cmp, dist[0]);
-      __m256i dist2 = _mm256_and_si256(cmp, dist[j]);
-      dist[0] = _mm256_or_si256(dist1, dist2);
-      __m256i ind1 = _mm256_set1_epi32(j);
-      ind = _mm256_or_si256(_mm256_andnot_si256(cmp, ind),
-                            _mm256_and_si256(cmp, ind1));
-    }
-
-    __m256i p1 = _mm256_packus_epi32(ind, v_zero);
-    __m256i px = _mm256_permute4x64_epi64(p1, 0x58);
-    __m256i p2 = _mm256_packus_epi16(px, v_zero);
-    __m128i d1 = _mm256_extracti128_si256(p2, 0);
-
-    _mm_storel_epi64((__m128i *)indices, d1);
-
-    if (total_dist) {
-      // Square, convert to 64 bit and add to sum.
-      dist[0] = _mm256_mullo_epi32(dist[0], dist[0]);
-      const __m256i dist1 = _mm256_unpacklo_epi32(dist[0], v_zero);
-      const __m256i dist2 = _mm256_unpackhi_epi32(dist[0], v_zero);
-      sum = _mm256_add_epi64(sum, dist1);
-      sum = _mm256_add_epi64(sum, dist2);
-    }
-
-    indices += 8;
-    data += 8;
-  }
-  if (total_dist) {
-    *total_dist = k_means_horizontal_sum_avx2(sum);
-  }
-}
-
-void av1_calc_indices_dim2_avx2(const int *data, const int *centroids,
-                                uint8_t *indices, int64_t *total_dist, int n,
-                                int k) {
-  __m256i dist[PALETTE_MAX_SIZE];
-  const __m256i v_zero = _mm256_setzero_si256();
-  const __m256i v_permute = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
-  __m256i sum = _mm256_setzero_si256();
-
-  for (int i = 0; i < n; i += 8) {
-    __m256i ind1 = _mm256_loadu_si256((__m256i *)data);
-    __m256i ind2 = _mm256_loadu_si256((__m256i *)(data + 8));
-    for (int j = 0; j < k; j++) {
-      __m128i cent0 = _mm_loadl_epi64((__m128i const *)&centroids[2 * j]);
-      __m256i cent1 = _mm256_inserti128_si256(v_zero, cent0, 0);
-      cent1 = _mm256_inserti128_si256(cent1, cent0, 1);
-      __m256i cent = _mm256_unpacklo_epi64(cent1, cent1);
-      __m256i d1 = _mm256_sub_epi32(ind1, cent);
-      __m256i d2 = _mm256_sub_epi32(ind2, cent);
-      __m256i d3 = _mm256_mullo_epi32(d1, d1);
-      __m256i d4 = _mm256_mullo_epi32(d2, d2);
-      __m256i d5 = _mm256_hadd_epi32(d3, d4);
-      dist[j] = _mm256_permutevar8x32_epi32(d5, v_permute);
-    }
-
-    __m256i ind = _mm256_setzero_si256();
-    for (int j = 1; j < k; j++) {
-      __m256i cmp = _mm256_cmpgt_epi32(dist[0], dist[j]);
-      __m256i dist1 = _mm256_andnot_si256(cmp, dist[0]);
-      __m256i dist2 = _mm256_and_si256(cmp, dist[j]);
-      dist[0] = _mm256_or_si256(dist1, dist2);
-      ind1 = _mm256_set1_epi32(j);
+      dist[0] = _mm256_min_epi32(dist[0], dist[j]);
+      const __m256i ind1 = _mm256_set1_epi32(j);
       ind = _mm256_or_si256(_mm256_andnot_si256(cmp, ind),
                             _mm256_and_si256(cmp, ind1));
     }
diff --git a/av1/encoder/x86/av1_k_means_sse2.c b/av1/encoder/x86/av1_k_means_sse2.c
index d2c7796..a284fa9 100644
--- a/av1/encoder/x86/av1_k_means_sse2.c
+++ b/av1/encoder/x86/av1_k_means_sse2.c
@@ -22,68 +22,49 @@
   return res;
 }
 
-static __m128i absolute_diff_epi32(__m128i a, __m128i b) {
-  const __m128i diff1 = _mm_sub_epi32(a, b);
-  const __m128i diff2 = _mm_sub_epi32(b, a);
-  const __m128i cmp = _mm_cmpgt_epi32(diff1, diff2);
-  const __m128i masked1 = _mm_and_si128(cmp, diff1);
-  const __m128i masked2 = _mm_andnot_si128(cmp, diff2);
-  return _mm_or_si128(masked1, masked2);
-}
-
-void av1_calc_indices_dim1_sse2(const int *data, const int *centroids,
+void av1_calc_indices_dim1_sse2(const int16_t *data, const int16_t *centroids,
                                 uint8_t *indices, int64_t *total_dist, int n,
                                 int k) {
   const __m128i v_zero = _mm_setzero_si128();
-  int l = 1;
   __m128i dist[PALETTE_MAX_SIZE];
-  __m128i ind[2];
   __m128i sum = _mm_setzero_si128();
 
-  for (int i = 0; i < n; i += 4) {
-    l = (l == 0) ? 1 : 0;
-    ind[l] = _mm_loadu_si128((__m128i *)data);
+  for (int i = 0; i < n; i += 8) {
+    __m128i in = _mm_loadu_si128((__m128i *)data);
     for (int j = 0; j < k; j++) {
-      __m128i cent = _mm_set1_epi32(centroids[j]);
-      dist[j] = absolute_diff_epi32(ind[l], cent);
+      __m128i cent = _mm_set1_epi16(centroids[j]);
+      __m128i d1 = _mm_sub_epi16(in, cent);
+      __m128i d2 = _mm_sub_epi16(cent, in);
+      dist[j] = _mm_max_epi16(d1, d2);
     }
 
-    ind[l] = _mm_setzero_si128();
+    __m128i ind = _mm_setzero_si128();
     for (int j = 1; j < k; j++) {
-      __m128i cmp = _mm_cmpgt_epi32(dist[0], dist[j]);
-      __m128i dist1 = _mm_andnot_si128(cmp, dist[0]);
-      __m128i dist2 = _mm_and_si128(cmp, dist[j]);
-      dist[0] = _mm_or_si128(dist1, dist2);
-      __m128i ind1 = _mm_set1_epi32(j);
-      ind[l] =
-          _mm_or_si128(_mm_andnot_si128(cmp, ind[l]), _mm_and_si128(cmp, ind1));
+      __m128i cmp = _mm_cmpgt_epi16(dist[0], dist[j]);
+      dist[0] = _mm_min_epi16(dist[0], dist[j]);
+      __m128i ind1 = _mm_set1_epi16(j);
+      ind = _mm_or_si128(_mm_andnot_si128(cmp, ind), _mm_and_si128(cmp, ind1));
     }
-    ind[l] = _mm_packus_epi16(ind[l], v_zero);
     if (total_dist) {
-      // Square and convert to 32 bit.
-      const __m128i d1 = _mm_packs_epi32(dist[0], v_zero);
-      const __m128i d2 = _mm_mullo_epi16(d1, d1);
-      const __m128i d3 = _mm_mulhi_epi16(d1, d1);
-      dist[0] = _mm_unpacklo_epi16(d2, d3);
+      // Square, convert to 32 bit and add together.
+      dist[0] = _mm_madd_epi16(dist[0], dist[0]);
       // Convert to 64 bit and add to sum.
       const __m128i dist1 = _mm_unpacklo_epi32(dist[0], v_zero);
       const __m128i dist2 = _mm_unpackhi_epi32(dist[0], v_zero);
       sum = _mm_add_epi64(sum, dist1);
       sum = _mm_add_epi64(sum, dist2);
     }
-    if (l == 1) {
-      __m128i p2 = _mm_packus_epi16(_mm_unpacklo_epi64(ind[0], ind[1]), v_zero);
-      _mm_storel_epi64((__m128i *)indices, p2);
-      indices += 8;
-    }
-    data += 4;
+    __m128i p2 = _mm_packus_epi16(ind, v_zero);
+    _mm_storel_epi64((__m128i *)indices, p2);
+    indices += 8;
+    data += 8;
   }
   if (total_dist) {
     *total_dist = k_means_horizontal_sum_sse2(sum);
   }
 }
 
-void av1_calc_indices_dim2_sse2(const int *data, const int *centroids,
+void av1_calc_indices_dim2_sse2(const int16_t *data, const int16_t *centroids,
                                 uint8_t *indices, int64_t *total_dist, int n,
                                 int k) {
   const __m128i v_zero = _mm_setzero_si128();
@@ -95,19 +76,11 @@
   for (int i = 0; i < n; i += 4) {
     l = (l == 0) ? 1 : 0;
     __m128i ind1 = _mm_loadu_si128((__m128i *)data);
-    __m128i ind2 = _mm_loadu_si128((__m128i *)(data + 4));
-    __m128i indl = _mm_unpacklo_epi32(ind1, ind2);
-    __m128i indh = _mm_unpackhi_epi32(ind1, ind2);
-    ind1 = _mm_unpacklo_epi32(indl, indh);
-    ind2 = _mm_unpackhi_epi32(indl, indh);
     for (int j = 0; j < k; j++) {
-      __m128i cent0 = _mm_set1_epi32(centroids[2 * j]);
-      __m128i cent1 = _mm_set1_epi32(centroids[2 * j + 1]);
-      __m128i d1 = absolute_diff_epi32(ind1, cent0);
-      __m128i d2 = absolute_diff_epi32(ind2, cent1);
-      __m128i d3 = _mm_madd_epi16(d1, d1);
-      __m128i d4 = _mm_madd_epi16(d2, d2);
-      dist[j] = _mm_add_epi32(d3, d4);
+      const int16_t cx = centroids[2 * j], cy = centroids[2 * j + 1];
+      const __m128i cent = _mm_set_epi16(cy, cx, cy, cx, cy, cx, cy, cx);
+      const __m128i d1 = _mm_sub_epi16(ind1, cent);
+      dist[j] = _mm_madd_epi16(d1, d1);
     }
 
     ind[l] = _mm_setzero_si128();
diff --git a/test/av1_k_means_test.cc b/test/av1_k_means_test.cc
index 5b6c22e..221dd10 100644
--- a/test/av1_k_means_test.cc
+++ b/test/av1_k_means_test.cc
@@ -28,12 +28,12 @@
 #include "third_party/googletest/src/googletest/include/gtest/gtest.h"
 
 namespace AV1Kmeans {
-typedef void (*av1_calc_indices_dim1_func)(const int *data,
-                                           const int *centroids,
+typedef void (*av1_calc_indices_dim1_func)(const int16_t *data,
+                                           const int16_t *centroids,
                                            uint8_t *indices,
                                            int64_t *total_dist, int n, int k);
-typedef void (*av1_calc_indices_dim2_func)(const int *data,
-                                           const int *centroids,
+typedef void (*av1_calc_indices_dim2_func)(const int16_t *data,
+                                           const int16_t *centroids,
                                            uint8_t *indices,
                                            int64_t *total_dist, int n, int k);
 
@@ -68,8 +68,8 @@
   }
 
   libaom_test::ACMRandom rnd_;
-  int data_[4096];
-  int centroids_[8];
+  int16_t data_[4096];
+  int16_t centroids_[8];
   uint8_t indices1_[4096];
   uint8_t indices2_[4096];
 };
@@ -178,8 +178,8 @@
   }
 
   libaom_test::ACMRandom rnd_;
-  int data_[4096 * 2];
-  int centroids_[8 * 2];
+  int16_t data_[4096 * 2];
+  int16_t centroids_[8 * 2];
   uint8_t indices1_[4096];
   uint8_t indices2_[4096];
 };