Add restoration tilesize to frame header

The restoration tilesize can be now chosen as either 256, 128
or 64 at the frame header.

Change-Id: I852fc42afedc053484d657bdca522de73aaacd67

av1/common/alloccommon.c

@@ -88,6 +88,7 @@ void av1_free_ref_frame_buffers(BufferPool *pool) {
 }
 
 #if CONFIG_LOOP_RESTORATION
+// Assumes cm->rst_info[p].restoration_tilesize is already initialized
 void av1_alloc_restoration_buffers(AV1_COMMON *cm) {
   int p;
   av1_alloc_restoration_struct(cm, &cm->rst_info[0], cm->width, cm->height);

av1/common/restoration.c

@@ -58,8 +58,8 @@ typedef void (*restore_func_highbd_type)(uint8_t *data8, int width, int height,
 
 int av1_alloc_restoration_struct(AV1_COMMON *cm, RestorationInfo *rst_info,
                                  int width, int height) {
-  const int ntiles = av1_get_rest_ntiles(width, height, NULL, NULL, NULL, NULL);
-  rst_info->frame_restoration_type = RESTORE_NONE;
+  const int ntiles = av1_get_rest_ntiles(
+      width, height, rst_info->restoration_tilesize, NULL, NULL, NULL, NULL);
   aom_free(rst_info->restoration_type);
   CHECK_MEM_ERROR(cm, rst_info->restoration_type,
                   (RestorationType *)aom_malloc(
@@ -1444,9 +1444,9 @@ 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.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->width, cm->height, cm->rst_info[AOM_PLANE_Y].restoration_tilesize,
+          &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];
@@ -1473,6 +1473,7 @@ static void loop_restoration_rows(YV12_BUFFER_CONFIG *frame, AV1_COMMON *cm,
       cm->rst_internal.ntiles = av1_get_rest_ntiles(
           ROUND_POWER_OF_TWO(cm->width, cm->subsampling_x),
           ROUND_POWER_OF_TWO(cm->height, cm->subsampling_y),
+          cm->rst_info[AOM_PLANE_U].restoration_tilesize,
           &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];
@@ -1501,6 +1502,7 @@ static void loop_restoration_rows(YV12_BUFFER_CONFIG *frame, AV1_COMMON *cm,
       cm->rst_internal.ntiles = av1_get_rest_ntiles(
           ROUND_POWER_OF_TWO(cm->width, cm->subsampling_x),
           ROUND_POWER_OF_TWO(cm->height, cm->subsampling_y),
+          cm->rst_info[AOM_PLANE_V].restoration_tilesize,
           &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];

av1/common/restoration.h

@@ -24,10 +24,9 @@ extern "C" {
 #define CLIP(x, lo, hi) ((x) < (lo) ? (lo) : (x) > (hi) ? (hi) : (x))
 #define RINT(x) ((x) < 0 ? (int)((x)-0.5) : (int)((x) + 0.5))
 
-#define RESTORATION_TILESIZE_SML 128
-#define RESTORATION_TILESIZE_BIG 256
+#define RESTORATION_TILESIZE_MAX 256
 #define RESTORATION_TILEPELS_MAX \
-  (RESTORATION_TILESIZE_BIG * RESTORATION_TILESIZE_BIG * 9 / 4)
+  (RESTORATION_TILESIZE_MAX * RESTORATION_TILESIZE_MAX * 9 / 4)
 
 #if USE_DOMAINTXFMRF
 #define DOMAINTXFMRF_PARAMS_BITS 6
@@ -56,7 +55,7 @@ extern "C" {
 // 2 for the restored versions of the frame and
 // 2 for each restoration operation
 #define SGRPROJ_OUTBUF_SIZE \
-  ((RESTORATION_TILESIZE_BIG * 3 / 2) * (RESTORATION_TILESIZE_BIG * 3 / 2 + 16))
+  ((RESTORATION_TILESIZE_MAX * 3 / 2) * (RESTORATION_TILESIZE_MAX * 3 / 2 + 16))
 #define SGRPROJ_TMPBUF_SIZE                         \
   (RESTORATION_TILEPELS_MAX * 2 * sizeof(int32_t) + \
    SGRPROJ_OUTBUF_SIZE * 2 * sizeof(int32_t))
@@ -160,6 +159,7 @@ typedef struct { int sigma_r; } DomaintxfmrfInfo;
 #endif  // USE_DOMAINTXFMRF
 
 typedef struct {
+  int restoration_tilesize;
   RestorationType frame_restoration_type;
   RestorationType *restoration_type;
   // Wiener filter
@@ -181,19 +181,11 @@ typedef struct {
   int32_t *tmpbuf;
 } RestorationInternal;
 
-static INLINE int get_rest_tilesize(int width, int height) {
-  if (width * height <= 352 * 288)
-    return RESTORATION_TILESIZE_SML;
-  else
-    return RESTORATION_TILESIZE_BIG;
-}
-
-static INLINE int av1_get_rest_ntiles(int width, int height, int *tile_width,
-                                      int *tile_height, int *nhtiles,
-                                      int *nvtiles) {
+static INLINE int av1_get_rest_ntiles(int width, int height, int tilesize,
+                                      int *tile_width, int *tile_height,
+                                      int *nhtiles, int *nvtiles) {
   int nhtiles_, nvtiles_;
   int tile_width_, tile_height_;
-  int tilesize = get_rest_tilesize(width, height);
   tile_width_ = (tilesize < 0) ? width : AOMMIN(tilesize, width);
   tile_height_ = (tilesize < 0) ? height : AOMMIN(tilesize, height);
   nhtiles_ = (width + (tile_width_ >> 1)) / tile_width_;

av1/decoder/decodeframe.c

@@ -2560,6 +2560,21 @@ static void decode_restoration_mode(AV1_COMMON *cm,
     cm->rst_info[p].frame_restoration_type =
         aom_rb_read_bit(rb) ? RESTORE_WIENER : RESTORE_NONE;
   }
+
+  cm->rst_info[0].restoration_tilesize = RESTORATION_TILESIZE_MAX;
+  cm->rst_info[1].restoration_tilesize = RESTORATION_TILESIZE_MAX;
+  cm->rst_info[2].restoration_tilesize = RESTORATION_TILESIZE_MAX;
+  if (cm->rst_info[0].frame_restoration_type != RESTORE_NONE ||
+      cm->rst_info[1].frame_restoration_type != RESTORE_NONE ||
+      cm->rst_info[2].frame_restoration_type != RESTORE_NONE) {
+    rsi = &cm->rst_info[0];
+    rsi->restoration_tilesize >>= aom_rb_read_bit(rb);
+    if (rsi->restoration_tilesize != RESTORATION_TILESIZE_MAX) {
+      rsi->restoration_tilesize >>= aom_rb_read_bit(rb);
+    }
+    cm->rst_info[1].restoration_tilesize = cm->rst_info[0].restoration_tilesize;
+    cm->rst_info[2].restoration_tilesize = cm->rst_info[0].restoration_tilesize;
+  }
 }
 
 static void read_wiener_filter(WienerInfo *wiener_info, aom_reader *rb) {
@@ -2610,12 +2625,13 @@ static void read_domaintxfmrf_filter(DomaintxfmrfInfo *domaintxfmrf_info,
 
 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(ROUND_POWER_OF_TWO(cm->width, cm->subsampling_x),
-                          ROUND_POWER_OF_TWO(cm->height, cm->subsampling_y),
-                          NULL, NULL, NULL, NULL);
+  const int ntiles = av1_get_rest_ntiles(cm->width, cm->height,
+                                         cm->rst_info[0].restoration_tilesize,
+                                         NULL, NULL, NULL, NULL);
+  const int ntiles_uv = av1_get_rest_ntiles(
+      ROUND_POWER_OF_TWO(cm->width, cm->subsampling_x),
+      ROUND_POWER_OF_TWO(cm->height, cm->subsampling_y),
+      cm->rst_info[1].restoration_tilesize, NULL, NULL, NULL, NULL);
   RestorationInfo *rsi = &cm->rst_info[0];
   if (rsi->frame_restoration_type != RESTORE_NONE) {
     if (rsi->frame_restoration_type == RESTORE_SWITCHABLE) {
@@ -4376,7 +4392,6 @@ static size_t read_uncompressed_header(AV1Decoder *pbi,
   setup_clpf(pbi, rb);
 #endif
 #if CONFIG_LOOP_RESTORATION
-  av1_alloc_restoration_buffers(cm);
   decode_restoration_mode(cm, rb);
 #endif  // CONFIG_LOOP_RESTORATION
   setup_quantization(cm, rb);
@@ -4579,7 +4594,12 @@ static int read_compressed_header(AV1Decoder *pbi, const uint8_t *data,
                        "Failed to allocate bool decoder 0");
 
 #if CONFIG_LOOP_RESTORATION
-  decode_restoration(cm, &r);
+  if (cm->rst_info[0].frame_restoration_type != RESTORE_NONE ||
+      cm->rst_info[1].frame_restoration_type != RESTORE_NONE ||
+      cm->rst_info[2].frame_restoration_type != RESTORE_NONE) {
+    av1_alloc_restoration_buffers(cm);
+    decode_restoration(cm, &r);
+  }
 #endif
 
 #if !CONFIG_EC_ADAPT

av1/encoder/bitstream.c

@@ -3402,6 +3402,19 @@ static void encode_restoration_mode(AV1_COMMON *cm,
       default: assert(0);
     }
   }
+  if (cm->rst_info[0].frame_restoration_type != RESTORE_NONE ||
+      cm->rst_info[1].frame_restoration_type != RESTORE_NONE ||
+      cm->rst_info[2].frame_restoration_type != RESTORE_NONE) {
+    rsi = &cm->rst_info[0];
+    if (rsi->restoration_tilesize == RESTORATION_TILESIZE_MAX) {
+      aom_wb_write_bit(wb, 0);
+    }
+    aom_wb_write_bit(wb, rsi->restoration_tilesize != RESTORATION_TILESIZE_MAX);
+    if (rsi->restoration_tilesize != RESTORATION_TILESIZE_MAX) {
+      aom_wb_write_bit(
+          wb, rsi->restoration_tilesize != (RESTORATION_TILESIZE_MAX >> 1));
+    }
+  }
 }
 
 static void write_wiener_filter(WienerInfo *wiener_info, aom_writer *wb) {
@@ -3436,12 +3449,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);
-  const int ntiles_uv =
-      av1_get_rest_ntiles(ROUND_POWER_OF_TWO(cm->width, cm->subsampling_x),
-                          ROUND_POWER_OF_TWO(cm->height, cm->subsampling_y),
-                          NULL, NULL, NULL, NULL);
+  const int ntiles = av1_get_rest_ntiles(cm->width, cm->height,
+                                         cm->rst_info[0].restoration_tilesize,
+                                         NULL, NULL, NULL, NULL);
+  const int ntiles_uv = av1_get_rest_ntiles(
+      ROUND_POWER_OF_TWO(cm->width, cm->subsampling_x),
+      ROUND_POWER_OF_TWO(cm->height, cm->subsampling_y),
+      cm->rst_info[1].restoration_tilesize, NULL, NULL, NULL, NULL);
   RestorationInfo *rsi = &cm->rst_info[0];
   if (rsi->frame_restoration_type != RESTORE_NONE) {
     if (rsi->frame_restoration_type == RESTORE_SWITCHABLE) {

av1/encoder/encoder.c

@@ -751,9 +751,6 @@ static void alloc_util_frame_buffers(AV1_COMP *cpi) {
   } else {
     cpi->extra_rstbuf = NULL;
   }
-  for (int i = 0; i < MAX_MB_PLANE; ++i)
-    av1_alloc_restoration_struct(cm, &cpi->rst_search[i], cm->width,
-                                 cm->height);
 #endif  // CONFIG_LOOP_RESTORATION
 
   if (aom_realloc_frame_buffer(&cpi->scaled_source, cm->width, cm->height,
@@ -3912,6 +3909,17 @@ static void init_motion_estimation(AV1_COMP *cpi) {
   }
 }
 
+#if CONFIG_LOOP_RESTORATION
+static void set_restoration_tilesize(int width, int height,
+                                     RestorationInfo *rst) {
+  (void)width;
+  (void)height;
+  rst[0].restoration_tilesize = (RESTORATION_TILESIZE_MAX >> 1);
+  rst[1].restoration_tilesize = rst[0].restoration_tilesize;
+  rst[2].restoration_tilesize = rst[0].restoration_tilesize;
+}
+#endif  // CONFIG_LOOP_RESTORATION
+
 static void set_frame_size(AV1_COMP *cpi) {
   int ref_frame;
   AV1_COMMON *const cm = &cpi->common;
@@ -3968,11 +3976,20 @@ static void set_frame_size(AV1_COMP *cpi) {
     aom_internal_error(&cm->error, AOM_CODEC_MEM_ERROR,
                        "Failed to allocate frame buffer");
 
-  alloc_util_frame_buffers(cpi);
-  init_motion_estimation(cpi);
 #if CONFIG_LOOP_RESTORATION
+  set_restoration_tilesize(cm->width, cm->height, cm->rst_info);
+  for (int i = 0; i < MAX_MB_PLANE; ++i)
+    cm->rst_info[i].frame_restoration_type = RESTORE_NONE;
   av1_alloc_restoration_buffers(cm);
+  for (int i = 0; i < MAX_MB_PLANE; ++i) {
+    cpi->rst_search[i].restoration_tilesize =
+        cm->rst_info[i].restoration_tilesize;
+    av1_alloc_restoration_struct(cm, &cpi->rst_search[i], cm->width,
+                                 cm->height);
+  }
 #endif  // CONFIG_LOOP_RESTORATION
+  alloc_util_frame_buffers(cpi);
+  init_motion_estimation(cpi);
 
   for (ref_frame = LAST_FRAME; ref_frame <= ALTREF_FRAME; ++ref_frame) {
     RefBuffer *const ref_buf = &cm->frame_refs[ref_frame - LAST_FRAME];

av1/encoder/pickrst.c

@@ -138,8 +138,9 @@ static int64_t try_restoration_tile(const YV12_BUFFER_CONFIG *src,
     width = src->uv_crop_width;
     height = src->uv_crop_height;
   }
-  ntiles = av1_get_rest_ntiles(width, height, &tile_width, &tile_height,
-                               &nhtiles, &nvtiles);
+  ntiles = av1_get_rest_ntiles(
+      width, height, cm->rst_info[components_pattern > 1].restoration_tilesize,
+      &tile_width, &tile_height, &nhtiles, &nvtiles);
   (void)ntiles;
 
   av1_loop_restoration_frame(cm->frame_to_show, cm, rsi, components_pattern,
@@ -344,8 +345,9 @@ static double search_sgrproj(const YV12_BUFFER_CONFIG *src, AV1_COMP *cpi,
   int tile_idx, tile_width, tile_height, nhtiles, nvtiles;
   int h_start, h_end, v_start, v_end;
   // Allocate for the src buffer at high precision
-  const int ntiles = av1_get_rest_ntiles(cm->width, cm->height, &tile_width,
-                                         &tile_height, &nhtiles, &nvtiles);
+  const int ntiles = av1_get_rest_ntiles(
+      cm->width, cm->height, cm->rst_info[0].restoration_tilesize, &tile_width,
+      &tile_height, &nhtiles, &nvtiles);
   rsi->frame_restoration_type = RESTORE_SGRPROJ;
 
   for (tile_idx = 0; tile_idx < ntiles; ++tile_idx) {
@@ -557,8 +559,9 @@ static double search_domaintxfmrf(const YV12_BUFFER_CONFIG *src, AV1_COMP *cpi,
   RestorationInfo *rsi = &cpi->rst_search[0];
   int tile_idx, 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);
+  const int ntiles = av1_get_rest_ntiles(
+      cm->width, cm->height, cm->rst_info[0].restoration_tilesize, &tile_width,
+      &tile_height, &nhtiles, &nvtiles);
 
   rsi->frame_restoration_type = RESTORE_DOMAINTXFMRF;
 
@@ -964,8 +967,9 @@ static double search_wiener_uv(const YV12_BUFFER_CONFIG *src, AV1_COMP *cpi,
   double score;
   int tile_idx, tile_width, tile_height, nhtiles, nvtiles;
   int h_start, h_end, v_start, v_end;
-  const int ntiles = av1_get_rest_ntiles(width, height, &tile_width,
-                                         &tile_height, &nhtiles, &nvtiles);
+  const int ntiles =
+      av1_get_rest_ntiles(width, height, cm->rst_info[1].restoration_tilesize,
+                          &tile_width, &tile_height, &nhtiles, &nvtiles);
   assert(width == dgd->uv_crop_width);
   assert(height == dgd->uv_crop_height);
 
@@ -1101,8 +1105,9 @@ static double search_wiener(const YV12_BUFFER_CONFIG *src, AV1_COMP *cpi,
   double score;
   int tile_idx, tile_width, tile_height, nhtiles, nvtiles;
   int h_start, h_end, v_start, v_end;
-  const int ntiles = av1_get_rest_ntiles(width, height, &tile_width,
-                                         &tile_height, &nhtiles, &nvtiles);
+  const int ntiles =
+      av1_get_rest_ntiles(width, height, cm->rst_info[0].restoration_tilesize,
+                          &tile_width, &tile_height, &nhtiles, &nvtiles);
   assert(width == dgd->y_crop_width);
   assert(height == dgd->y_crop_height);
   assert(width == src->y_crop_width);
@@ -1214,8 +1219,9 @@ static double search_norestore(const YV12_BUFFER_CONFIG *src, AV1_COMP *cpi,
   AV1_COMMON *const cm = &cpi->common;
   int tile_idx, 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);
+  const int ntiles = av1_get_rest_ntiles(
+      cm->width, cm->height, cm->rst_info[0].restoration_tilesize, &tile_width,
+      &tile_height, &nhtiles, &nvtiles);
   (void)info;
   (void)dst_frame;
   (void)partial_frame;
@@ -1246,8 +1252,9 @@ static double search_switchable_restoration(
   MACROBLOCK *x = &cpi->td.mb;
   double cost_switchable = 0;
   int r, bits, tile_idx;
-  const int ntiles =
-      av1_get_rest_ntiles(cm->width, cm->height, NULL, NULL, NULL, NULL);
+  const int ntiles = av1_get_rest_ntiles(cm->width, cm->height,
+                                         cm->rst_info[0].restoration_tilesize,
+                                         NULL, NULL, NULL, NULL);
   (void)partial_frame;
 
   rsi->frame_restoration_type = RESTORE_SWITCHABLE;
@@ -1285,8 +1292,9 @@ void av1_pick_filter_restoration(const YV12_BUFFER_CONFIG *src, AV1_COMP *cpi,
   double best_cost_restore;
   RestorationType r, best_restore;
 
-  const int ntiles =
-      av1_get_rest_ntiles(cm->width, cm->height, NULL, NULL, NULL, NULL);
+  const int ntiles = av1_get_rest_ntiles(cm->width, cm->height,
+                                         cm->rst_info[0].restoration_tilesize,
+                                         NULL, NULL, NULL, NULL);
 
   for (r = 0; r < RESTORE_SWITCHABLE_TYPES; r++) {
     tile_cost[r] = (double *)aom_malloc(sizeof(*tile_cost[0]) * ntiles);