Make lpf_sb work with loopfilter_level

Make lpf_sb compatible with loopfilter_level, when USE_GUESS_LEVEL = 1.

Filter levels will be selected based on q index and applied for
filtering on Y, U, V planes separately.

Current model only allows to guess one filter level.
Now Y_vert = Y_horz = U = V. In the future, we need to retrain the
model and get filter levels for Y_vert, Y_horz, U and V separately.

When USE_GUESS_LEVEL = 0, lpf_sb can't work with loopfilter_level yet.

Change-Id: Icd774a147c07a4035cf8204a8754b2a99668bbfd
diff --git a/av1/encoder/bitstream.c b/av1/encoder/bitstream.c
index ba271c7..3c305e2 100644
--- a/av1/encoder/bitstream.c
+++ b/av1/encoder/bitstream.c
@@ -2298,11 +2298,26 @@
 #if CONFIG_LPF_SB
   // send filter level for each superblock (64x64)
   if (bsize == cm->sb_size && !USE_GUESS_LEVEL) {
+    int lvl_idx;
     if (mi_row == 0 && mi_col == 0) {
+#if CONFIG_LOOPFILTER_LEVEL
+      aom_write_literal(w, cm->mi[0].mbmi.filt_lvl[0], 6);
+      aom_write_literal(w, cm->mi[0].mbmi.filt_lvl[1], 6);
+      if (cm->mi[0].mbmi.filt_lvl[0] || cm->mi[0].mbmi.filt_lvl[1]) {
+        aom_write_literal(w, cm->mi[0].mbmi.filt_lvl[2], 6);
+        aom_write_literal(w, cm->mi[0].mbmi.filt_lvl[3], 6);
+      }
+      for (lvl_idx = 0; lvl_idx < 4; ++lvl_idx) {
+        cm->mi_grid_visible[0]->mbmi.reuse_sb_lvl[lvl_idx] = 0;
+        cm->mi_grid_visible[0]->mbmi.delta[lvl_idx] = 0;
+        cm->mi_grid_visible[0]->mbmi.sign[lvl_idx] = 0;
+      }
+#else
       aom_write_literal(w, cm->mi[0].mbmi.filt_lvl, 6);
       cm->mi_grid_visible[0]->mbmi.reuse_sb_lvl = 0;
       cm->mi_grid_visible[0]->mbmi.delta = 0;
       cm->mi_grid_visible[0]->mbmi.sign = 0;
+#endif
     } else {
       int prev_mi_row, prev_mi_col;
       if (mi_col - MAX_MIB_SIZE < 0) {
@@ -2317,6 +2332,44 @@
       MB_MODE_INFO *prev_mbmi =
           &cm->mi_grid_visible[prev_mi_row * cm->mi_stride + prev_mi_col]->mbmi;
 
+#if CONFIG_LOOPFILTER_LEVEL
+      for (lvl_idx = 0; lvl_idx < 4; ++lvl_idx) {
+        const uint8_t curr_lvl = curr_mbmi->filt_lvl[lvl_idx];
+        const uint8_t prev_lvl = prev_mbmi->filt_lvl[lvl_idx];
+
+        const int reuse_prev_lvl = curr_lvl == prev_lvl;
+        const int reuse_ctx = prev_mbmi->reuse_sb_lvl[lvl_idx];
+        curr_mbmi->reuse_sb_lvl[lvl_idx] = reuse_prev_lvl;
+        aom_write_symbol(w, reuse_prev_lvl,
+                         xd->tile_ctx->lpf_reuse_cdf[lvl_idx][reuse_ctx], 2);
+        cpi->td.counts->lpf_reuse[lvl_idx][reuse_ctx][reuse_prev_lvl]++;
+
+        if (reuse_prev_lvl) {
+          curr_mbmi->delta[lvl_idx] = 0;
+          curr_mbmi->sign[lvl_idx] = 0;
+        } else {
+          const unsigned int delta = abs(curr_lvl - prev_lvl) / LPF_STEP;
+          const int delta_ctx = prev_mbmi->delta[lvl_idx];
+          curr_mbmi->delta[lvl_idx] = delta;
+          aom_write_symbol(w, delta,
+                           xd->tile_ctx->lpf_delta_cdf[lvl_idx][delta_ctx],
+                           DELTA_RANGE);
+          cpi->td.counts->lpf_delta[lvl_idx][delta_ctx][delta]++;
+
+          if (delta) {
+            const int sign = curr_lvl > prev_lvl;
+            const int sign_ctx = prev_mbmi->sign[lvl_idx];
+            curr_mbmi->sign[lvl_idx] = sign;
+            aom_write_symbol(
+                w, sign,
+                xd->tile_ctx->lpf_sign_cdf[lvl_idx][reuse_ctx][sign_ctx], 2);
+            cpi->td.counts->lpf_sign[lvl_idx][reuse_ctx][sign_ctx][sign]++;
+          } else {
+            curr_mbmi->sign[lvl_idx] = 0;
+          }
+        }
+      }
+#else
       const uint8_t curr_lvl = curr_mbmi->filt_lvl;
       const uint8_t prev_lvl = prev_mbmi->filt_lvl;
 
@@ -2349,6 +2402,7 @@
           curr_mbmi->sign = 0;
         }
       }
+#endif  // CONFIG_LOOPFILTER_LEVEL
     }
   }
 #endif
@@ -2586,12 +2640,18 @@
 
 // Encode the loop filter level and type
 #if CONFIG_LOOPFILTER_LEVEL
-  aom_wb_write_literal(wb, lf->filter_level[0], 6);
-  aom_wb_write_literal(wb, lf->filter_level[1], 6);
-  if (lf->filter_level[0] || lf->filter_level[1]) {
-    aom_wb_write_literal(wb, lf->filter_level_u, 6);
-    aom_wb_write_literal(wb, lf->filter_level_v, 6);
+#if CONFIG_LPF_SB
+  if (USE_GUESS_LEVEL) {
+#endif  // CONFIG_LPF_SB
+    aom_wb_write_literal(wb, lf->filter_level[0], 6);
+    aom_wb_write_literal(wb, lf->filter_level[1], 6);
+    if (lf->filter_level[0] || lf->filter_level[1]) {
+      aom_wb_write_literal(wb, lf->filter_level_u, 6);
+      aom_wb_write_literal(wb, lf->filter_level_v, 6);
+    }
+#if CONFIG_LPF_SB
   }
+#endif  // CONFIG_LPF_SB
 #else
 #if CONFIG_LPF_SB
   if (USE_GUESS_LEVEL) aom_wb_write_literal(wb, lf->filter_level, 6);
diff --git a/av1/encoder/encodeframe.c b/av1/encoder/encodeframe.c
index f984652..f1d516e 100644
--- a/av1/encoder/encodeframe.c
+++ b/av1/encoder/encodeframe.c
@@ -3256,6 +3256,23 @@
     }
 #if CONFIG_LPF_SB
     if (USE_LOOP_FILTER_SUPERBLOCK) {
+#if CONFIG_LOOPFILTER_LEVEL
+      int filter_lvl[4];
+
+      if (USE_GUESS_LEVEL) {
+        struct loopfilter *lf = &cpi->common.lf;
+        filter_lvl[0] = lf->filter_level[0];
+        filter_lvl[1] = lf->filter_level[1];
+        filter_lvl[2] = lf->filter_level_u;
+        filter_lvl[3] = lf->filter_level_v;
+        av1_loop_filter_frame(cm->frame_to_show, cm, xd, filter_lvl[0],
+                              filter_lvl[1], 0, 1, mi_row, mi_col);
+        av1_loop_filter_frame(cm->frame_to_show, cm, xd, filter_lvl[2],
+                              filter_lvl[2], 1, 1, mi_row, mi_col);
+        av1_loop_filter_frame(cm->frame_to_show, cm, xd, filter_lvl[3],
+                              filter_lvl[3], 2, 1, mi_row, mi_col);
+      }
+#else
       int filter_lvl;
 
       if (USE_GUESS_LEVEL) {
@@ -3283,6 +3300,7 @@
       // if filter_lvl is 0, we still need to set mi info
       if (filter_lvl == 0)
         av1_loop_filter_sb_level_init(cm, mi_row, mi_col, filter_lvl);
+#endif  // CONFIG_LOOPFILTER_LEVEL
     }
 #endif  // CONFIG_LPF_SB
   }
diff --git a/av1/encoder/encoder.c b/av1/encoder/encoder.c
index 0aac2cb..f6228bd 100644
--- a/av1/encoder/encoder.c
+++ b/av1/encoder/encoder.c
@@ -4709,10 +4709,10 @@
 #endif
 #endif  // CONFIG_LPF_SB
   {
-#if CONFIG_LPF_SB
+#if CONFIG_LPF_SB && !CONFIG_LOOPFILTER_LEVEL
     av1_loop_filter_frame(cm->frame_to_show, cm, xd, lf->filter_level, 0, 0, 0,
                           0);
-#else
+#elif !CONFIG_LPF_SB
 #if CONFIG_LOOPFILTER_LEVEL
     av1_loop_filter_frame(cm->frame_to_show, cm, xd, lf->filter_level[0],
                           lf->filter_level[1], 0, 0);
diff --git a/av1/encoder/picklpf.c b/av1/encoder/picklpf.c
index b064ddb..12f07de 100644
--- a/av1/encoder/picklpf.c
+++ b/av1/encoder/picklpf.c
@@ -27,7 +27,7 @@
 #include "av1/encoder/encoder.h"
 #include "av1/encoder/picklpf.h"
 
-#if CONFIG_LPF_SB
+#if CONFIG_LPF_SB && !CONFIG_LOOPFILTER_LEVEL
 #if CONFIG_HIGHBITDEPTH
 static int compute_sb_y_sse_highbd(const YV12_BUFFER_CONFIG *src,
                                    const YV12_BUFFER_CONFIG *frame,
@@ -115,7 +115,7 @@
   }
 }
 
-#if CONFIG_LPF_SB
+#if CONFIG_LPF_SB && !CONFIG_LOOPFILTER_LEVEL
 // TODO(chengchen): reduce memory usage by copy superblock instead of frame
 static int try_filter_superblock(const YV12_BUFFER_CONFIG *sd,
                                  AV1_COMP *const cpi, int filt_level,
@@ -230,7 +230,7 @@
   return filt_best;
 }
 
-#else  // CONFIG_LPF_SB
+#elif !CONFIG_LPF_SB  // CONFIG_LPF_SB
 static int64_t try_filter_frame(const YV12_BUFFER_CONFIG *sd,
                                 AV1_COMP *const cpi, int filt_level,
                                 int partial_frame
@@ -401,6 +401,7 @@
                            LPF_PICK_METHOD method) {
   AV1_COMMON *const cm = &cpi->common;
   struct loopfilter *const lf = &cm->lf;
+  (void)sd;
 
   lf->sharpness_level = cm->frame_type == KEY_FRAME ? 0 : cpi->oxcf.sharpness;
 
@@ -450,13 +451,16 @@
     if (cm->bit_depth != AOM_BITS_8 && cm->frame_type == KEY_FRAME)
       filt_guess -= 4;
 #if CONFIG_LOOPFILTER_LEVEL
+    // TODO(chengchen): retrain the model for Y, U, V filter levels
     lf->filter_level[0] = clamp(filt_guess, min_filter_level, max_filter_level);
     lf->filter_level[1] = clamp(filt_guess, min_filter_level, max_filter_level);
+    lf->filter_level_u = clamp(filt_guess, min_filter_level, max_filter_level);
+    lf->filter_level_v = clamp(filt_guess, min_filter_level, max_filter_level);
 #else
     lf->filter_level = clamp(filt_guess, min_filter_level, max_filter_level);
 #endif
   } else {
-#if CONFIG_LPF_SB
+#if CONFIG_LPF_SB && !CONFIG_LOOPFILTER_LEVEL
     int mi_row, mi_col;
     // TODO(chengchen): init last_lvl using previous frame's info?
     int last_lvl = 0;
@@ -485,7 +489,7 @@
         }
       }
     }
-#else  // CONFIG_LPF_SB
+#elif !CONFIG_LPF_SB  // !CONFIG_LPF_SB
 #if CONFIG_LOOPFILTER_LEVEL
     lf->filter_level[0] = lf->filter_level[1] = search_filter_level(
         sd, cpi, method == LPF_PICK_FROM_SUBIMAGE, NULL, 0, 2);