Pickrst: Size restoration line buffers to frame width

av2/common/alloccommon.c

@@ -296,6 +296,56 @@ void av2_alloc_cdef_buffers(AV2_COMMON *const cm,
                       cdef_info->allocated_mi_rows);
 }
 
+void av2_alloc_restoration_line_buffers(AV2_COMMON *const cm,
+                                        RestorationLineBuffers **rlbs_ptr) {
+  if (!*rlbs_ptr) {
+    CHECK_MEM_ERROR(cm, *rlbs_ptr, avm_calloc(1, sizeof(**rlbs_ptr)));
+  }
+
+  RestorationLineBuffers *rlbs = *rlbs_ptr;
+  const int frame_w = cm->mi_params.mi_cols << MI_SIZE_LOG2;
+  for (PLANE_TYPE ptype = PLANE_TYPE_Y;
+       ptype <= (cm->seq_params.monochrome ? PLANE_TYPE_Y : PLANE_TYPE_UV);
+       ++ptype) {
+    // allocate for line buffer width.
+    const int ss_x = ptype && cm->seq_params.subsampling_x;
+    const int plane_w =
+        ((frame_w + ss_x) >> ss_x) + 2 * RESTORATION_BORDER_HORZ;
+    const int stride = ALIGN_POWER_OF_TWO(plane_w, 5);
+
+    if (stride > rlbs->line_stride[ptype]) {
+      for (int i = 0; i < RESTORATION_BORDER_VERT; ++i) {
+        avm_free(rlbs->tmp_save_above[ptype][i]);
+        rlbs->tmp_save_above[ptype][i] = NULL;
+        CHECK_MEM_ERROR(
+            cm, rlbs->tmp_save_above[ptype][i],
+            avm_memalign(32, stride * sizeof(*rlbs->tmp_save_above[ptype][i])));
+
+        avm_free(rlbs->tmp_save_below[ptype][i]);
+        rlbs->tmp_save_below[ptype][i] = NULL;
+        CHECK_MEM_ERROR(
+            cm, rlbs->tmp_save_below[ptype][i],
+            avm_memalign(32, stride * sizeof(*rlbs->tmp_save_below[ptype][i])));
+      }
+      rlbs->line_stride[ptype] = stride;
+    }
+  }
+}
+
+void av2_free_restoration_line_buffers(RestorationLineBuffers *rlbs) {
+  if (rlbs) {
+    for (PLANE_TYPE ptype = PLANE_TYPE_Y; ptype < PLANE_TYPES; ++ptype) {
+      for (int i = 0; i < RESTORATION_BORDER_VERT; ++i) {
+        avm_free(rlbs->tmp_save_above[ptype][i]);
+        rlbs->tmp_save_above[ptype][i] = NULL;
+        avm_free(rlbs->tmp_save_below[ptype][i]);
+        rlbs->tmp_save_below[ptype][i] = NULL;
+      }
+    }
+    avm_free(rlbs);
+  }
+}
+
 // Assumes cm->rst_info[p].restoration_unit_size is already initialized
 void av2_alloc_restoration_buffers(AV2_COMMON *cm) {
   const int num_planes = av2_num_planes(cm);
@@ -307,9 +357,8 @@ void av2_alloc_restoration_buffers(AV2_COMMON *cm) {
     translate_pcwiener_filters_to_wienerns(cm);
   }
 
-  if (cm->rlbs == NULL) {
-    CHECK_MEM_ERROR(cm, cm->rlbs, avm_malloc(sizeof(RestorationLineBuffers)));
-  }
+  av2_alloc_restoration_line_buffers(cm, &cm->rlbs);
+
   if (cm->lru_stripe_buf == NULL) {
     CHECK_MEM_ERROR(
         cm, cm->lru_stripe_buf,
@@ -373,7 +422,7 @@ void av2_free_restoration_buffers(AV2_COMMON *cm) {
   int p;
   for (p = 0; p < MAX_MB_PLANE; ++p)
     av2_free_restoration_struct(&cm->rst_info[p]);
-  avm_free(cm->rlbs);
+  av2_free_restoration_line_buffers(cm->rlbs);
   cm->rlbs = NULL;
   avm_free(cm->lru_stripe_buf);
   cm->lru_stripe_buf = NULL;

av2/common/alloccommon.h

@@ -27,6 +27,7 @@ struct CommonContexts;
 struct CommonModeInfoParams;
 struct AV2CdefWorker;
 struct AV2CdefSyncData;
+struct RestorationLineBuffers;
 
 void av2_remove_common(struct AV2Common *cm);
 
@@ -56,6 +57,11 @@ void av2_alloc_restoration_boundary_buffers(struct AV2Common *cm,
                                             int num_planes);
 void av2_free_restoration_buffers(struct AV2Common *cm);
 
+void av2_alloc_restoration_line_buffers(
+    struct AV2Common *const cm, struct RestorationLineBuffers **rlbs_ptr);
+
+void av2_free_restoration_line_buffers(struct RestorationLineBuffers *rlbs);
+
 int av2_get_MBs(int width, int height);
 
 #ifdef __cplusplus

av2/common/restoration.c

@@ -1665,13 +1665,6 @@ uint16_t *wienerns_copy_luma_with_virtual_lines(struct AV2Common *cm,
   int height_y = frame_buf->heights[AVM_PLANE_Y];
   int width_uv = frame_buf->widths[1];
   int height_uv = frame_buf->heights[1];
-
-  if (width_y > RESTORATION_LINEBUFFER_WIDTH)
-    avm_internal_error(
-        &cm->error, AVM_CODEC_ERROR,
-        "picture width is larger than 8192 * 8, need to disable "
-        "cross-component wienerns in this software implementation");
-
   int in_stride = frame_buf->strides[AVM_PLANE_Y];
   int border = WIENERNS_UV_BRD;
   int resized_luma_stride = width_uv + 2 * WIENERNS_UV_BRD;

av2/common/restoration.h

@@ -260,23 +260,12 @@ typedef struct {
 #define RESTORATION_LINEBUFFER_WIDTH \
   (MAX_SUPPORTED_PIC_WIDTH_IN_CCALF_IMP * 3 / 2 + 2 * RESTORATION_BORDER_HORZ)
 
-// Similarly, the column buffers (used when we're at a vertical tile edge
-// that we can't filter across) need space for one processing unit's worth
-// of pixels, plus the top/bottom border width
-#define RESTORATION_COLBUFFER_HEIGHT \
-  (RESTORATION_PROC_UNIT_SIZE + 2 * RESTORATION_BORDER_VERT)
-
-typedef struct {
+typedef struct RestorationLineBuffers {
   // Temporary buffers to save/restore 3 lines above/below the restoration
   // stripe.
-  uint16_t tmp_save_above[2][RESTORATION_BORDER_VERT]
-                         [RESTORATION_LINEBUFFER_WIDTH];
-  uint16_t tmp_save_below[2][RESTORATION_BORDER_VERT]
-                         [RESTORATION_LINEBUFFER_WIDTH];
-  uint16_t tmp_save_left[2][RESTORATION_COLBUFFER_HEIGHT]
-                        [RESTORATION_BORDER_HORZ];
-  uint16_t tmp_save_right[2][RESTORATION_COLBUFFER_HEIGHT]
-                         [RESTORATION_BORDER_HORZ];
+  uint16_t *tmp_save_above[PLANE_TYPES][RESTORATION_BORDER_VERT];
+  uint16_t *tmp_save_below[PLANE_TYPES][RESTORATION_BORDER_VERT];
+  int line_stride[PLANE_TYPES];
 } RestorationLineBuffers;
 /*!\endcond */
 

av2/encoder/pickrst.c

@@ -142,6 +142,7 @@ typedef struct {
 
   int luma_stride;
 
+  RestorationLineBuffers *rlbs;
   // Temporary storage used by *wienerns_filter* functions.
   double *wienerns_tmpbuf;
 
@@ -325,9 +326,7 @@ static int64_t try_restoration_unit(const RestSearchCtxt *rsc,
   const int plane = rsc->plane;
   const int is_uv = plane > 0;
   const RestorationInfo *rsi = &cm->rst_info[plane];
-  RestorationLineBuffers *rlbs = avm_malloc(sizeof(RestorationLineBuffers));
-  if (rlbs == NULL)
-    fprintf(stderr, "rlbs buffer does not allocate successfully\n");
+  RestorationLineBuffers *rlbs = rsc->rlbs;
   const int bit_depth = cm->seq_params.bit_depth;
 
   const YV12_BUFFER_CONFIG *fts = &cm->cur_frame->buf;
@@ -342,7 +341,6 @@ static int64_t try_restoration_unit(const RestSearchCtxt *rsc,
       rsc->plane_width, cm->seq_params.disable_loopfilters_across_tiles,
       optimized_lr);
 
-  if (rlbs != NULL) avm_free(rlbs);
   return sse_restoration_unit(limits, rsc->src, rsc->dst, plane);
 }
 
@@ -3845,6 +3843,9 @@ void av2_pick_filter_restoration(const YV12_BUFFER_CONFIG *src, AV2_COMP *cpi) {
   uint16_t *luma_virtual = NULL;
   uint16_t *luma_virtual_buf;
 
+  rsc.rlbs = NULL;
+  av2_alloc_restoration_line_buffers(cm, &rsc.rlbs);
+
   luma_virtual_buf = wienerns_copy_luma_with_virtual_lines(cm, &luma_virtual);
   rsc.luma = luma_virtual;
 
@@ -4027,6 +4028,7 @@ void av2_pick_filter_restoration(const YV12_BUFFER_CONFIG *src, AV2_COMP *cpi) {
   avm_free(rusi);
   free(luma_buf);
   free(luma_virtual_buf);
+  av2_free_restoration_line_buffers(rsc.rlbs);
   avm_free(rsc.wienerns_tmpbuf);
   avm_vector_destroy(&wienerns_stats);
   avm_vector_destroy(&unit_stack);