Add UV wiener loop restoration

Enables Wiener based loop restoration only for the UV
frames. The selfguided and domaintranform filters do not
work very well for UV components, hence they are disabled.
For each UV frame a single set of wiener parameters are
sent. They are applied tile-wise, but all tiles use the
same parameters.

BDRATE (Global PSNR) results:
-----------------------------
lowres: -1.266% (up from -0.666%, good improvement)
midres: -1.815% (up from -1.792%, tiny improvement)

Tiling on UV components will be explored subsequently.

Change-Id: Ib5be93121c4e88e05edf3c36c46488df3cfcd1e2
diff --git a/av1/common/alloccommon.c b/av1/common/alloccommon.c
index 471ae6c..50446b2 100644
--- a/av1/common/alloccommon.c
+++ b/av1/common/alloccommon.c
@@ -89,7 +89,9 @@
 
 #if CONFIG_LOOP_RESTORATION
 void av1_alloc_restoration_buffers(AV1_COMMON *cm) {
-  av1_alloc_restoration_struct(&cm->rst_info, cm->width, cm->height);
+  int p;
+  for (p = 0; p < MAX_MB_PLANE; ++p)
+    av1_alloc_restoration_struct(&cm->rst_info[p], cm->width, cm->height);
   cm->rst_internal.tmpbuf =
       (int32_t *)aom_realloc(cm->rst_internal.tmpbuf, RESTORATION_TMPBUF_SIZE);
   if (cm->rst_internal.tmpbuf == NULL)
@@ -98,7 +100,9 @@
 }
 
 void av1_free_restoration_buffers(AV1_COMMON *cm) {
-  av1_free_restoration_struct(&cm->rst_info);
+  int p;
+  for (p = 0; p < MAX_MB_PLANE; ++p)
+    av1_free_restoration_struct(&cm->rst_info[p]);
   aom_free(cm->rst_internal.tmpbuf);
   cm->rst_internal.tmpbuf = NULL;
 }
diff --git a/av1/common/onyxc_int.h b/av1/common/onyxc_int.h
index 4d88c2c..7c9416d 100644
--- a/av1/common/onyxc_int.h
+++ b/av1/common/onyxc_int.h
@@ -306,7 +306,7 @@
 
   loop_filter_info_n lf_info;
 #if CONFIG_LOOP_RESTORATION
-  RestorationInfo rst_info;
+  RestorationInfo rst_info[MAX_MB_PLANE];
   RestorationInternal rst_internal;
 #endif  // CONFIG_LOOP_RESTORATION
 
diff --git a/av1/common/restoration.c b/av1/common/restoration.c
index 1c00dd2..d725302 100644
--- a/av1/common/restoration.c
+++ b/av1/common/restoration.c
@@ -23,9 +23,6 @@
 
 static int domaintxfmrf_vtable[DOMAINTXFMRF_ITERS][DOMAINTXFMRF_PARAMS][256];
 
-// Whether to filter only y or not
-static const int override_y_only[RESTORE_TYPES] = { 1, 1, 1, 1, 1 };
-
 static const int domaintxfmrf_params[DOMAINTXFMRF_PARAMS] = {
   32,  40,  48,  56,  64,  68,  72,  76,  80,  82,  84,  86,  88,
   90,  92,  94,  96,  97,  98,  99,  100, 101, 102, 103, 104, 105,
@@ -100,53 +97,12 @@
 
 void av1_loop_restoration_precal() { GenDomainTxfmRFVtable(); }
 
-void av1_loop_restoration_init(RestorationInternal *rst, RestorationInfo *rsi,
-                               int kf, int width, int height) {
-  int i, tile_idx;
-  rst->rsi = rsi;
+static void loop_restoration_init(RestorationInternal *rst, int kf, int width,
+                                  int height) {
   rst->keyframe = kf;
-  rst->subsampling_x = 0;
-  rst->subsampling_y = 0;
   rst->ntiles =
       av1_get_rest_ntiles(width, height, &rst->tile_width, &rst->tile_height,
                           &rst->nhtiles, &rst->nvtiles);
-  if (rsi->frame_restoration_type == RESTORE_WIENER) {
-    for (tile_idx = 0; tile_idx < rst->ntiles; ++tile_idx) {
-      if (rsi->wiener_info[tile_idx].level) {
-        rsi->wiener_info[tile_idx].vfilter[WIENER_HALFWIN] =
-            rsi->wiener_info[tile_idx].hfilter[WIENER_HALFWIN] =
-                WIENER_FILT_STEP;
-        for (i = 0; i < WIENER_HALFWIN; ++i) {
-          rsi->wiener_info[tile_idx].vfilter[WIENER_WIN - 1 - i] =
-              rsi->wiener_info[tile_idx].vfilter[i];
-          rsi->wiener_info[tile_idx].hfilter[WIENER_WIN - 1 - i] =
-              rsi->wiener_info[tile_idx].hfilter[i];
-          rsi->wiener_info[tile_idx].vfilter[WIENER_HALFWIN] -=
-              2 * rsi->wiener_info[tile_idx].vfilter[i];
-          rsi->wiener_info[tile_idx].hfilter[WIENER_HALFWIN] -=
-              2 * rsi->wiener_info[tile_idx].hfilter[i];
-        }
-      }
-    }
-  } else if (rsi->frame_restoration_type == RESTORE_SWITCHABLE) {
-    for (tile_idx = 0; tile_idx < rst->ntiles; ++tile_idx) {
-      if (rsi->restoration_type[tile_idx] == RESTORE_WIENER) {
-        rsi->wiener_info[tile_idx].vfilter[WIENER_HALFWIN] =
-            rsi->wiener_info[tile_idx].hfilter[WIENER_HALFWIN] =
-                WIENER_FILT_STEP;
-        for (i = 0; i < WIENER_HALFWIN; ++i) {
-          rsi->wiener_info[tile_idx].vfilter[WIENER_WIN - 1 - i] =
-              rsi->wiener_info[tile_idx].vfilter[i];
-          rsi->wiener_info[tile_idx].hfilter[WIENER_WIN - 1 - i] =
-              rsi->wiener_info[tile_idx].hfilter[i];
-          rsi->wiener_info[tile_idx].vfilter[WIENER_HALFWIN] -=
-              2 * rsi->wiener_info[tile_idx].vfilter[i];
-          rsi->wiener_info[tile_idx].hfilter[WIENER_HALFWIN] -=
-              2 * rsi->wiener_info[tile_idx].hfilter[i];
-        }
-      }
-    }
-  }
 }
 
 static void extend_frame(uint8_t *data, int width, int height, int stride) {
@@ -1043,9 +999,10 @@
 }
 #endif  // CONFIG_AOM_HIGHBITDEPTH
 
-void av1_loop_restoration_rows(YV12_BUFFER_CONFIG *frame, AV1_COMMON *cm,
-                               int start_mi_row, int end_mi_row, int y_only,
-                               YV12_BUFFER_CONFIG *dst) {
+static void loop_restoration_rows(YV12_BUFFER_CONFIG *frame, AV1_COMMON *cm,
+                                  int start_mi_row, int end_mi_row,
+                                  int components_pattern, RestorationInfo *rsi,
+                                  YV12_BUFFER_CONFIG *dst) {
   const int ywidth = frame->y_crop_width;
   const int ystride = frame->y_stride;
   const int uvwidth = frame->uv_crop_width;
@@ -1064,29 +1021,44 @@
     loop_domaintxfmrf_filter_highbd, loop_switchable_filter_highbd
   };
 #endif  // CONFIG_AOM_HIGHBITDEPTH
-  restore_func_type restore_func =
-      restore_funcs[cm->rst_internal.rsi->frame_restoration_type];
+  restore_func_type restore_func;
 #if CONFIG_AOM_HIGHBITDEPTH
-  restore_func_highbd_type restore_func_highbd =
-      restore_funcs_highbd[cm->rst_internal.rsi->frame_restoration_type];
+  restore_func_highbd_type restore_func_highbd;
 #endif  // CONFIG_AOM_HIGHBITDEPTH
   YV12_BUFFER_CONFIG dst_;
 
   yend = AOMMIN(yend, cm->height);
   uvend = AOMMIN(uvend, cm->subsampling_y ? (cm->height + 1) >> 1 : cm->height);
 
-  if (cm->rst_internal.rsi->frame_restoration_type == RESTORE_NONE) {
-    if (dst) {
-      if (y_only)
-        aom_yv12_copy_y(frame, dst);
-      else
-        aom_yv12_copy_frame(frame, dst);
+  if (components_pattern == (1 << AOM_PLANE_Y)) {
+    // Only y
+    if (rsi[0].frame_restoration_type == RESTORE_NONE) {
+      if (dst) aom_yv12_copy_y(frame, dst);
+      return;
     }
-    return;
+  } else if (components_pattern == (1 << AOM_PLANE_U)) {
+    // Only U
+    if (rsi[1].frame_restoration_type == RESTORE_NONE) {
+      if (dst) aom_yv12_copy_u(frame, dst);
+      return;
+    }
+  } else if (components_pattern == (1 << AOM_PLANE_V)) {
+    // Only V
+    if (rsi[2].frame_restoration_type == RESTORE_NONE) {
+      if (dst) aom_yv12_copy_v(frame, dst);
+      return;
+    }
+  } else if (components_pattern ==
+             ((1 << AOM_PLANE_Y) | (1 << AOM_PLANE_U) | (1 << AOM_PLANE_V))) {
+    // All components
+    if (rsi[0].frame_restoration_type == RESTORE_NONE &&
+        rsi[1].frame_restoration_type == RESTORE_NONE &&
+        rsi[2].frame_restoration_type == RESTORE_NONE) {
+      if (dst) aom_yv12_copy_frame(frame, dst);
+      return;
+    }
   }
 
-  if (y_only == 0)
-    y_only = override_y_only[cm->rst_internal.rsi->frame_restoration_type];
   if (!dst) {
     dst = &dst_;
     memset(dst, 0, sizeof(YV12_BUFFER_CONFIG));
@@ -1100,68 +1072,103 @@
                          "Failed to allocate restoration dst buffer");
   }
 
+  if ((components_pattern >> AOM_PLANE_Y) & 1) {
+    if (rsi[0].frame_restoration_type != RESTORE_NONE) {
+      cm->rst_internal.subsampling_x = 0;
+      cm->rst_internal.subsampling_y = 0;
+      cm->rst_internal.rsi = &rsi[0];
+      restore_func =
+          restore_funcs[cm->rst_internal.rsi->frame_restoration_type];
 #if CONFIG_AOM_HIGHBITDEPTH
-  if (cm->use_highbitdepth)
-    restore_func_highbd(frame->y_buffer + ystart * ystride, ywidth,
-                        yend - ystart, ystride, &cm->rst_internal,
-                        cm->bit_depth, dst->y_buffer + ystart * dst->y_stride,
-                        dst->y_stride);
-  else
+      restore_func_highbd =
+          restore_funcs_highbd[cm->rst_internal.rsi->frame_restoration_type];
+      if (cm->use_highbitdepth)
+        restore_func_highbd(
+            frame->y_buffer + ystart * ystride, ywidth, yend - ystart, ystride,
+            &cm->rst_internal, cm->bit_depth,
+            dst->y_buffer + ystart * dst->y_stride, dst->y_stride);
+      else
 #endif  // CONFIG_AOM_HIGHBITDEPTH
-    restore_func(frame->y_buffer + ystart * ystride, ywidth, yend - ystart,
-                 ystride, &cm->rst_internal,
-                 dst->y_buffer + ystart * dst->y_stride, dst->y_stride);
-  if (!y_only) {
+        restore_func(frame->y_buffer + ystart * ystride, ywidth, yend - ystart,
+                     ystride, &cm->rst_internal,
+                     dst->y_buffer + ystart * dst->y_stride, dst->y_stride);
+    } else {
+      aom_yv12_copy_y(frame, dst);
+    }
+  }
+
+  if ((components_pattern >> AOM_PLANE_U) & 1) {
     cm->rst_internal.subsampling_x = cm->subsampling_x;
     cm->rst_internal.subsampling_y = cm->subsampling_y;
+    cm->rst_internal.rsi = &rsi[1];
+    if (rsi[1].frame_restoration_type != RESTORE_NONE) {
+      restore_func =
+          restore_funcs[cm->rst_internal.rsi->frame_restoration_type];
 #if CONFIG_AOM_HIGHBITDEPTH
-    if (cm->use_highbitdepth) {
-      restore_func_highbd(
-          frame->u_buffer + uvstart * uvstride, uvwidth, uvend - uvstart,
-          uvstride, &cm->rst_internal, cm->bit_depth,
-          dst->u_buffer + uvstart * dst->uv_stride, dst->uv_stride);
-      restore_func_highbd(
-          frame->v_buffer + uvstart * uvstride, uvwidth, uvend - uvstart,
-          uvstride, &cm->rst_internal, cm->bit_depth,
-          dst->v_buffer + uvstart * dst->uv_stride, dst->uv_stride);
+      restore_func_highbd =
+          restore_funcs_highbd[cm->rst_internal.rsi->frame_restoration_type];
+      if (cm->use_highbitdepth)
+        restore_func_highbd(
+            frame->u_buffer + uvstart * uvstride, uvwidth, uvend - uvstart,
+            uvstride, &cm->rst_internal, cm->bit_depth,
+            dst->u_buffer + uvstart * dst->uv_stride, dst->uv_stride);
+      else
+#endif  // CONFIG_AOM_HIGHBITDEPTH
+        restore_func(frame->u_buffer + uvstart * uvstride, uvwidth,
+                     uvend - uvstart, uvstride, &cm->rst_internal,
+                     dst->u_buffer + uvstart * dst->uv_stride, dst->uv_stride);
     } else {
-#endif  // CONFIG_AOM_HIGHBITDEPTH
-      restore_func(frame->u_buffer + uvstart * uvstride, uvwidth,
-                   uvend - uvstart, uvstride, &cm->rst_internal,
-                   dst->u_buffer + uvstart * dst->uv_stride, dst->uv_stride);
-      restore_func(frame->v_buffer + uvstart * uvstride, uvwidth,
-                   uvend - uvstart, uvstride, &cm->rst_internal,
-                   dst->v_buffer + uvstart * dst->uv_stride, dst->uv_stride);
-#if CONFIG_AOM_HIGHBITDEPTH
+      aom_yv12_copy_u(frame, dst);
     }
+  }
+
+  if ((components_pattern >> AOM_PLANE_V) & 1) {
+    cm->rst_internal.subsampling_x = cm->subsampling_x;
+    cm->rst_internal.subsampling_y = cm->subsampling_y;
+    cm->rst_internal.rsi = &rsi[2];
+    if (rsi[2].frame_restoration_type != RESTORE_NONE) {
+      restore_func =
+          restore_funcs[cm->rst_internal.rsi->frame_restoration_type];
+#if CONFIG_AOM_HIGHBITDEPTH
+      restore_func_highbd =
+          restore_funcs_highbd[cm->rst_internal.rsi->frame_restoration_type];
+      if (cm->use_highbitdepth)
+        restore_func_highbd(
+            frame->v_buffer + uvstart * uvstride, uvwidth, uvend - uvstart,
+            uvstride, &cm->rst_internal, cm->bit_depth,
+            dst->v_buffer + uvstart * dst->uv_stride, dst->uv_stride);
+      else
 #endif  // CONFIG_AOM_HIGHBITDEPTH
+        restore_func(frame->v_buffer + uvstart * uvstride, uvwidth,
+                     uvend - uvstart, uvstride, &cm->rst_internal,
+                     dst->v_buffer + uvstart * dst->uv_stride, dst->uv_stride);
+    } else {
+      aom_yv12_copy_v(frame, dst);
+    }
   }
 
   if (dst == &dst_) {
-    if (y_only)
-      aom_yv12_copy_y(dst, frame);
-    else
-      aom_yv12_copy_frame(dst, frame);
+    if ((components_pattern >> AOM_PLANE_Y) & 1) aom_yv12_copy_y(dst, frame);
+    if ((components_pattern >> AOM_PLANE_U) & 1) aom_yv12_copy_u(dst, frame);
+    if ((components_pattern >> AOM_PLANE_V) & 1) aom_yv12_copy_v(dst, frame);
     aom_free_frame_buffer(dst);
   }
 }
 
 void av1_loop_restoration_frame(YV12_BUFFER_CONFIG *frame, AV1_COMMON *cm,
-                                RestorationInfo *rsi, int y_only,
+                                RestorationInfo *rsi, int components_pattern,
                                 int partial_frame, YV12_BUFFER_CONFIG *dst) {
   int start_mi_row, end_mi_row, mi_rows_to_filter;
-  if (rsi->frame_restoration_type != RESTORE_NONE) {
-    start_mi_row = 0;
-    mi_rows_to_filter = cm->mi_rows;
-    if (partial_frame && cm->mi_rows > 8) {
-      start_mi_row = cm->mi_rows >> 1;
-      start_mi_row &= 0xfffffff8;
-      mi_rows_to_filter = AOMMAX(cm->mi_rows / 8, 8);
-    }
-    end_mi_row = start_mi_row + mi_rows_to_filter;
-    av1_loop_restoration_init(&cm->rst_internal, rsi,
-                              cm->frame_type == KEY_FRAME, cm->width,
-                              cm->height);
-    av1_loop_restoration_rows(frame, cm, start_mi_row, end_mi_row, y_only, dst);
+  start_mi_row = 0;
+  mi_rows_to_filter = cm->mi_rows;
+  if (partial_frame && cm->mi_rows > 8) {
+    start_mi_row = cm->mi_rows >> 1;
+    start_mi_row &= 0xfffffff8;
+    mi_rows_to_filter = AOMMAX(cm->mi_rows / 8, 8);
   }
+  end_mi_row = start_mi_row + mi_rows_to_filter;
+  loop_restoration_init(&cm->rst_internal, cm->frame_type == KEY_FRAME,
+                        cm->width, cm->height);
+  loop_restoration_rows(frame, cm, start_mi_row, end_mi_row, components_pattern,
+                        rsi, dst);
 }
diff --git a/av1/common/restoration.h b/av1/common/restoration.h
index 8bcbb91..cf44962 100644
--- a/av1/common/restoration.h
+++ b/av1/common/restoration.h
@@ -233,14 +233,9 @@
                                          int32_t *tmpbuf);
 #endif  // CONFIG_AOM_HIGHBITDEPTH
 void decode_xq(int *xqd, int *xq);
-void av1_loop_restoration_init(RestorationInternal *rst, RestorationInfo *rsi,
-                               int kf, int width, int height);
 void av1_loop_restoration_frame(YV12_BUFFER_CONFIG *frame, struct AV1Common *cm,
-                                RestorationInfo *rsi, int y_only,
+                                RestorationInfo *rsi, int components_pattern,
                                 int partial_frame, YV12_BUFFER_CONFIG *dst);
-void av1_loop_restoration_rows(YV12_BUFFER_CONFIG *frame, struct AV1Common *cm,
-                               int start_mi_row, int end_mi_row, int y_only,
-                               YV12_BUFFER_CONFIG *dst);
 void av1_loop_restoration_precal();
 #ifdef __cplusplus
 }  // extern "C"