From 52bd722bfa59d7df80a994bf0cdeb1f30119398e Mon Sep 17 00:00:00 2001
From: Rachel Barker <[EMAIL REDACTED]>
Date: Fri, 9 Feb 2024 17:22:35 +0000
Subject: [PATCH] Add AVX2 disflow code
AVX2 implementation is based on the SSE4.1 implementation, but
processes 2 rows at a time.
Speedup relative to C code is 10.53x
Speedup relative to SSE4.1 code is 1.29x
Change-Id: I8e7343668d85113c138ae93bff423ba698a1bcd8
---
aom_dsp/aom_dsp.cmake | 3 +-
aom_dsp/aom_dsp_rtcd_defs.pl | 2 +-
aom_dsp/flow_estimation/x86/disflow_avx2.c | 418 +++++++++++++++++++++
aom_dsp/x86/synonyms.h | 7 +
aom_dsp/x86/synonyms_avx2.h | 10 +
test/disflow_test.cc | 5 +
6 files changed, 443 insertions(+), 2 deletions(-)
create mode 100644 aom_dsp/flow_estimation/x86/disflow_avx2.c
diff --git a/aom_dsp/aom_dsp.cmake b/aom_dsp/aom_dsp.cmake
index d9fd7f777b..de987cbd23 100644
--- a/aom_dsp/aom_dsp.cmake
+++ b/aom_dsp/aom_dsp.cmake
@@ -200,7 +200,8 @@ if(CONFIG_AV1_ENCODER)
"${AOM_ROOT}/aom_dsp/flow_estimation/x86/disflow_sse4.c")
list(APPEND AOM_DSP_ENCODER_INTRIN_AVX2
- "${AOM_ROOT}/aom_dsp/flow_estimation/x86/corner_match_avx2.c")
+ "${AOM_ROOT}/aom_dsp/flow_estimation/x86/corner_match_avx2.c"
+ "${AOM_ROOT}/aom_dsp/flow_estimation/x86/disflow_avx2.c")
list(APPEND AOM_DSP_ENCODER_INTRIN_NEON
"${AOM_ROOT}/aom_dsp/flow_estimation/arm/disflow_neon.c")
diff --git a/aom_dsp/aom_dsp_rtcd_defs.pl b/aom_dsp/aom_dsp_rtcd_defs.pl
index 62ad9f4ce8..28c0c59323 100755
--- a/aom_dsp/aom_dsp_rtcd_defs.pl
+++ b/aom_dsp/aom_dsp_rtcd_defs.pl
@@ -1796,7 +1796,7 @@ ()
specialize qw/aom_compute_correlation sse4_1 avx2/;
add_proto qw/void aom_compute_flow_at_point/, "const uint8_t *src, const uint8_t *ref, int x, int y, int width, int height, int stride, double *u, double *v";
- specialize qw/aom_compute_flow_at_point sse4_1 neon/;
+ specialize qw/aom_compute_flow_at_point sse4_1 avx2 neon/;
}
} # CONFIG_AV1_ENCODER
diff --git a/aom_dsp/flow_estimation/x86/disflow_avx2.c b/aom_dsp/flow_estimation/x86/disflow_avx2.c
new file mode 100644
index 0000000000..e67017b173
--- /dev/null
+++ b/aom_dsp/flow_estimation/x86/disflow_avx2.c
@@ -0,0 +1,418 @@
+/*
+ * Copyright (c) 2022, Alliance for Open Media. All rights reserved
+ *
+ * This source code is subject to the terms of the BSD 3-Clause Clear License
+ * and the Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear
+ * License was not distributed with this source code in the LICENSE file, you
+ * can obtain it at aomedia.org/license/software-license/bsd-3-c-c/. 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
+ * aomedia.org/license/patent-license/.
+ */
+
+#include <assert.h>
+#include <math.h>
+#include <immintrin.h>
+
+#include "aom_dsp/aom_dsp_common.h"
+#include "aom_dsp/flow_estimation/disflow.h"
+#include "aom_dsp/x86/synonyms.h"
+#include "aom_dsp/x86/synonyms_avx2.h"
+
+#include "config/aom_dsp_rtcd.h"
+
+#if DISFLOW_PATCH_SIZE != 8
+#error "Need to change disflow_avx2.c if DISFLOW_PATCH_SIZE != 8"
+#endif
+
+// Compute horizontal and vertical kernels and return them packed into a
+// register. The coefficient ordering is:
+// h0, h1, v0, v1, h2, h3, v2, v3
+// This is chosen because it takes less work than fully separating the kernels,
+// but it is separated enough that we can pick out each coefficient pair in the
+// main compute_flow_at_point function
+static INLINE __m128i compute_cubic_kernels(double u, double v) {
+ const __m128d x = _mm_set_pd(v, u);
+
+ const __m128d x2 = _mm_mul_pd(x, x);
+ const __m128d x3 = _mm_mul_pd(x2, x);
+
+ // Macro to multiply a value v by a constant coefficient c
+#define MULC(c, v) _mm_mul_pd(_mm_set1_pd(c), v)
+
+ // Compute floating-point kernel
+ // Note: To ensure results are bit-identical to the C code, we need to perform
+ // exactly the same sequence of operations here as in the C code.
+ __m128d k0 = _mm_sub_pd(_mm_add_pd(MULC(-0.5, x), x2), MULC(0.5, x3));
+ __m128d k1 =
+ _mm_add_pd(_mm_sub_pd(_mm_set1_pd(1.0), MULC(2.5, x2)), MULC(1.5, x3));
+ __m128d k2 =
+ _mm_sub_pd(_mm_add_pd(MULC(0.5, x), MULC(2.0, x2)), MULC(1.5, x3));
+ __m128d k3 = _mm_add_pd(MULC(-0.5, x2), MULC(0.5, x3));
+#undef MULC
+
+ // Integerize
+ __m128d prec = _mm_set1_pd((double)(1 << DISFLOW_INTERP_BITS));
+
+ k0 = _mm_round_pd(_mm_mul_pd(k0, prec),
+ _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);
+ k1 = _mm_round_pd(_mm_mul_pd(k1, prec),
+ _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);
+ k2 = _mm_round_pd(_mm_mul_pd(k2, prec),
+ _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);
+ k3 = _mm_round_pd(_mm_mul_pd(k3, prec),
+ _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);
+
+ const __m128i c0 = _mm_cvtpd_epi32(k0);
+ const __m128i c1 = _mm_cvtpd_epi32(k1);
+ const __m128i c2 = _mm_cvtpd_epi32(k2);
+ const __m128i c3 = _mm_cvtpd_epi32(k3);
+
+ // Rearrange results and convert down to 16 bits, giving the target output
+ // ordering
+ const __m128i c01 = _mm_unpacklo_epi32(c0, c1);
+ const __m128i c23 = _mm_unpacklo_epi32(c2, c3);
+ return _mm_packs_epi32(c01, c23);
+}
+
+// Compare two regions of width x height pixels, one rooted at position
+// (x, y) in src and the other at (x + u, y + v) in ref.
+// This function returns the sum of squared pixel differences between
+// the two regions.
+//
+// TODO(rachelbarker): Test speed/quality impact of using bilinear interpolation
+// instad of bicubic interpolation
+static INLINE void compute_flow_vector(const uint8_t *src, const uint8_t *ref,
+ int width, int height, int stride, int x,
+ int y, double u, double v,
+ const int16_t *dx, const int16_t *dy,
+ int *b) {
+ const __m256i zero = _mm256_setzero_si256();
+
+ // Accumulate 8 32-bit partial sums for each element of b
+ // These will be flattened at the end.
+ __m256i b0_acc = _mm256_setzero_si256();
+ __m256i b1_acc = _mm256_setzero_si256();
+
+ // Split offset into integer and fractional parts, and compute cubic
+ // interpolation kernels
+ const int u_int = (int)floor(u);
+ const int v_int = (int)floor(v);
+ const double u_frac = u - floor(u);
+ const double v_frac = v - floor(v);
+
+ const __m128i kernels = compute_cubic_kernels(u_frac, v_frac);
+
+ // Storage for intermediate values between the two convolution directions
+ // In the AVX2 implementation, this needs a dummy row at the end, because
+ // we generate 2 rows at a time but the total number of rows is odd.
+ // So we generate one more row than we actually need.
+ DECLARE_ALIGNED(32, int16_t,
+ tmp_[DISFLOW_PATCH_SIZE * (DISFLOW_PATCH_SIZE + 4)]);
+ int16_t *tmp = tmp_ + DISFLOW_PATCH_SIZE; // Offset by one row
+
+ // Clamp coordinates so that all pixels we fetch will remain within the
+ // allocated border region, but allow them to go far enough out that
+ // the border pixels' values do not change.
+ // Since we are calculating an 8x8 block, the bottom-right pixel
+ // in the block has coordinates (x0 + 7, y0 + 7). Then, the cubic
+ // interpolation has 4 taps, meaning that the output of pixel
+ // (x_w, y_w) depends on the pixels in the range
+ // ([x_w - 1, x_w + 2], [y_w - 1, y_w + 2]).
+ //
+ // Thus the most extreme coordinates which will be fetched are
+ // (x0 - 1, y0 - 1) and (x0 + 9, y0 + 9).
+ const int x0 = clamp(x + u_int, -9, width);
+ const int y0 = clamp(y + v_int, -9, height);
+
+ // Horizontal convolution
+
+ // Prepare the kernel vectors
+ // We split the kernel into two vectors with kernel indices:
+ // 0, 1, 0, 1, 0, 1, 0, 1, and
+ // 2, 3, 2, 3, 2, 3, 2, 3
+ __m256i h_kernel_01 = _mm256_broadcastd_epi32(kernels);
+ __m256i h_kernel_23 = _mm256_broadcastd_epi32(_mm_srli_si128(kernels, 8));
+
+ __m256i round_const_h = _mm256_set1_epi32(1 << (DISFLOW_INTERP_BITS - 6 - 1));
+
+ for (int i = -1; i < DISFLOW_PATCH_SIZE + 2; i += 2) {
+ const int y_w = y0 + i;
+ const uint8_t *ref_row = &ref[y_w * stride + (x0 - 1)];
+ int16_t *tmp_row = &tmp[i * DISFLOW_PATCH_SIZE];
+
+ // Load this row of pixels.
+ // For an 8x8 patch, we need to load the 8 image pixels + 3 extras,
+ // for a total of 11 pixels. Here we load 16 pixels, but only use
+ // the first 11.
+ __m256i row =
+ _mm256_loadu2_m128i((__m128i *)(ref_row + stride), (__m128i *)ref_row);
+
+ // Expand pixels to int16s
+ // We must use unpacks here, as we have one row in each 128-bit lane
+ // and want to handle each of those independently.
+ // This is in contrast to _mm256_cvtepu8_epi16(), which takes a single
+ // 128-bit input and widens it to 256 bits.
+ __m256i px_0to7_i16 = _mm256_unpacklo_epi8(row, zero);
+ __m256i px_4to10_i16 =
+ _mm256_unpacklo_epi8(_mm256_srli_si256(row, 4), zero);
+
+ // Compute first four outputs
+ // input pixels 0, 1, 1, 2, 2, 3, 3, 4
+ // * kernel 0, 1, 0, 1, 0, 1, 0, 1
+ __m256i px0 =
+ _mm256_unpacklo_epi16(px_0to7_i16, _mm256_srli_si256(px_0to7_i16, 2));
+ // input pixels 2, 3, 3, 4, 4, 5, 5, 6
+ // * kernel 2, 3, 2, 3, 2, 3, 2, 3
+ __m256i px1 = _mm256_unpacklo_epi16(_mm256_srli_si256(px_0to7_i16, 4),
+ _mm256_srli_si256(px_0to7_i16, 6));
+ // Convolve with kernel and sum 2x2 boxes to form first 4 outputs
+ __m256i sum0 = _mm256_add_epi32(_mm256_madd_epi16(px0, h_kernel_01),
+ _mm256_madd_epi16(px1, h_kernel_23));
+
+ __m256i out0 = _mm256_srai_epi32(_mm256_add_epi32(sum0, round_const_h),
+ DISFLOW_INTERP_BITS - 6);
+
+ // Compute second four outputs
+ __m256i px2 =
+ _mm256_unpacklo_epi16(px_4to10_i16, _mm256_srli_si256(px_4to10_i16, 2));
+ __m256i px3 = _mm256_unpacklo_epi16(_mm256_srli_si256(px_4to10_i16, 4),
+ _mm256_srli_si256(px_4to10_i16, 6));
+ __m256i sum1 = _mm256_add_epi32(_mm256_madd_epi16(px2, h_kernel_01),
+ _mm256_madd_epi16(px3, h_kernel_23));
+
+ // Round by just enough bits that the result is
+ // guaranteed to fit into an i16. Then the next stage can use 16 x 16 -> 32
+ // bit multiplies, which should be a fair bit faster than 32 x 32 -> 32
+ // as it does now
+ // This means shifting down so we have 6 extra bits, for a maximum value
+ // of +18360, which can occur if u_frac == 0.5 and the input pixels are
+ // {0, 255, 255, 0}.
+ __m256i out1 = _mm256_srai_epi32(_mm256_add_epi32(sum1, round_const_h),
+ DISFLOW_INTERP_BITS - 6);
+
+ _mm256_storeu_si256((__m256i *)tmp_row, _mm256_packs_epi32(out0, out1));
+ }
+
+ // Vertical convolution
+ const int round_bits = DISFLOW_INTERP_BITS + 6 - DISFLOW_DERIV_SCALE_LOG2;
+ __m256i round_const_v = _mm256_set1_epi32(1 << (round_bits - 1));
+
+ __m256i v_kernel_01 = _mm256_broadcastd_epi32(_mm_srli_si128(kernels, 4));
+ __m256i v_kernel_23 = _mm256_broadcastd_epi32(_mm_srli_si128(kernels, 12));
+
+ for (int i = 0; i < DISFLOW_PATCH_SIZE; i += 2) {
+ int16_t *tmp_row = &tmp[i * DISFLOW_PATCH_SIZE];
+
+ // Load 5 rows of 8 x 16-bit values, and pack into 4 registers
+ // holding rows {0, 1}, {1, 2}, {2, 3}, {3, 4}
+ __m128i row0 = _mm_loadu_si128((__m128i *)(tmp_row - DISFLOW_PATCH_SIZE));
+ __m128i row1 = _mm_loadu_si128((__m128i *)tmp_row);
+ __m128i row2 = _mm_loadu_si128((__m128i *)(tmp_row + DISFLOW_PATCH_SIZE));
+ __m128i row3 =
+ _mm_loadu_si128((__m128i *)(tmp_row + 2 * DISFLOW_PATCH_SIZE));
+ __m128i row4 =
+ _mm_loadu_si128((__m128i *)(tmp_row + 3 * DISFLOW_PATCH_SIZE));
+
+ __m256i px0 = _mm256_set_m128i(row1, row0);
+ __m256i px1 = _mm256_set_m128i(row2, row1);
+ __m256i px2 = _mm256_set_m128i(row3, row2);
+ __m256i px3 = _mm256_set_m128i(row4, row3);
+
+ // We want to calculate px0 * v_kernel[0] + px1 * v_kernel[1] + ... ,
+ // but each multiply expands its output to 32 bits. So we need to be
+ // a little clever about how we do this
+ __m256i sum0 = _mm256_add_epi32(
+ _mm256_madd_epi16(_mm256_unpacklo_epi16(px0, px1), v_kernel_01),
+ _mm256_madd_epi16(_mm256_unpacklo_epi16(px2, px3), v_kernel_23));
+ __m256i sum1 = _mm256_add_epi32(
+ _mm256_madd_epi16(_mm256_unpackhi_epi16(px0, px1), v_kernel_01),
+ _mm256_madd_epi16(_mm256_unpackhi_epi16(px2, px3), v_kernel_23));
+
+ __m256i sum0_rounded =
+ _mm256_srai_epi32(_mm256_add_epi32(sum0, round_const_v), round_bits);
+ __m256i sum1_rounded =
+ _mm256_srai_epi32(_mm256_add_epi32(sum1, round_const_v), round_bits);
+
+ __m256i warped = _mm256_packs_epi32(sum0_rounded, sum1_rounded);
+ __m128i src_pixels_u8 = xx_loadu_2x64(&src[(y + i + 1) * stride + x],
+ &src[(y + i) * stride + x]);
+ __m256i src_pixels =
+ _mm256_slli_epi16(_mm256_cvtepu8_epi16(src_pixels_u8), 3);
+
+ // Calculate delta from the target patch
+ __m256i dt = _mm256_sub_epi16(warped, src_pixels);
+
+ // Load 2x8 elements each of dx and dt, to pair with the 2x8 elements of dt
+ // that we have just computed. Then compute 2x8 partial sums of dx * dt
+ // and dy * dt, implicitly sum to give 2x4 partial sums of each, and
+ // accumulate.
+ __m256i dx_row = _mm256_loadu_si256((__m256i *)&dx[i * DISFLOW_PATCH_SIZE]);
+ __m256i dy_row = _mm256_loadu_si256((__m256i *)&dy[i * DISFLOW_PATCH_SIZE]);
+ b0_acc = _mm256_add_epi32(b0_acc, _mm256_madd_epi16(dx_row, dt));
+ b1_acc = _mm256_add_epi32(b1_acc, _mm256_madd_epi16(dy_row, dt));
+ }
+
+ // Flatten the two sets of partial sums to find the final value of b
+ // We need to set b[0] = sum(b0_acc), b[1] = sum(b1_acc).
+ // We need to do 14 additions in total; a `hadd` instruction can take care
+ // of eight of them, then a vertical sum can do four more, leaving two
+ // scalar additions.
+ __m256i partial_sum_256 = _mm256_hadd_epi32(b0_acc, b1_acc);
+ __m128i partial_sum =
+ _mm_add_epi32(_mm256_extracti128_si256(partial_sum_256, 0),
+ _mm256_extracti128_si256(partial_sum_256, 1));
+ b[0] = _mm_extract_epi32(partial_sum, 0) + _mm_extract_epi32(partial_sum, 1);
+ b[1] = _mm_extract_epi32(partial_sum, 2) + _mm_extract_epi32(partial_sum, 3);
+}
+
+// Compute the x and y gradients of the source patch in a single pass,
+// and store into dx and dy respectively.
+static INLINE void sobel_filter(const uint8_t *src, int src_stride, int16_t *dx,
+ int16_t *dy) {
+ const __m256i zero = _mm256_setzero_si256();
+
+ // Loop setup: Load the first two rows (of 10 input rows) and apply
+ // the horizontal parts of the two filters
+ __m256i row_m1_0 = _mm256_loadu2_m128i((__m128i *)(src - 1),
+ (__m128i *)(src - src_stride - 1));
+ __m256i row_m1_0_a = _mm256_unpacklo_epi8(row_m1_0, zero);
+ __m256i row_m1_0_b =
+ _mm256_unpacklo_epi8(_mm256_srli_si256(row_m1_0, 1), zero);
+ __m256i row_m1_0_c =
+ _mm256_unpacklo_epi8(_mm256_srli_si256(row_m1_0, 2), zero);
+
+ __m256i row_m1_0_hsmooth =
+ _mm256_add_epi16(_mm256_add_epi16(row_m1_0_a, row_m1_0_c),
+ _mm256_slli_epi16(row_m1_0_b, 1));
+ __m256i row_m1_0_hdiff = _mm256_sub_epi16(row_m1_0_a, row_m1_0_c);
+
+ // Main loop: For each pair of output rows (i, i+1):
+ // * Load rows (i+1, i+2) and apply both horizontal filters
+ // * Apply vertical filters and store results
+ // * Shift rows for next iteration
+ for (int i = 0; i < DISFLOW_PATCH_SIZE; i += 2) {
+ // Load rows (i+1, i+2) and apply both horizontal filters
+ const __m256i row_p1_p2 =
+ _mm256_loadu2_m128i((__m128i *)(src + (i + 2) * src_stride - 1),
+ (__m128i *)(src + (i + 1) * src_stride - 1));
+ const __m256i row_p1_p2_a = _mm256_unpacklo_epi8(row_p1_p2, zero);
+ const __m256i row_p1_p2_b =
+ _mm256_unpacklo_epi8(_mm256_srli_si256(row_p1_p2, 1), zero);
+ const __m256i row_p1_p2_c =
+ _mm256_unpacklo_epi8(_mm256_srli_si256(row_p1_p2, 2), zero);
+
+ const __m256i row_p1_p2_hsmooth =
+ _mm256_add_epi16(_mm256_add_epi16(row_p1_p2_a, row_p1_p2_c),
+ _mm256_slli_epi16(row_p1_p2_b, 1));
+ const __m256i row_p1_p2_hdiff = _mm256_sub_epi16(row_p1_p2_a, row_p1_p2_c);
+
+ // Apply vertical filters and store results
+ // dx = vertical smooth(horizontal diff(input))
+ // dy = vertical diff(horizontal smooth(input))
+ const __m256i row_0_p1_hdiff =
+ _mm256_permute2x128_si256(row_m1_0_hdiff, row_p1_p2_hdiff, 0x21);
+ const __m256i dx_row =
+ _mm256_add_epi16(_mm256_add_epi16(row_m1_0_hdiff, row_p1_p2_hdiff),
+ _mm256_slli_epi16(row_0_p1_hdiff, 1));
+ const __m256i dy_row =
+ _mm256_sub_epi16(row_m1_0_hsmooth, row_p1_p2_hsmooth);
+
+ _mm256_storeu_si256((__m256i *)(dx + i * DISFLOW_PATCH_SIZE), dx_row);
+ _mm256_storeu_si256((__m256i *)(dy + i * DISFLOW_PATCH_SIZE), dy_row);
+
+ // Shift rows for next iteration
+ // This allows a lot of work to be reused, reducing the number of
+ // horizontal filtering operations from 2*3*8 = 48 to 2*10 = 20
+ row_m1_0_hsmooth = row_p1_p2_hsmooth;
+ row_m1_0_hdiff = row_p1_p2_hdiff;
+ }
+}
+
+static INLINE void compute_flow_matrix(const int16_t *dx, int dx_stride,
+ const int16_t *dy, int dy_stride,
+ double *M) {
+ __m256i acc[4] = { 0 };
+
+ for (int i = 0; i < DISFLOW_PATCH_SIZE; i += 2) {
+ __m256i dx_row = _mm256_loadu_si256((__m256i *)&dx[i * dx_stride]);
+ __m256i dy_row = _mm256_loadu_si256((__m256i *)&dy[i * dy_stride]);
+
+ acc[0] = _mm256_add_epi32(acc[0], _mm256_madd_epi16(dx_row, dx_row));
+ acc[1] = _mm256_add_epi32(acc[1], _mm256_madd_epi16(dx_row, dy_row));
+ // Don't compute acc[2], as it should be equal to acc[1]
+ acc[3] = _mm256_add_epi32(acc[3], _mm256_madd_epi16(dy_row, dy_row));
+ }
+
+ // Condense sums
+ __m256i partial_sum_0 = _mm256_hadd_epi32(acc[0], acc[1]);
+ __m256i partial_sum_1 = _mm256_hadd_epi32(acc[1], acc[3]);
+ __m256i result_256 = _mm256_hadd_epi32(partial_sum_0, partial_sum_1);
+ __m128i result = _mm_add_epi32(_mm256_extracti128_si256(result_256, 0),
+ _mm256_extracti128_si256(result_256, 1));
+
+ // Apply regularization
+ // We follow the standard regularization method of adding `k * I` before
+ // inverting. This ensures that the matrix will be invertible.
+ //
+ // Setting the regularization strength k to 1 seems to work well here, as
+ // typical values coming from the other equations are very large (1e5 to
+ // 1e6, with an upper limit of around 6e7, at the time of writing).
+ // It also preserves the property that all matrix values are whole numbers,
+ // which is convenient for integerized SIMD implementation.
+ result = _mm_add_epi32(result, _mm_set_epi32(1, 0, 0, 1));
+
+ // Convert results to doubles and store
+ _mm256_storeu_pd(M, _mm256_cvtepi32_pd(result));
+}
+
+// Try to invert the matrix M
+// Note: Due to the nature of how a least-squares matrix is constructed, all of
+// the eigenvalues will be >= 0, and therefore det M >= 0 as well.
+// The regularization term `+ k * I` further ensures that det M >= k^2.
+// As mentioned in compute_flow_matrix(), here we use k = 1, so det M >= 1.
+// So we don't have to worry about non-invertible matrices here.
+static INLINE void invert_2x2(const double *M, double *M_inv) {
+ double det = (M[0] * M[3]) - (M[1] * M[2]);
+ assert(det >= 1);
+ const double det_inv = 1 / det;
+
+ M_inv[0] = M[3] * det_inv;
+ M_inv[1] = -M[1] * det_inv;
+ M_inv[2] = -M[2] * det_inv;
+ M_inv[3] = M[0] * det_inv;
+}
+
+void aom_compute_flow_at_point_avx2(const uint8_t *src, const uint8_t *ref,
+ int x, int y, int width, int height,
+ int stride, double *u, double *v) {
+ DECLARE_ALIGNED(32, double, M[4]);
+ DECLARE_ALIGNED(32, double, M_inv[4]);
+ DECLARE_ALIGNED(32, int16_t, dx[DISFLOW_PATCH_SIZE * DISFLOW_PATCH_SIZE]);
+ DECLARE_ALIGNED(32, int16_t, dy[DISFLOW_PATCH_SIZE * DISFLOW_PATCH_SIZE]);
+ int b[2];
+
+ // Compute gradients within this patch
+ const uint8_t *src_patch = &src[y * stride + x];
+ sobel_filter(src_patch, stride, dx, dy);
+
+ compute_flow_matrix(dx, DISFLOW_PATCH_SIZE, dy, DISFLOW_PATCH_SIZE, M);
+ invert_2x2(M, M_inv);
+
+ for (int itr = 0; itr < DISFLOW_MAX_ITR; itr++) {
+ compute_flow_vector(src, ref, width, height, stride, x, y, *u, *v, dx, dy,
+ b);
+
+ // Solve flow equations to find a better estimate for the flow vector
+ // at this point
+ const double step_u = M_inv[0] * b[0] + M_inv[1] * b[1];
+ const double step_v = M_inv[2] * b[0] + M_inv[3] * b[1];
+ *u += fclamp(step_u * DISFLOW_STEP_SIZE, -2, 2);
+ *v += fclamp(step_v * DISFLOW_STEP_SIZE, -2, 2);
+
+ if (fabs(step_u) + fabs(step_v) < DISFLOW_STEP_SIZE_THRESOLD) {
+ // Stop iteration when we're close to convergence
+ break;
+ }
+ }
+}
diff --git a/aom_dsp/x86/synonyms.h b/aom_dsp/x86/synonyms.h
index 6744ec51d0..0d51cdff48 100644
--- a/aom_dsp/x86/synonyms.h
+++ b/aom_dsp/x86/synonyms.h
@@ -46,6 +46,13 @@ static INLINE __m128i xx_loadu_128(const void *a) {
return _mm_loadu_si128((const __m128i *)a);
}
+// Load 64 bits from each of hi and low, and pack into an SSE register
+// Since directly loading as `int64_t`s and using _mm_set_epi64 may violate
+// the strict aliasing rule, this takes a different approach
+static INLINE __m128i xx_loadu_2x64(const void *hi, const void *lo) {
+ return _mm_unpacklo_epi64(_mm_loadu_si64(lo), _mm_loadu_si64(hi));
+}
+
static INLINE void xx_storel_32(void *const a, const __m128i v) {
const int val = _mm_cvtsi128_si32(v);
memcpy(a, &val, sizeof(val));
diff --git a/aom_dsp/x86/synonyms_avx2.h b/aom_dsp/x86/synonyms_avx2.h
index b729e5f410..d4e8f69111 100644
--- a/aom_dsp/x86/synonyms_avx2.h
+++ b/aom_dsp/x86/synonyms_avx2.h
@@ -43,6 +43,16 @@ static INLINE void yy_storeu_256(void *const a, const __m256i v) {
_mm256_storeu_si256((__m256i *)a, v);
}
+// Fill an AVX register using an interleaved pair of values, ie. set the
+// 16 channels to {a, b} repeated 8 times, using the same channel ordering
+// as when a register is stored to / loaded from memory.
+//
+// This is useful for rearranging filter kernels for use with the _mm_madd_epi16
+// instruction
+static INLINE __m256i yy_set2_epi16(int16_t a, int16_t b) {
+ return _mm256_setr_epi16(a, b, a, b, a, b, a, b, a, b, a, b, a, b, a, b);
+}
+
// The _mm256_set1_epi64x() intrinsic is undefined for some Visual Studio
// compilers. The following function is equivalent to _mm256_set1_epi64x()
// acting on a 32-bit integer.
diff --git a/test/disflow_test.cc b/test/disflow_test.cc
index 124c9a96c7..4f004480e2 100644
--- a/test/disflow_test.cc
+++ b/test/disflow_test.cc
@@ -114,6 +114,11 @@ INSTANTIATE_TEST_SUITE_P(SSE4_1, ComputeFlowTest,
::testing::Values(aom_compute_flow_at_point_sse4_1));
#endif
+#if HAVE_AVX2
+INSTANTIATE_TEST_SUITE_P(AVX2, ComputeFlowTest,
+ ::testing::Values(aom_compute_flow_at_point_avx2));
+#endif
+
#if HAVE_NEON
INSTANTIATE_TEST_SUITE_P(NEON, ComputeFlowTest,
::testing::Values(aom_compute_flow_at_point_neon));