Take bitdepth into account in SAO decision

Author: Jovasa

src/sao.c

@@ -385,7 +385,7 @@ static void sao_search_edge_sao(const encoder_state_t * const state,
     // Call calc_sao_edge_dir once for luma and twice for chroma.
     for (i = 0; i < buf_cnt; ++i) {
       FILL(cat_sum_cnt, 0);
-      kvz_calc_sao_edge_dir(data[i], recdata[i], edge_class,
+      kvz_calc_sao_edge_dir(state->encoder_control, data[i], recdata[i], edge_class,
                         block_width, block_height, cat_sum_cnt);
 
 
@@ -514,7 +514,7 @@ static void sao_search_best_mode(const encoder_state_t * const state, const kvz_
     unsigned buf_i;
 
     for (buf_i = 0; buf_i < buf_cnt; ++buf_i) {
-      ddistortion += kvz_sao_edge_ddistortion(data[buf_i], recdata[buf_i], 
+      ddistortion += kvz_sao_edge_ddistortion(state->encoder_control, data[buf_i], recdata[buf_i], 
                                           block_width, block_height,
                                           edge_sao.eo_class, &edge_sao.offsets[5 * buf_i]);
     }
@@ -577,7 +577,7 @@ static void sao_search_best_mode(const encoder_state_t * const state, const kvz_
         switch (merge_cand->type) {
           case SAO_TYPE_EDGE:
                 for (buf_i = 0; buf_i < buf_cnt; ++buf_i) {
-                  ddistortion += kvz_sao_edge_ddistortion(data[buf_i], recdata[buf_i],
+                  ddistortion += kvz_sao_edge_ddistortion(state->encoder_control, data[buf_i], recdata[buf_i],
                     block_width, block_height,
                     merge_cand->eo_class, &merge_cand->offsets[5 * buf_i]);
                 }

src/strategies/avx2/sao-avx2.c

@@ -284,7 +284,8 @@ static INLINE __m256i FIX_W32 do_one_edge_ymm(const __m256i a,
   return calc_diff_off_delta(diff_lo, diff_hi, offset, orig);
 }
 
-static int32_t sao_edge_ddistortion_avx2(const uint8_t *orig_data,
+static int32_t sao_edge_ddistortion_avx2(const encoder_control_t* const encoder, 
+                                         const uint8_t *orig_data,
                                          const uint8_t *rec_data,
                                                int32_t  block_width,
                                                int32_t  block_height,
@@ -316,7 +317,8 @@ static int32_t sao_edge_ddistortion_avx2(const uint8_t *orig_data,
   assert(NUM_SAO_EDGE_CATEGORIES == 5);
 
   if (offsets_ok != 0xffff) {
-    return sao_edge_ddistortion_generic(orig_data,
+    return sao_edge_ddistortion_generic(encoder, 
+                                        orig_data,
                                         rec_data,
                                         block_width,
                                         block_height,
@@ -420,7 +422,8 @@ static void FIX_W32 calc_edge_dir_one_ymm(const __m256i  a,
   }
 }
 
-static void calc_sao_edge_dir_avx2(const uint8_t *orig_data,
+static void calc_sao_edge_dir_avx2(const encoder_control_t* const encoder, 
+                                   const uint8_t *orig_data,
                                    const uint8_t *rec_data,
                                          int32_t  eo_class,
                                          int32_t  block_width,

src/strategies/generic/sao-generic.c

@@ -47,7 +47,8 @@
  * \param dir_offsets
  * \param is_chroma  0 for luma, 1 for chroma. Indicates
  */
-static void calc_sao_edge_dir_generic(const kvz_pixel *orig_data,
+static void calc_sao_edge_dir_generic(const encoder_control_t* const encoder, 
+                                      const kvz_pixel *orig_data,
                                       const kvz_pixel *rec_data,
                                       int eo_class,
                                       int block_width,
@@ -62,7 +63,7 @@ static void calc_sao_edge_dir_generic(const kvz_pixel *orig_data,
   // Don't sample the edge pixels because this function doesn't have access to
   // their neighbours.
 
-  
+  const int offset = encoder->bitdepth != 8 ? 1 << (encoder->bitdepth - 9) : 0;
 
   for (y = 1; y < block_height - 1; ++y) {
     for (x = 1; x < block_width - 1; ++x) {
@@ -73,7 +74,7 @@ static void calc_sao_edge_dir_generic(const kvz_pixel *orig_data,
 
       int eo_cat = sao_calc_eo_cat(a, b, c);
 
-      cat_sum_cnt[0][eo_cat] += orig_data[y * block_width + x] - c;
+      cat_sum_cnt[0][eo_cat] += (orig_data[y * block_width + x] - c + offset) >> (encoder->bitdepth - 8);
       cat_sum_cnt[1][eo_cat] += 1;
     }
   }

src/strategies/generic/sao_shared_generics.h

@@ -49,7 +49,8 @@ static int sao_calc_eo_cat(kvz_pixel a, kvz_pixel b, kvz_pixel c)
   return sao_eo_idx_to_eo_category[eo_idx];
 }
 
-static int sao_edge_ddistortion_generic(const kvz_pixel *orig_data,
+static int sao_edge_ddistortion_generic(const encoder_control_t* const encoder, 
+                                        const kvz_pixel *orig_data,
                                         const kvz_pixel *rec_data,
                                               int32_t    block_width,
                                               int32_t    block_height,
@@ -61,6 +62,8 @@ static int sao_edge_ddistortion_generic(const kvz_pixel *orig_data,
   vector2d_t a_ofs = g_sao_edge_offsets[eo_class][0];
   vector2d_t b_ofs = g_sao_edge_offsets[eo_class][1];
 
+  const int bit_offset = encoder->bitdepth != 8 ? 1 << (encoder->bitdepth - 9) : 0;
+
   for (y = 1; y < block_height - 1; y++) {
     for (x = 1; x < block_width - 1; x++) {
       uint32_t c_pos =  y            * block_width + x;
@@ -76,7 +79,7 @@ static int sao_edge_ddistortion_generic(const kvz_pixel *orig_data,
       int32_t offset = offsets[eo_cat];
 
       if (offset != 0) {
-        int32_t diff   = orig - c;
+        int32_t diff   = (orig - c + bit_offset) >> (encoder->bitdepth - 8);
         int32_t delta  = diff - offset;
         int32_t curr   = delta * delta - diff * diff;
 

src/strategies/strategies-sao.h

@@ -46,11 +46,13 @@
 
 
 // Declare function pointers.
-typedef int (sao_edge_ddistortion_func)(const kvz_pixel *orig_data, const kvz_pixel *rec_data,
+typedef int (sao_edge_ddistortion_func)(const encoder_control_t* const encoder, 
+  const kvz_pixel *orig_data, const kvz_pixel *rec_data,
   int block_width, int block_height,
   int eo_class, int offsets[NUM_SAO_EDGE_CATEGORIES]);
 
-typedef void (calc_sao_edge_dir_func)(const kvz_pixel *orig_data, const kvz_pixel *rec_data,
+typedef void (calc_sao_edge_dir_func)(const encoder_control_t* const encoder, 
+  const kvz_pixel *orig_data, const kvz_pixel *rec_data,
   int eo_class, int block_width, int block_height,
   int cat_sum_cnt[2][NUM_SAO_EDGE_CATEGORIES]);