diff --git a/av1/common/av1_rtcd_defs.pl b/av1/common/av1_rtcd_defs.pl
index 5a05548..ca299aa 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, int n, int k";
+  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";
   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, int n, int k";
+  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";
   specialize qw/av1_calc_indices_dim2 sse2 avx2/;
 
   # ENCODEMB INVOKE
diff --git a/av1/encoder/k_means_template.h b/av1/encoder/k_means_template.h
index 87677a1..fe3e9ef 100644
--- a/av1/encoder/k_means_template.h
+++ b/av1/encoder/k_means_template.h
@@ -34,7 +34,10 @@
 }
 
 void RENAME(av1_calc_indices)(const int *data, const int *centroids,
-                              uint8_t *indices, int n, int k) {
+                              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;
@@ -46,6 +49,9 @@
         indices[i] = j;
       }
     }
+    if (dist) {
+      dist += min_dist;
+    }
   }
 }
 
@@ -80,17 +86,6 @@
   }
 }
 
-static int64_t RENAME(calc_total_dist)(const int *data, const int *centroids,
-                                       const uint8_t *indices, int n, int k) {
-  int64_t dist = 0;
-  (void)k;
-  for (int 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 int *data, int *centroids, uint8_t *indices,
                          int n, int k, int max_itr) {
   int centroids_tmp[AV1_K_MEANS_DIM * PALETTE_MAX_SIZE];
@@ -103,11 +98,10 @@
   assert(n <= MAX_PALETTE_BLOCK_WIDTH * MAX_PALETTE_BLOCK_HEIGHT);
 
 #if AV1_K_MEANS_DIM - 2
-  av1_calc_indices_dim1(data, centroids, indices, n, k);
+  av1_calc_indices_dim1(data, centroids, indices, &this_dist, n, k);
 #else
-  av1_calc_indices_dim2(data, centroids, indices, n, k);
+  av1_calc_indices_dim2(data, centroids, indices, &this_dist, n, k);
 #endif
-  this_dist = RENAME(calc_total_dist)(data, centroids, indices, n, k);
 
   for (i = 0; i < max_itr; ++i) {
     const int64_t prev_dist = this_dist;
@@ -116,12 +110,12 @@
 
     RENAME(calc_centroids)(data, meta_centroids[l], meta_indices[prev_l], n, k);
 #if AV1_K_MEANS_DIM - 2
-    av1_calc_indices_dim1(data, meta_centroids[l], meta_indices[l], n, k);
+    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], n, k);
+    av1_calc_indices_dim2(data, meta_centroids[l], meta_indices[l], &this_dist,
+                          n, k);
 #endif
-    this_dist =
-        RENAME(calc_total_dist)(data, meta_centroids[l], meta_indices[l], n, k);
 
     if (this_dist > prev_dist) {
       best_l = prev_l;
diff --git a/av1/encoder/palette.h b/av1/encoder/palette.h
index 34d2ddc..6f33f44 100644
--- a/av1/encoder/palette.h
+++ b/av1/encoder/palette.h
@@ -56,9 +56,9 @@
   assert(n > 0);
   assert(k > 0);
   if (dim == 1) {
-    av1_calc_indices_dim1(data, centroids, indices, n, k);
+    av1_calc_indices_dim1(data, centroids, indices, /*total_dist=*/NULL, n, k);
   } else if (dim == 2) {
-    av1_calc_indices_dim2(data, centroids, indices, n, k);
+    av1_calc_indices_dim2(data, centroids, indices, /*total_dist=*/NULL, n, k);
   } else {
     assert(0 && "Untemplated k means dimension");
   }
diff --git a/av1/encoder/x86/av1_k_means_avx2.c b/av1/encoder/x86/av1_k_means_avx2.c
index 759f515..2177358 100644
--- a/av1/encoder/x86/av1_k_means_avx2.c
+++ b/av1/encoder/x86/av1_k_means_avx2.c
@@ -13,10 +13,18 @@
 #include "config/aom_dsp_rtcd.h"
 #include "aom_dsp/x86/synonyms.h"
 
+static int64_t k_means_horizontal_sum_avx2(__m256i a) {
+  int64_t dists[4];
+  _mm256_store_si256((__m256i *)dists, a);
+  return a[0] + a[1] + a[2] + a[3];
+}
+
 void av1_calc_indices_dim1_avx2(const int *data, const int *centroids,
-                                uint8_t *indices, int n, int k) {
+                                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 += 8) {
     __m256i ind = _mm256_loadu_si256((__m256i *)data);
@@ -44,16 +52,29 @@
 
     _mm_storel_epi64((__m128i *)indices, d1);
 
+    if (total_dist) {
+      // Convert to 64 bit and add to sum.
+      const __m256i dist1 = _mm256_unpacklo_epi32(v_zero, dist[0]);
+      const __m256i dist2 = _mm256_unpackhi_epi32(v_zero, dist[0]);
+      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, int n, int k) {
+                                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);
@@ -89,7 +110,18 @@
 
     _mm_storel_epi64((__m128i *)indices, d1);
 
+    if (total_dist) {
+      // Convert to 64 bit and add to sum.
+      const __m256i dist1 = _mm256_unpacklo_epi32(v_zero, dist[0]);
+      const __m256i dist2 = _mm256_unpackhi_epi32(v_zero, dist[0]);
+      sum = _mm256_add_epi64(sum, dist1);
+      sum = _mm256_add_epi64(sum, dist2);
+    }
+
     indices += 8;
     data += 16;
   }
+  if (total_dist) {
+    *total_dist = k_means_horizontal_sum_avx2(sum);
+  }
 }
diff --git a/av1/encoder/x86/av1_k_means_sse2.c b/av1/encoder/x86/av1_k_means_sse2.c
index f03c459..df823a0 100644
--- a/av1/encoder/x86/av1_k_means_sse2.c
+++ b/av1/encoder/x86/av1_k_means_sse2.c
@@ -14,12 +14,22 @@
 #include "config/aom_dsp_rtcd.h"
 #include "aom_dsp/x86/synonyms.h"
 
+static int64_t k_means_horizontal_sum_sse2(__m128i a) {
+  const __m128i sum1 = _mm_unpackhi_epi64(a, a);
+  const __m128i sum2 = _mm_add_epi64(a, sum1);
+  int64_t res;
+  _mm_storel_epi64((__m128i_u *)(&res), sum2);
+  return res;
+}
+
 void av1_calc_indices_dim1_sse2(const int *data, const int *centroids,
-                                uint8_t *indices, int n, int k) {
+                                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;
@@ -44,6 +54,13 @@
           _mm_or_si128(_mm_andnot_si128(cmp, ind[l]), _mm_and_si128(cmp, ind1));
     }
     ind[l] = _mm_packus_epi16(ind[l], v_zero);
+    if (total_dist) {
+      // Convert to 64 bit and add to sum.
+      const __m128i dist1 = _mm_unpacklo_epi32(v_zero, dist[0]);
+      const __m128i dist2 = _mm_unpackhi_epi32(v_zero, dist[0]);
+      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);
@@ -51,14 +68,19 @@
     }
     data += 4;
   }
+  if (total_dist) {
+    *total_dist = k_means_horizontal_sum_sse2(sum);
+  }
 }
 
 void av1_calc_indices_dim2_sse2(const int *data, const int *centroids,
-                                uint8_t *indices, int n, int k) {
+                                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;
@@ -89,6 +111,13 @@
           _mm_or_si128(_mm_andnot_si128(cmp, ind[l]), _mm_and_si128(cmp, ind1));
     }
     ind[l] = _mm_packus_epi16(ind[l], v_zero);
+    if (total_dist) {
+      // Convert to 64 bit and add to sum.
+      const __m128i dist1 = _mm_unpacklo_epi32(v_zero, dist[0]);
+      const __m128i dist2 = _mm_unpackhi_epi32(v_zero, dist[0]);
+      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);
@@ -96,4 +125,7 @@
     }
     data += 8;
   }
+  if (total_dist) {
+    *total_dist = k_means_horizontal_sum_sse2(sum);
+  }
 }
diff --git a/test/av1_k_means_test.cc b/test/av1_k_means_test.cc
index 8a85323..2ef837c 100644
--- a/test/av1_k_means_test.cc
+++ b/test/av1_k_means_test.cc
@@ -30,10 +30,12 @@
 namespace AV1Kmeans {
 typedef void (*av1_calc_indices_dim1_func)(const int *data,
                                            const int *centroids,
-                                           uint8_t *indices, int n, int k);
+                                           uint8_t *indices,
+                                           int64_t *total_dist, int n, int k);
 typedef void (*av1_calc_indices_dim2_func)(const int *data,
                                            const int *centroids,
-                                           uint8_t *indices, int n, int k);
+                                           uint8_t *indices,
+                                           int64_t *total_dist, int n, int k);
 
 typedef std::tuple<av1_calc_indices_dim1_func, BLOCK_SIZE>
     av1_calc_indices_dim1Param;
@@ -92,8 +94,9 @@
   const int w = block_size_wide[bsize];
   const int h = block_size_high[bsize];
   const int n = w * h;
-  av1_calc_indices_dim1_c(data_, centroids_, indices1_, n, k);
-  test_impl(data_, centroids_, indices2_, n, k);
+  av1_calc_indices_dim1_c(data_, centroids_, indices1_, /*total_dist=*/nullptr,
+                          n, k);
+  test_impl(data_, centroids_, indices2_, /*total_dist=*/nullptr, n, k);
 
   ASSERT_EQ(CheckResult(n), true)
       << " block " << bsize << " index " << n << " Centroids " << k;
@@ -113,7 +116,7 @@
     aom_usec_timer_start(&timer);
     av1_calc_indices_dim1_func func = funcs[i];
     for (int j = 0; j < num_loops; ++j) {
-      func(data_, centroids_, indices1_, n, k);
+      func(data_, centroids_, indices1_, /*total_dist=*/nullptr, n, k);
     }
     aom_usec_timer_mark(&timer);
     double time = static_cast<double>(aom_usec_timer_elapsed(&timer));
@@ -200,8 +203,9 @@
   const int w = block_size_wide[bsize];
   const int h = block_size_high[bsize];
   const int n = w * h;
-  av1_calc_indices_dim2_c(data_, centroids_, indices1_, n, k);
-  test_impl(data_, centroids_, indices2_, n, k);
+  av1_calc_indices_dim2_c(data_, centroids_, indices1_, /*total_dist=*/nullptr,
+                          n, k);
+  test_impl(data_, centroids_, indices2_, /*total_dist=*/nullptr, n, k);
 
   ASSERT_EQ(CheckResult(n), true)
       << " block " << bsize << " index " << n << " Centroids " << k;
@@ -221,7 +225,7 @@
     aom_usec_timer_start(&timer);
     av1_calc_indices_dim2_func func = funcs[i];
     for (int j = 0; j < num_loops; ++j) {
-      func(data_, centroids_, indices1_, n, k);
+      func(data_, centroids_, indices1_, nullptr, n, k);
     }
     aom_usec_timer_mark(&timer);
     double time = static_cast<double>(aom_usec_timer_elapsed(&timer));