Refactor UV restoration to use same tilesize as Y

Change-Id: I56e741551f74624a84250d7565520db9c5127d1b

av1/common/alloccommon.c

@@ -90,8 +90,11 @@ void av1_free_ref_frame_buffers(BufferPool *pool) {
 #if CONFIG_LOOP_RESTORATION
 void av1_alloc_restoration_buffers(AV1_COMMON *cm) {
   int p;
-  for (p = 0; p < MAX_MB_PLANE; ++p)
-    av1_alloc_restoration_struct(&cm->rst_info[p], cm->width, cm->height);
+  av1_alloc_restoration_struct(&cm->rst_info[0], cm->width, cm->height);
+  for (p = 1; p < MAX_MB_PLANE; ++p)
+    av1_alloc_restoration_struct(&cm->rst_info[p],
+                                 cm->width >> cm->subsampling_x,
+                                 cm->height >> cm->subsampling_y);
   cm->rst_internal.tmpbuf =
       (int32_t *)aom_realloc(cm->rst_internal.tmpbuf, RESTORATION_TMPBUF_SIZE);
   if (cm->rst_internal.tmpbuf == NULL)

av1/common/restoration.c

@@ -97,12 +97,8 @@ static void GenDomainTxfmRFVtable() {
 
 void av1_loop_restoration_precal() { GenDomainTxfmRFVtable(); }
 
-static void loop_restoration_init(RestorationInternal *rst, int kf, int width,
-                                  int height) {
+static void loop_restoration_init(RestorationInternal *rst, int kf) {
   rst->keyframe = kf;
-  rst->ntiles =
-      av1_get_rest_ntiles(width, height, &rst->tile_width, &rst->tile_height,
-                          &rst->nhtiles, &rst->nvtiles);
 }
 
 void extend_frame(uint8_t *data, int width, int height, int stride) {
@@ -127,8 +123,8 @@ static void loop_copy_tile(uint8_t *data, int tile_idx, int subtile_idx,
                            int subtile_bits, int width, int height, int stride,
                            RestorationInternal *rst, uint8_t *dst,
                            int dst_stride) {
-  const int tile_width = rst->tile_width >> rst->subsampling_x;
-  const int tile_height = rst->tile_height >> rst->subsampling_y;
+  const int tile_width = rst->tile_width;
+  const int tile_height = rst->tile_height;
   int i;
   int h_start, h_end, v_start, v_end;
   av1_get_rest_tile_limits(tile_idx, subtile_idx, subtile_bits, rst->nhtiles,
@@ -143,8 +139,8 @@ static void loop_wiener_filter_tile(uint8_t *data, int tile_idx, int width,
                                     int height, int stride,
                                     RestorationInternal *rst, uint8_t *dst,
                                     int dst_stride) {
-  const int tile_width = rst->tile_width >> rst->subsampling_x;
-  const int tile_height = rst->tile_height >> rst->subsampling_y;
+  const int tile_width = rst->tile_width;
+  const int tile_height = rst->tile_height;
   int i, j;
   int h_start, h_end, v_start, v_end;
   DECLARE_ALIGNED(16, InterpKernel, hkernel);
@@ -631,8 +627,8 @@ static void loop_sgrproj_filter_tile(uint8_t *data, int tile_idx, int width,
                                      int height, int stride,
                                      RestorationInternal *rst, uint8_t *dst,
                                      int dst_stride) {
-  const int tile_width = rst->tile_width >> rst->subsampling_x;
-  const int tile_height = rst->tile_height >> rst->subsampling_y;
+  const int tile_width = rst->tile_width;
+  const int tile_height = rst->tile_height;
   int h_start, h_end, v_start, v_end;
   uint8_t *data_p, *dst_p;
 
@@ -790,8 +786,8 @@ static void loop_domaintxfmrf_filter_tile(uint8_t *data, int tile_idx,
                                           int width, int height, int stride,
                                           RestorationInternal *rst,
                                           uint8_t *dst, int dst_stride) {
-  const int tile_width = rst->tile_width >> rst->subsampling_x;
-  const int tile_height = rst->tile_height >> rst->subsampling_y;
+  const int tile_width = rst->tile_width;
+  const int tile_height = rst->tile_height;
   int h_start, h_end, v_start, v_end;
   int32_t *tmpbuf = (int32_t *)rst->tmpbuf;
 
@@ -866,8 +862,8 @@ static void loop_copy_tile_highbd(uint16_t *data, int tile_idx, int subtile_idx,
                                   int subtile_bits, int width, int height,
                                   int stride, RestorationInternal *rst,
                                   uint16_t *dst, int dst_stride) {
-  const int tile_width = rst->tile_width >> rst->subsampling_x;
-  const int tile_height = rst->tile_height >> rst->subsampling_y;
+  const int tile_width = rst->tile_width;
+  const int tile_height = rst->tile_height;
   int i;
   int h_start, h_end, v_start, v_end;
   av1_get_rest_tile_limits(tile_idx, subtile_idx, subtile_bits, rst->nhtiles,
@@ -883,8 +879,8 @@ static void loop_wiener_filter_tile_highbd(uint16_t *data, int tile_idx,
                                            RestorationInternal *rst,
                                            int bit_depth, uint16_t *dst,
                                            int dst_stride) {
-  const int tile_width = rst->tile_width >> rst->subsampling_x;
-  const int tile_height = rst->tile_height >> rst->subsampling_y;
+  const int tile_width = rst->tile_width;
+  const int tile_height = rst->tile_height;
   int h_start, h_end, v_start, v_end;
   int i, j;
   DECLARE_ALIGNED(16, InterpKernel, hkernel);
@@ -978,8 +974,8 @@ static void loop_sgrproj_filter_tile_highbd(uint16_t *data, int tile_idx,
                                             RestorationInternal *rst,
                                             int bit_depth, uint16_t *dst,
                                             int dst_stride) {
-  const int tile_width = rst->tile_width >> rst->subsampling_x;
-  const int tile_height = rst->tile_height >> rst->subsampling_y;
+  const int tile_width = rst->tile_width;
+  const int tile_height = rst->tile_height;
   int h_start, h_end, v_start, v_end;
   uint16_t *data_p, *dst_p;
 
@@ -1108,8 +1104,8 @@ void av1_domaintxfmrf_restoration_highbd(uint16_t *dgd, int width, int height,
 static void loop_domaintxfmrf_filter_tile_highbd(
     uint16_t *data, int tile_idx, int width, int height, int stride,
     RestorationInternal *rst, int bit_depth, uint16_t *dst, int dst_stride) {
-  const int tile_width = rst->tile_width >> rst->subsampling_x;
-  const int tile_height = rst->tile_height >> rst->subsampling_y;
+  const int tile_width = rst->tile_width;
+  const int tile_height = rst->tile_height;
   int h_start, h_end, v_start, v_end;
   int32_t *tmpbuf = (int32_t *)rst->tmpbuf;
 
@@ -1243,8 +1239,10 @@ static void loop_restoration_rows(YV12_BUFFER_CONFIG *frame, AV1_COMMON *cm,
 
   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.ntiles = av1_get_rest_ntiles(
+          cm->width, cm->height, &cm->rst_internal.tile_width,
+          &cm->rst_internal.tile_height, &cm->rst_internal.nhtiles,
+          &cm->rst_internal.nvtiles);
       cm->rst_internal.rsi = &rsi[0];
       restore_func =
           restore_funcs[cm->rst_internal.rsi->frame_restoration_type];
@@ -1267,10 +1265,12 @@ static void loop_restoration_rows(YV12_BUFFER_CONFIG *frame, AV1_COMMON *cm,
   }
 
   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) {
+    if (rsi[AOM_PLANE_U].frame_restoration_type != RESTORE_NONE) {
+      cm->rst_internal.ntiles = av1_get_rest_ntiles(
+          cm->width >> cm->subsampling_x, cm->height >> cm->subsampling_y,
+          &cm->rst_internal.tile_width, &cm->rst_internal.tile_height,
+          &cm->rst_internal.nhtiles, &cm->rst_internal.nvtiles);
+      cm->rst_internal.rsi = &rsi[AOM_PLANE_U];
       restore_func =
           restore_funcs[cm->rst_internal.rsi->frame_restoration_type];
 #if CONFIG_AOM_HIGHBITDEPTH
@@ -1292,10 +1292,12 @@ static void loop_restoration_rows(YV12_BUFFER_CONFIG *frame, AV1_COMMON *cm,
   }
 
   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) {
+    if (rsi[AOM_PLANE_V].frame_restoration_type != RESTORE_NONE) {
+      cm->rst_internal.ntiles = av1_get_rest_ntiles(
+          cm->width >> cm->subsampling_x, cm->height >> cm->subsampling_y,
+          &cm->rst_internal.tile_width, &cm->rst_internal.tile_height,
+          &cm->rst_internal.nhtiles, &cm->rst_internal.nvtiles);
+      cm->rst_internal.rsi = &rsi[AOM_PLANE_V];
       restore_func =
           restore_funcs[cm->rst_internal.rsi->frame_restoration_type];
 #if CONFIG_AOM_HIGHBITDEPTH
@@ -1336,8 +1338,7 @@ void av1_loop_restoration_frame(YV12_BUFFER_CONFIG *frame, AV1_COMMON *cm,
     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_init(&cm->rst_internal, cm->frame_type == KEY_FRAME);
   loop_restoration_rows(frame, cm, start_mi_row, end_mi_row, components_pattern,
                         rsi, dst);
 }

av1/common/restoration.h

@@ -151,8 +151,6 @@ typedef struct {
 typedef struct {
   RestorationInfo *rsi;
   int keyframe;
-  int subsampling_x;
-  int subsampling_y;
   int ntiles;
   int tile_width, tile_height;
   int nhtiles, nvtiles;

av1/decoder/decodeframe.c

@@ -2358,6 +2358,9 @@ static void decode_restoration(AV1_COMMON *cm, aom_reader *rb) {
   int i, p;
   const int ntiles =
       av1_get_rest_ntiles(cm->width, cm->height, NULL, NULL, NULL, NULL);
+  const int ntiles_uv = av1_get_rest_ntiles(cm->width >> cm->subsampling_x,
+                                            cm->height >> cm->subsampling_x,
+                                            NULL, NULL, NULL, NULL);
   RestorationInfo *rsi = &cm->rst_info[0];
   if (rsi->frame_restoration_type != RESTORE_NONE) {
     if (rsi->frame_restoration_type == RESTORE_SWITCHABLE) {
@@ -2417,7 +2420,7 @@ static void decode_restoration(AV1_COMMON *cm, aom_reader *rb) {
       rsi->restoration_type[0] = RESTORE_WIENER;
       rsi->wiener_info[0].level = 1;
       read_wiener_filter(&rsi->wiener_info[0], rb);
-      for (i = 1; i < ntiles; ++i) {
+      for (i = 1; i < ntiles_uv; ++i) {
         rsi->restoration_type[i] = RESTORE_WIENER;
         memcpy(&rsi->wiener_info[i], &rsi->wiener_info[0],
                sizeof(rsi->wiener_info[0]));

av1/encoder/bitstream.c

@@ -3086,11 +3086,13 @@ static void write_domaintxfmrf_filter(DomaintxfmrfInfo *domaintxfmrf_info,
 
 static void encode_restoration(AV1_COMMON *cm, aom_writer *wb) {
   int i, p;
+  const int ntiles =
+      av1_get_rest_ntiles(cm->width, cm->height, NULL, NULL, NULL, NULL);
   RestorationInfo *rsi = &cm->rst_info[0];
   if (rsi->frame_restoration_type != RESTORE_NONE) {
     if (rsi->frame_restoration_type == RESTORE_SWITCHABLE) {
       // RESTORE_SWITCHABLE
-      for (i = 0; i < cm->rst_internal.ntiles; ++i) {
+      for (i = 0; i < ntiles; ++i) {
         av1_write_token(
             wb, av1_switchable_restore_tree, cm->fc->switchable_restore_prob,
             &switchable_restore_encodings[rsi->restoration_type[i]]);
@@ -3103,14 +3105,14 @@ static void encode_restoration(AV1_COMMON *cm, aom_writer *wb) {
         }
       }
     } else if (rsi->frame_restoration_type == RESTORE_WIENER) {
-      for (i = 0; i < cm->rst_internal.ntiles; ++i) {
+      for (i = 0; i < ntiles; ++i) {
         aom_write(wb, rsi->wiener_info[i].level != 0, RESTORE_NONE_WIENER_PROB);
         if (rsi->wiener_info[i].level) {
           write_wiener_filter(&rsi->wiener_info[i], wb);
         }
       }
     } else if (rsi->frame_restoration_type == RESTORE_SGRPROJ) {
-      for (i = 0; i < cm->rst_internal.ntiles; ++i) {
+      for (i = 0; i < ntiles; ++i) {
         aom_write(wb, rsi->sgrproj_info[i].level != 0,
                   RESTORE_NONE_SGRPROJ_PROB);
         if (rsi->sgrproj_info[i].level) {
@@ -3118,7 +3120,7 @@ static void encode_restoration(AV1_COMMON *cm, aom_writer *wb) {
         }
       }
     } else if (rsi->frame_restoration_type == RESTORE_DOMAINTXFMRF) {
-      for (i = 0; i < cm->rst_internal.ntiles; ++i) {
+      for (i = 0; i < ntiles; ++i) {
         aom_write(wb, rsi->domaintxfmrf_info[i].level != 0,
                   RESTORE_NONE_DOMAINTXFMRF_PROB);
         if (rsi->domaintxfmrf_info[i].level) {

av1/encoder/pickrst.c

@@ -44,6 +44,10 @@ static int64_t sse_restoration_tile(const YV12_BUFFER_CONFIG *src,
                                     int width, int v_start, int height,
                                     int components_pattern) {
   int64_t filt_err = 0;
+  (void)cm;
+  // Y and UV components cannot be mixed
+  assert(components_pattern == 1 || components_pattern == 2 ||
+         components_pattern == 4 || components_pattern == 6);
 #if CONFIG_AOM_HIGHBITDEPTH
   if (cm->use_highbitdepth) {
     if ((components_pattern >> AOM_PLANE_Y) & 1) {
@@ -51,14 +55,12 @@ static int64_t sse_restoration_tile(const YV12_BUFFER_CONFIG *src,
           aom_highbd_get_y_sse_part(src, dst, h_start, width, v_start, height);
     }
     if ((components_pattern >> AOM_PLANE_U) & 1) {
-      filt_err += aom_highbd_get_u_sse_part(
-          src, dst, h_start >> cm->subsampling_x, width >> cm->subsampling_x,
-          v_start >> cm->subsampling_y, height >> cm->subsampling_y);
+      filt_err +=
+          aom_highbd_get_u_sse_part(src, dst, h_start, width, v_start, height);
     }
     if ((components_pattern >> AOM_PLANE_V) & 1) {
-      filt_err += aom_highbd_get_v_sse_part(
-          src, dst, h_start >> cm->subsampling_x, width >> cm->subsampling_x,
-          v_start >> cm->subsampling_y, height >> cm->subsampling_y);
+      filt_err +=
+          aom_highbd_get_v_sse_part(src, dst, h_start, width, v_start, height);
     }
     return filt_err;
   }
@@ -67,14 +69,10 @@ static int64_t sse_restoration_tile(const YV12_BUFFER_CONFIG *src,
     filt_err += aom_get_y_sse_part(src, dst, h_start, width, v_start, height);
   }
   if ((components_pattern >> AOM_PLANE_U) & 1) {
-    filt_err += aom_get_u_sse_part(
-        src, dst, h_start >> cm->subsampling_x, width >> cm->subsampling_x,
-        v_start >> cm->subsampling_y, height >> cm->subsampling_y);
+    filt_err += aom_get_u_sse_part(src, dst, h_start, width, v_start, height);
   }
   if ((components_pattern >> AOM_PLANE_V) & 1) {
-    filt_err += aom_get_u_sse_part(
-        src, dst, h_start >> cm->subsampling_x, width >> cm->subsampling_x,
-        v_start >> cm->subsampling_y, height >> cm->subsampling_y);
+    filt_err += aom_get_v_sse_part(src, dst, h_start, width, v_start, height);
   }
   return filt_err;
 }
@@ -119,16 +117,28 @@ static int64_t try_restoration_tile(const YV12_BUFFER_CONFIG *src,
   int64_t filt_err;
   int tile_width, tile_height, nhtiles, nvtiles;
   int h_start, h_end, v_start, v_end;
-  const int ntiles = av1_get_rest_ntiles(cm->width, cm->height, &tile_width,
-                                         &tile_height, &nhtiles, &nvtiles);
+  int ntiles, width, height;
+
+  // Y and UV components cannot be mixed
+  assert(components_pattern == 1 || components_pattern == 2 ||
+         components_pattern == 4 || components_pattern == 6);
+
+  if (components_pattern == 1) {  // Y only
+    width = src->y_crop_width;
+    height = src->y_crop_height;
+  } else {  // Color
+    width = src->uv_crop_width;
+    height = src->uv_crop_height;
+  }
+  ntiles = av1_get_rest_ntiles(width, height, &tile_width, &tile_height,
+                               &nhtiles, &nvtiles);
   (void)ntiles;
 
   av1_loop_restoration_frame(cm->frame_to_show, cm, rsi, components_pattern,
                              partial_frame, dst_frame);
   av1_get_rest_tile_limits(tile_idx, subtile_idx, subtile_bits, nhtiles,
-                           nvtiles, tile_width, tile_height, cm->width,
-                           cm->height, 0, 0, &h_start, &h_end, &v_start,
-                           &v_end);
+                           nvtiles, tile_width, tile_height, width, height, 0,
+                           0, &h_start, &h_end, &v_start, &v_end);
   filt_err = sse_restoration_tile(src, dst_frame, cm, h_start, h_end - h_start,
                                   v_start, v_end - v_start, components_pattern);
 
@@ -951,7 +961,7 @@ static double search_wiener_uv(const YV12_BUFFER_CONFIG *src, AV1_COMP *cpi,
   const int dgd_stride = dgd->uv_stride;
   double score;
   int tile_idx, tile_width, tile_height, nhtiles, nvtiles;
-  const int ntiles = av1_get_rest_ntiles(cm->width, cm->height, &tile_width,
+  const int ntiles = av1_get_rest_ntiles(width, height, &tile_width,
                                          &tile_height, &nhtiles, &nvtiles);
 
   assert(width == dgd->uv_crop_width);
@@ -1361,6 +1371,8 @@ void av1_pick_filter_restoration(const YV12_BUFFER_CONFIG *src, AV1_COMP *cpi,
          cm->rst_info[0].frame_restoration_type,
          cm->rst_info[1].frame_restoration_type,
          cm->rst_info[2].frame_restoration_type);
+         */
+  /*
   printf("Frame %d/%d frame_restore_type %d : %f %f %f %f %f\n",
          cm->current_video_frame, cm->show_frame,
          cm->rst_info[0].frame_restoration_type, cost_restore[0],