Add refinement search for sgrproj filter

Change-Id: I8deea1d2a76ce8c36085d19362ac2508fd3c438f

av1/encoder/pickrst.c

@@ -217,6 +217,62 @@ static int64_t get_pixel_proj_error(uint8_t *src8, int width, int height,
   return err;
 }
 
+#define USE_SGRPROJ_REFINEMENT_SEARCH 1
+static int64_t finer_search_pixel_proj_error(
+    uint8_t *src8, int width, int height, int src_stride, uint8_t *dat8,
+    int dat_stride, int bit_depth, int32_t *flt1, int flt1_stride,
+    int32_t *flt2, int flt2_stride, int start_step, int *xqd) {
+  int64_t err = get_pixel_proj_error(src8, width, height, src_stride, dat8,
+                                     dat_stride, bit_depth, flt1, flt1_stride,
+                                     flt2, flt2_stride, xqd);
+  (void)start_step;
+#if USE_SGRPROJ_REFINEMENT_SEARCH
+  int64_t err2;
+  int tap_min[] = { SGRPROJ_PRJ_MIN0, SGRPROJ_PRJ_MIN1 };
+  int tap_max[] = { SGRPROJ_PRJ_MAX0, SGRPROJ_PRJ_MAX1 };
+  for (int s = start_step; s >= 1; s >>= 1) {
+    for (int p = 0; p < 2; ++p) {
+      int skip = 0;
+      do {
+        if (xqd[p] - s >= tap_min[p]) {
+          xqd[p] -= s;
+          err2 = get_pixel_proj_error(src8, width, height, src_stride, dat8,
+                                      dat_stride, bit_depth, flt1, flt1_stride,
+                                      flt2, flt2_stride, xqd);
+          if (err2 > err) {
+            xqd[p] += s;
+          } else {
+            err = err2;
+            skip = 1;
+            // At the highest step size continue moving in the same direction
+            if (s == start_step) continue;
+          }
+        }
+        break;
+      } while (1);
+      if (skip) break;
+      do {
+        if (xqd[p] + s <= tap_max[p]) {
+          xqd[p] += s;
+          err2 = get_pixel_proj_error(src8, width, height, src_stride, dat8,
+                                      dat_stride, bit_depth, flt1, flt1_stride,
+                                      flt2, flt2_stride, xqd);
+          if (err2 > err) {
+            xqd[p] -= s;
+          } else {
+            err = err2;
+            // At the highest step size continue moving in the same direction
+            if (s == start_step) continue;
+          }
+        }
+        break;
+      } while (1);
+    }
+  }
+#endif  // USE_SGRPROJ_REFINEMENT_SEARCH
+  return err;
+}
+
 static void get_proj_subspace(uint8_t *src8, int width, int height,
                               int src_stride, uint8_t *dat8, int dat_stride,
                               int bit_depth, int32_t *flt1, int flt1_stride,
@@ -335,9 +391,9 @@ static void search_selfguided_restoration(uint8_t *dat8, int width, int height,
                       bit_depth, flt1, width, flt2, width, exq);
     aom_clear_system_state();
     encode_xq(exq, exqd);
-    err =
-        get_pixel_proj_error(src8, width, height, src_stride, dat8, dat_stride,
-                             bit_depth, flt1, width, flt2, width, exqd);
+    err = finer_search_pixel_proj_error(src8, width, height, src_stride, dat8,
+                                        dat_stride, bit_depth, flt1, width,
+                                        flt2, width, 2, exqd);
     if (besterr == -1 || err < besterr) {
       bestep = ep;
       besterr = err;
@@ -418,7 +474,8 @@ static double search_sgrproj(const YV12_BUFFER_CONFIG *src, AV1_COMP *cpi,
                              tile_height, width, height, 0, 0, &h_start, &h_end,
                              &v_start, &v_end);
     err = sse_restoration_tile(src, cm->frame_to_show, cm, h_start,
-                               h_end - h_start, v_start, v_end - v_start, 1);
+                               h_end - h_start, v_start, v_end - v_start,
+                               (1 << plane));
     // #bits when a tile is not restored
     bits = av1_cost_bit(RESTORE_NONE_SGRPROJ_PROB, 0);
     cost_norestore = RDCOST_DBL(x->rdmult, x->rddiv, (bits >> 4), err);
@@ -809,6 +866,7 @@ static int64_t finer_tile_search_wiener(const YV12_BUFFER_CONFIG *src,
                                         YV12_BUFFER_CONFIG *dst_frame) {
   int64_t err = try_restoration_tile(src, cpi, rsi, 1 << plane, partial_frame,
                                      tile_idx, 0, 0, dst_frame);
+  (void)start_step;
 #if USE_WIENER_REFINEMENT_SEARCH
   int64_t err2;
   int tap_min[] = { WIENER_FILT_TAP0_MINV, WIENER_FILT_TAP1_MINV,