1 files changed, 414 insertions, 0 deletions
diff --git a/third_party/aom/aom_dsp/pyramid.c b/third_party/aom/aom_dsp/pyramid.c
new file mode 100644
index 0000000000..324a18baea
--- /dev/null
+++ b/third_party/aom/aom_dsp/pyramid.c
@@ -0,0 +1,414 @@
+/*
+ * Copyright (c) 2022, Alliance for Open Media. All rights reserved
+ *
+ * This source code is subject to the terms of the BSD 2 Clause License and
+ * the Alliance for Open Media Patent License 1.0. If the BSD 2 Clause License
+ * was not distributed with this source code in the LICENSE file, you can
+ * obtain it at www.aomedia.org/license/software. If the Alliance for Open
+ * Media Patent License 1.0 was not distributed with this source code in the
+ * PATENTS file, you can obtain it at www.aomedia.org/license/patent.
+ */
+
+#include "aom_dsp/pyramid.h"
+#include "aom_mem/aom_mem.h"
+#include "aom_ports/bitops.h"
+#include "aom_util/aom_thread.h"
+
+// TODO(rachelbarker): Move needed code from av1/ to aom_dsp/
+#include "av1/common/resize.h"
+
+#include <assert.h>
+#include <string.h>
+
+// Lifecycle:
+// * Frame buffer alloc code calls aom_get_pyramid_alloc_size()
+//   to work out how much space is needed for a given number of pyramid
+//   levels. This is counted in the size checked against the max allocation
+//   limit
+// * Then calls aom_alloc_pyramid() to actually create the pyramid
+// * Pyramid is initially marked as invalid (no data)
+// * Whenever pyramid is needed, we check the valid flag. If set, use existing
+//   data. If not set, compute full pyramid
+// * Whenever frame buffer is reused, clear the valid flag
+// * Whenever frame buffer is resized, reallocate pyramid
+
+size_t aom_get_pyramid_alloc_size(int width, int height, int n_levels,
+                                  bool image_is_16bit) {
+  // Limit number of levels on small frames
+  const int msb = get_msb(AOMMIN(width, height));
+  const int max_levels = AOMMAX(msb - MIN_PYRAMID_SIZE_LOG2, 1);
+  n_levels = AOMMIN(n_levels, max_levels);
+
+  size_t alloc_size = 0;
+  alloc_size += sizeof(ImagePyramid);
+  alloc_size += n_levels * sizeof(PyramidLayer);
+
+  // Calculate how much memory is needed for downscaled frame buffers
+  size_t buffer_size = 0;
+
+  // Work out if we need to allocate a few extra bytes for alignment.
+  // aom_memalign() will ensure that the start of the allocation is aligned
+  // to a multiple of PYRAMID_ALIGNMENT. But we want the first image pixel
+  // to be aligned, not the first byte of the allocation.
+  //
+  // In the loop below, we ensure that the stride of every image is a multiple
+  // of PYRAMID_ALIGNMENT. Thus the allocated size of each pyramid level will
+  // also be a multiple of PYRAMID_ALIGNMENT. Thus, as long as we can get the
+  // first pixel in the first pyramid layer aligned properly, that will
+  // automatically mean that the first pixel of every row of every layer is
+  // properly aligned too.
+  //
+  // Thus all we need to consider is the first pixel in the first layer.
+  // This is located at offset
+  //   extra_bytes + level_stride * PYRAMID_PADDING + PYRAMID_PADDING
+  // bytes into the buffer. Since level_stride is a multiple of
+  // PYRAMID_ALIGNMENT, we can ignore that. So we need
+  //   extra_bytes + PYRAMID_PADDING = multiple of PYRAMID_ALIGNMENT
+  //
+  // To solve this, we can round PYRAMID_PADDING up to the next multiple
+  // of PYRAMID_ALIGNMENT, then subtract the orginal value to calculate
+  // how many extra bytes are needed.
+  size_t first_px_offset =
+      (PYRAMID_PADDING + PYRAMID_ALIGNMENT - 1) & ~(PYRAMID_ALIGNMENT - 1);
+  size_t extra_bytes = first_px_offset - PYRAMID_PADDING;
+  buffer_size += extra_bytes;
+
+  // If the original image is stored in an 8-bit buffer, then we can point the
+  // lowest pyramid level at that buffer rather than allocating a new one.
+  int first_allocated_level = image_is_16bit ? 0 : 1;
+
+  for (int level = first_allocated_level; level < n_levels; level++) {
+    int level_width = width >> level;
+    int level_height = height >> level;
+
+    // Allocate padding for each layer
+    int padded_width = level_width + 2 * PYRAMID_PADDING;
+    int padded_height = level_height + 2 * PYRAMID_PADDING;
+
+    // Align the layer stride to be a multiple of PYRAMID_ALIGNMENT
+    // This ensures that, as long as the top-left pixel in this pyramid level is
+    // properly aligned, then so will the leftmost pixel in every row of the
+    // pyramid level.
+    int level_stride =
+        (padded_width + PYRAMID_ALIGNMENT - 1) & ~(PYRAMID_ALIGNMENT - 1);
+
+    buffer_size += level_stride * padded_height;
+  }
+
+  alloc_size += buffer_size;
+
+  return alloc_size;
+}
+
+ImagePyramid *aom_alloc_pyramid(int width, int height, int n_levels,
+                                bool image_is_16bit) {
+  // Limit number of levels on small frames
+  const int msb = get_msb(AOMMIN(width, height));
+  const int max_levels = AOMMAX(msb - MIN_PYRAMID_SIZE_LOG2, 1);
+  n_levels = AOMMIN(n_levels, max_levels);
+
+  ImagePyramid *pyr = aom_calloc(1, sizeof(*pyr));
+  if (!pyr) {
+    return NULL;
+  }
+
+  pyr->layers = aom_calloc(n_levels, sizeof(*pyr->layers));
+  if (!pyr->layers) {
+    aom_free(pyr);
+    return NULL;
+  }
+
+  pyr->valid = false;
+  pyr->n_levels = n_levels;
+
+  // Compute sizes and offsets for each pyramid level
+  // These are gathered up first, so that we can allocate all pyramid levels
+  // in a single buffer
+  size_t buffer_size = 0;
+  size_t *layer_offsets = aom_calloc(n_levels, sizeof(*layer_offsets));
+  if (!layer_offsets) {
+    aom_free(pyr->layers);
+    aom_free(pyr);
+    return NULL;
+  }
+
+  // Work out if we need to allocate a few extra bytes for alignment.
+  // aom_memalign() will ensure that the start of the allocation is aligned
+  // to a multiple of PYRAMID_ALIGNMENT. But we want the first image pixel
+  // to be aligned, not the first byte of the allocation.
+  //
+  // In the loop below, we ensure that the stride of every image is a multiple
+  // of PYRAMID_ALIGNMENT. Thus the allocated size of each pyramid level will
+  // also be a multiple of PYRAMID_ALIGNMENT. Thus, as long as we can get the
+  // first pixel in the first pyramid layer aligned properly, that will
+  // automatically mean that the first pixel of every row of every layer is
+  // properly aligned too.
+  //
+  // Thus all we need to consider is the first pixel in the first layer.
+  // This is located at offset
+  //   extra_bytes + level_stride * PYRAMID_PADDING + PYRAMID_PADDING
+  // bytes into the buffer. Since level_stride is a multiple of
+  // PYRAMID_ALIGNMENT, we can ignore that. So we need
+  //   extra_bytes + PYRAMID_PADDING = multiple of PYRAMID_ALIGNMENT
+  //
+  // To solve this, we can round PYRAMID_PADDING up to the next multiple
+  // of PYRAMID_ALIGNMENT, then subtract the orginal value to calculate
+  // how many extra bytes are needed.
+  size_t first_px_offset =
+      (PYRAMID_PADDING + PYRAMID_ALIGNMENT - 1) & ~(PYRAMID_ALIGNMENT - 1);
+  size_t extra_bytes = first_px_offset - PYRAMID_PADDING;
+  buffer_size += extra_bytes;
+
+  // If the original image is stored in an 8-bit buffer, then we can point the
+  // lowest pyramid level at that buffer rather than allocating a new one.
+  int first_allocated_level = image_is_16bit ? 0 : 1;
+
+  for (int level = first_allocated_level; level < n_levels; level++) {
+    PyramidLayer *layer = &pyr->layers[level];
+
+    int level_width = width >> level;
+    int level_height = height >> level;
+
+    // Allocate padding for each layer
+    int padded_width = level_width + 2 * PYRAMID_PADDING;
+    int padded_height = level_height + 2 * PYRAMID_PADDING;
+
+    // Align the layer stride to be a multiple of PYRAMID_ALIGNMENT
+    // This ensures that, as long as the top-left pixel in this pyramid level is
+    // properly aligned, then so will the leftmost pixel in every row of the
+    // pyramid level.
+    int level_stride =
+        (padded_width + PYRAMID_ALIGNMENT - 1) & ~(PYRAMID_ALIGNMENT - 1);
+
+    size_t level_alloc_start = buffer_size;
+    size_t level_start =
+        level_alloc_start + PYRAMID_PADDING * level_stride + PYRAMID_PADDING;
+
+    buffer_size += level_stride * padded_height;
+
+    layer_offsets[level] = level_start;
+    layer->width = level_width;
+    layer->height = level_height;
+    layer->stride = level_stride;
+  }
+
+  pyr->buffer_alloc =
+      aom_memalign(PYRAMID_ALIGNMENT, buffer_size * sizeof(*pyr->buffer_alloc));
+  if (!pyr->buffer_alloc) {
+    aom_free(pyr->layers);
+    aom_free(pyr);
+    aom_free(layer_offsets);
+    return NULL;
+  }
+
+  // Fill in pointers for each level
+  // If image is 8-bit, then the lowest level is left unconfigured for now,
+  // and will be set up properly when the pyramid is filled in
+  for (int level = first_allocated_level; level < n_levels; level++) {
+    PyramidLayer *layer = &pyr->layers[level];
+    layer->buffer = pyr->buffer_alloc + layer_offsets[level];
+  }
+
+#if CONFIG_MULTITHREAD
+  pthread_mutex_init(&pyr->mutex, NULL);
+#endif  // CONFIG_MULTITHREAD
+
+  aom_free(layer_offsets);
+  return pyr;
+}
+
+// Fill the border region of a pyramid frame.
+// This must be called after the main image area is filled out.
+// `img_buf` should point to the first pixel in the image area,
+// ie. it should be pyr->level_buffer + pyr->level_loc[level].
+static INLINE void fill_border(uint8_t *img_buf, const int width,
+                               const int height, const int stride) {
+  // Fill left and right areas
+  for (int row = 0; row < height; row++) {
+    uint8_t *row_start = &img_buf[row * stride];
+    uint8_t left_pixel = row_start[0];
+    memset(row_start - PYRAMID_PADDING, left_pixel, PYRAMID_PADDING);
+    uint8_t right_pixel = row_start[width - 1];
+    memset(row_start + width, right_pixel, PYRAMID_PADDING);
+  }
+
+  // Fill top area
+  for (int row = -PYRAMID_PADDING; row < 0; row++) {
+    uint8_t *row_start = &img_buf[row * stride];
+    memcpy(row_start - PYRAMID_PADDING, img_buf - PYRAMID_PADDING,
+           width + 2 * PYRAMID_PADDING);
+  }
+
+  // Fill bottom area
+  uint8_t *last_row_start = &img_buf[(height - 1) * stride];
+  for (int row = height; row < height + PYRAMID_PADDING; row++) {
+    uint8_t *row_start = &img_buf[row * stride];
+    memcpy(row_start - PYRAMID_PADDING, last_row_start - PYRAMID_PADDING,
+           width + 2 * PYRAMID_PADDING);
+  }
+}
+
+// Compute coarse to fine pyramids for a frame
+// This must only be called while holding frame_pyr->mutex
+static INLINE bool fill_pyramid(const YV12_BUFFER_CONFIG *frame, int bit_depth,
+                                ImagePyramid *frame_pyr) {
+  int n_levels = frame_pyr->n_levels;
+  const int frame_width = frame->y_crop_width;
+  const int frame_height = frame->y_crop_height;
+  const int frame_stride = frame->y_stride;
+  assert((frame_width >> n_levels) >= 0);
+  assert((frame_height >> n_levels) >= 0);
+
+  PyramidLayer *first_layer = &frame_pyr->layers[0];
+  if (frame->flags & YV12_FLAG_HIGHBITDEPTH) {
+    // For frames stored in a 16-bit buffer, we need to downconvert to 8 bits
+    assert(first_layer->width == frame_width);
+    assert(first_layer->height == frame_height);
+
+    uint16_t *frame_buffer = CONVERT_TO_SHORTPTR(frame->y_buffer);
+    uint8_t *pyr_buffer = first_layer->buffer;
+    int pyr_stride = first_layer->stride;
+    for (int y = 0; y < frame_height; y++) {
+      uint16_t *frame_row = frame_buffer + y * frame_stride;
+      uint8_t *pyr_row = pyr_buffer + y * pyr_stride;
+      for (int x = 0; x < frame_width; x++) {
+        pyr_row[x] = frame_row[x] >> (bit_depth - 8);
+      }
+    }
+
+    fill_border(pyr_buffer, frame_width, frame_height, pyr_stride);
+  } else {
+    // For frames stored in an 8-bit buffer, we need to configure the first
+    // pyramid layer to point at the original image buffer
+    first_layer->buffer = frame->y_buffer;
+    first_layer->width = frame_width;
+    first_layer->height = frame_height;
+    first_layer->stride = frame_stride;
+  }
+
+  // Fill in the remaining levels through progressive downsampling
+  for (int level = 1; level < n_levels; ++level) {
+    PyramidLayer *prev_layer = &frame_pyr->layers[level - 1];
+    uint8_t *prev_buffer = prev_layer->buffer;
+    int prev_stride = prev_layer->stride;
+
+    PyramidLayer *this_layer = &frame_pyr->layers[level];
+    uint8_t *this_buffer = this_layer->buffer;
+    int this_width = this_layer->width;
+    int this_height = this_layer->height;
+    int this_stride = this_layer->stride;
+
+    // Compute the this pyramid level by downsampling the current level.
+    //
+    // We downsample by a factor of exactly 2, clipping the rightmost and
+    // bottommost pixel off of the current level if needed. We do this for
+    // two main reasons:
+    //
+    // 1) In the disflow code, when stepping from a higher pyramid level to a
+    //    lower pyramid level, we need to not just interpolate the flow field
+    //    but also to scale each flow vector by the upsampling ratio.
+    //    So it is much more convenient if this ratio is simply 2.
+    //
+    // 2) Up/downsampling by a factor of 2 can be implemented much more
+    //    efficiently than up/downsampling by a generic ratio.
+    //    TODO(rachelbarker): Use optimized downsample-by-2 function
+    if (!av1_resize_plane(prev_buffer, this_height << 1, this_width << 1,
+                          prev_stride, this_buffer, this_height, this_width,
+                          this_stride))
+      return false;
+    fill_border(this_buffer, this_width, this_height, this_stride);
+  }
+  return true;
+}
+
+// Fill out a downsampling pyramid for a given frame.
+//
+// The top level (index 0) will always be an 8-bit copy of the input frame,
+// regardless of the input bit depth. Additional levels are then downscaled
+// by powers of 2.
+//
+// For small input frames, the number of levels actually constructed
+// will be limited so that the smallest image is at least MIN_PYRAMID_SIZE
+// pixels along each side.
+//
+// However, if the input frame has a side of length < MIN_PYRAMID_SIZE,
+// we will still construct the top level.
+bool aom_compute_pyramid(const YV12_BUFFER_CONFIG *frame, int bit_depth,
+                         ImagePyramid *pyr) {
+  assert(pyr);
+
+  // Per the comments in the ImagePyramid struct, we must take this mutex
+  // before reading or writing the "valid" flag, and hold it while computing
+  // the pyramid, to ensure proper behaviour if multiple threads call this
+  // function simultaneously
+#if CONFIG_MULTITHREAD
+  pthread_mutex_lock(&pyr->mutex);
+#endif  // CONFIG_MULTITHREAD
+
+  if (!pyr->valid) {
+    pyr->valid = fill_pyramid(frame, bit_depth, pyr);
+  }
+  bool valid = pyr->valid;
+
+  // At this point, the pyramid is guaranteed to be valid, and can be safely
+  // read from without holding the mutex any more
+
+#if CONFIG_MULTITHREAD
+  pthread_mutex_unlock(&pyr->mutex);
+#endif  // CONFIG_MULTITHREAD
+  return valid;
+}
+
+#ifndef NDEBUG
+// Check if a pyramid has already been computed.
+// This is mostly a debug helper - as it is necessary to hold pyr->mutex
+// while reading the valid flag, we cannot just write:
+//   assert(pyr->valid);
+// This function allows the check to be correctly written as:
+//   assert(aom_is_pyramid_valid(pyr));
+bool aom_is_pyramid_valid(ImagePyramid *pyr) {
+  assert(pyr);
+
+  // Per the comments in the ImagePyramid struct, we must take this mutex
+  // before reading or writing the "valid" flag, and hold it while computing
+  // the pyramid, to ensure proper behaviour if multiple threads call this
+  // function simultaneously
+#if CONFIG_MULTITHREAD
+  pthread_mutex_lock(&pyr->mutex);
+#endif  // CONFIG_MULTITHREAD
+
+  bool valid = pyr->valid;
+
+#if CONFIG_MULTITHREAD
+  pthread_mutex_unlock(&pyr->mutex);
+#endif  // CONFIG_MULTITHREAD
+
+  return valid;
+}
+#endif
+
+// Mark a pyramid as no longer containing valid data.
+// This must be done whenever the corresponding frame buffer is reused
+void aom_invalidate_pyramid(ImagePyramid *pyr) {
+  if (pyr) {
+#if CONFIG_MULTITHREAD
+    pthread_mutex_lock(&pyr->mutex);
+#endif  // CONFIG_MULTITHREAD
+    pyr->valid = false;
+#if CONFIG_MULTITHREAD
+    pthread_mutex_unlock(&pyr->mutex);
+#endif  // CONFIG_MULTITHREAD
+  }
+}
+
+// Release the memory associated with a pyramid
+void aom_free_pyramid(ImagePyramid *pyr) {
+  if (pyr) {
+#if CONFIG_MULTITHREAD
+    pthread_mutex_destroy(&pyr->mutex);
+#endif  // CONFIG_MULTITHREAD
+    aom_free(pyr->buffer_alloc);
+    aom_free(pyr->layers);
+    aom_free(pyr);
+  }
+}