Google Git
Sign in
aomedia / aom / HEAD / . / av1 / common / selfguided_hwy.h
blob: 33433e416774940741204cd249adf3b1c07b45c9 [file] [log] [blame]
/*
* Copyright (c) 2025, 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.
*/
#ifndef AV1_COMMON_SELFGUIDED_HWY_H_
#define AV1_COMMON_SELFGUIDED_HWY_H_
#include "av1/common/restoration.h"
#include "config/aom_config.h"
#include "config/av1_rtcd.h"
#include "third_party/highway/hwy/highway.h"
HWY_BEFORE_NAMESPACE();
namespace {
namespace HWY_NAMESPACE {
namespace hn = hwy::HWY_NAMESPACE;
template <int NumBlocks>
struct ScanTraits {};
template <>
struct ScanTraits<1> {
template <typename D>
HWY_ATTR HWY_INLINE static hn::VFromD<D> AddBlocks(D int32_tag,
hn::VFromD<D> v) {
(void)int32_tag;
return v;
}
};
template <>
struct ScanTraits<2> {
template <typename D>
HWY_ATTR HWY_INLINE static hn::VFromD<D> AddBlocks(D int32_tag,
hn::VFromD<D> v) {
constexpr hn::Half<D> half_tag;
const int32_t s = hn::ExtractLane(v, 3);
const auto s01 = hn::Set(half_tag, s);
const auto s02 = hn::InsertBlock<1>(hn::Zero(int32_tag), s01);
return hn::Add(v, s02);
}
};
template <>
struct ScanTraits<4> {
template <typename D>
HWY_ATTR HWY_INLINE static hn::VFromD<D> AddBlocks(D int32_tag,
hn::VFromD<D> v) {
HWY_ALIGN static const int32_t kA[] = {
0, 0, 0, 0, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19, 19,
};
HWY_ALIGN static const int32_t kB[] = {
0, 0, 0, 0, 0, 0, 0, 0, 23, 23, 23, 23, 23, 23, 23, 23,
};
HWY_ALIGN static const int32_t kC[] = {
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 27, 27, 27, 27,
};
const auto a = hn::SetTableIndices(int32_tag, kA);
const auto b = hn::SetTableIndices(int32_tag, kB);
const auto c = hn::SetTableIndices(int32_tag, kC);
const auto s01 =
hn::TwoTablesLookupLanes(int32_tag, hn::Zero(int32_tag), v, a);
const auto s02 =
hn::TwoTablesLookupLanes(int32_tag, hn::Zero(int32_tag), v, b);
const auto s03 =
hn::TwoTablesLookupLanes(int32_tag, hn::Zero(int32_tag), v, c);
v = hn::Add(v, s01);
v = hn::Add(v, s02);
v = hn::Add(v, s03);
return v;
}
};
// Compute the scan of a register holding 32-bit integers. If the register holds
// x0..x7 then the scan will hold x0, x0+x1, x0+x1+x2, ..., x0+x1+...+x7
//
// For the AVX2 example below, let [...] represent a 128-bit block, and let a,
// ..., h be 32-bit integers (assumed small enough to be able to add them
// without overflow).
//
// Use -> as shorthand for summing, i.e. h->a = h + g + f + e + d + c + b + a.
//
// x = [h g f e][d c b a]
// x01 = [g f e 0][c b a 0]
// x02 = [g+h f+g e+f e][c+d b+c a+b a]
// x03 = [e+f e 0 0][a+b a 0 0]
// x04 = [e->h e->g e->f e][a->d a->c a->b a]
// s = a->d
// s01 = [a->d a->d a->d a->d]
// s02 = [a->d a->d a->d a->d][0 0 0 0]
// ret = [a->h a->g a->f a->e][a->d a->c a->b a]
template <typename D>
HWY_ATTR HWY_INLINE hn::VFromD<D> Scan32(D int32_tag, hn::VFromD<D> x) {
const auto x01 = hn::ShiftLeftBytes<4>(x);
const auto x02 = hn::Add(x, x01);
const auto x03 = hn::ShiftLeftBytes<8>(x02);
const auto x04 = hn::Add(x02, x03);
return ScanTraits<int32_tag.MaxBlocks()>::AddBlocks(int32_tag, x04);
}
// Compute two integral images from src. B sums elements; A sums their
// squares. The images are offset by one pixel, so will have width and height
// equal to width + 1, height + 1 and the first row and column will be zero.
//
// A+1 and B+1 should be aligned to 32 bytes. buf_stride should be a multiple
// of 8.
template <typename T, typename D>
HWY_ATTR HWY_INLINE void IntegralImages(D int32_tag, const T *HWY_RESTRICT src,
int src_stride, int width, int height,
int32_t *HWY_RESTRICT A,
int32_t *HWY_RESTRICT B,
int buf_stride) {
constexpr hn::Rebind<T, D> uint_tag;
constexpr hn::Repartition<int16_t, D> int16_tag;
// Write out the zero top row
hwy::ZeroBytes(A, 4 * (width + 8));
hwy::ZeroBytes(B, 4 * (width + 8));
for (int i = 0; i < height; ++i) {
// Zero the left column.
A[(i + 1) * buf_stride] = B[(i + 1) * buf_stride] = 0;
// ldiff is the difference H - D where H is the output sample immediately
// to the left and D is the output sample above it. These are scalars,
// replicated across the eight lanes.
auto ldiff1 = hn::Zero(int32_tag);
auto ldiff2 = hn::Zero(int32_tag);
for (int j = 0; j < width; j += hn::MaxLanes(int32_tag)) {
const int ABj = 1 + j;
const auto above1 = hn::Load(int32_tag, B + ABj + i * buf_stride);
const auto above2 = hn::Load(int32_tag, A + ABj + i * buf_stride);
const auto x1 = hn::PromoteTo(
int32_tag, hn::LoadU(uint_tag, src + j + i * src_stride));
const auto x2 = hn::WidenMulPairwiseAdd(
int32_tag, hn::BitCast(int16_tag, x1), hn::BitCast(int16_tag, x1));
const auto sc1 = Scan32(int32_tag, x1);
const auto sc2 = Scan32(int32_tag, x2);
const auto row1 = hn::Add(hn::Add(sc1, above1), ldiff1);
const auto row2 = hn::Add(hn::Add(sc2, above2), ldiff2);
hn::Store(row1, int32_tag, B + ABj + (i + 1) * buf_stride);
hn::Store(row2, int32_tag, A + ABj + (i + 1) * buf_stride);
// Calculate the new H - D.
ldiff1 = hn::Set(int32_tag, hn::ExtractLane(hn::Sub(row1, above1),
hn::MaxLanes(int32_tag) - 1));
ldiff2 = hn::Set(int32_tag, hn::ExtractLane(hn::Sub(row2, above2),
hn::MaxLanes(int32_tag) - 1));
}
}
}
template <typename D>
HWY_ATTR HWY_INLINE hn::VFromD<D> BoxSumFromII(D int32_tag,
const int32_t *HWY_RESTRICT ii,
int stride, int r) {
const auto tl = hn::LoadU(int32_tag, ii - (r + 1) - (r + 1) * stride);
const auto tr = hn::LoadU(int32_tag, ii + (r + 0) - (r + 1) * stride);
const auto bl = hn::LoadU(int32_tag, ii - (r + 1) + r * stride);
const auto br = hn::LoadU(int32_tag, ii + (r + 0) + r * stride);
const auto u = hn::Sub(tr, tl);
const auto v = hn::Sub(br, bl);
return hn::Sub(v, u);
}
template <typename D>
HWY_ATTR HWY_INLINE hn::VFromD<D> RoundForShift(D int32_tag,
unsigned int shift) {
return hn::Set(int32_tag, (1 << shift) >> 1);
}
template <typename D>
HWY_ATTR HWY_INLINE hn::VFromD<D> ComputeP(D int32_tag, hn::VFromD<D> sum1,
hn::VFromD<D> sum2, int bit_depth,
int n) {
constexpr hn::Repartition<int16_t, D> int16_tag;
if (bit_depth > 8) {
const auto rounding_a = RoundForShift(int32_tag, 2 * (bit_depth - 8));
const auto rounding_b = RoundForShift(int32_tag, bit_depth - 8);
const auto a =
hn::ShiftRightSame(hn::Add(sum2, rounding_a), 2 * (bit_depth - 8));
const auto b = hn::ShiftRightSame(hn::Add(sum1, rounding_b), bit_depth - 8);
// b < 2^14, so we can use a 16-bit madd rather than a 32-bit
// mullo to square it
const auto b_16 = hn::BitCast(int16_tag, b);
const auto bb = hn::WidenMulPairwiseAdd(int32_tag, b_16, b_16);
const auto an = hn::Max(hn::Mul(a, hn::Set(int32_tag, n)), bb);
return hn::Sub(an, bb);
}
const auto sum1_16 = hn::BitCast(int16_tag, sum1);
const auto bb = hn::WidenMulPairwiseAdd(int32_tag, sum1_16, sum1_16);
const auto an = hn::Mul(sum2, hn::Set(int32_tag, n));
return hn::Sub(an, bb);
}
// Calculate 8 values of the "cross sum" starting at buf. This is a 3x3 filter
// where the outer four corners have weight 3 and all other pixels have weight
// 4.
//
// Pixels are indexed as follows:
// xtl xt xtr
// xl x xr
// xbl xb xbr
//
// buf points to x
//
// fours = xl + xt + xr + xb + x
// threes = xtl + xtr + xbr + xbl
// cross_sum = 4 * fours + 3 * threes
// = 4 * (fours + threes) - threes
// = (fours + threes) << 2 - threes
template <typename D>
HWY_ATTR HWY_INLINE hn::VFromD<D> CrossSum(D int32_tag,
const int32_t *HWY_RESTRICT buf,
int stride) {
const auto xtl = hn::LoadU(int32_tag, buf - 1 - stride);
const auto xt = hn::LoadU(int32_tag, buf - stride);
const auto xtr = hn::LoadU(int32_tag, buf + 1 - stride);
const auto xl = hn::LoadU(int32_tag, buf - 1);
const auto x = hn::LoadU(int32_tag, buf);
const auto xr = hn::LoadU(int32_tag, buf + 1);
const auto xbl = hn::LoadU(int32_tag, buf - 1 + stride);
const auto xb = hn::LoadU(int32_tag, buf + stride);
const auto xbr = hn::LoadU(int32_tag, buf + 1 + stride);
const auto fours = hn::Add(xl, hn::Add(xt, hn::Add(xr, hn::Add(xb, x))));
const auto threes = hn::Add(xtl, hn::Add(xtr, hn::Add(xbr, xbl)));
return hn::Sub(hn::ShiftLeft<2>(hn::Add(fours, threes)), threes);
}
// The final filter for self-guided restoration. Computes a weighted average
// across A, B with "cross sums" (see CrossSum implementation above).
template <typename DL>
HWY_ATTR HWY_INLINE void FinalFilter(
DL int32_tag, int32_t *HWY_RESTRICT dst, int dst_stride,
const int32_t *HWY_RESTRICT A, const int32_t *HWY_RESTRICT B,
int buf_stride, const void *HWY_RESTRICT dgd8, int dgd_stride, int width,
int height, int highbd) {
constexpr hn::Repartition<uint8_t, hn::Half<DL>> uint8_half_tag;
constexpr hn::Repartition<int16_t, DL> int16_tag;
constexpr int nb = 5;
constexpr int kShift = SGRPROJ_SGR_BITS + nb - SGRPROJ_RST_BITS;
const auto rounding = RoundForShift(int32_tag, kShift);
const uint8_t *HWY_RESTRICT dgd_real =
highbd ? reinterpret_cast<const uint8_t *>(CONVERT_TO_SHORTPTR(dgd8))
: reinterpret_cast<const uint8_t *>(dgd8);
for (int i = 0; i < height; ++i) {
for (int j = 0; j < width; j += hn::MaxLanes(int32_tag)) {
const auto a = CrossSum(int32_tag, A + i * buf_stride + j, buf_stride);
const auto b = CrossSum(int32_tag, B + i * buf_stride + j, buf_stride);
const auto raw = hn::LoadU(uint8_half_tag,
dgd_real + ((i * dgd_stride + j) << highbd));
const auto src =
highbd ? hn::PromoteTo(
int32_tag,
hn::BitCast(
hn::Repartition<int16_t, decltype(uint8_half_tag)>(),
raw))
: hn::PromoteTo(int32_tag, hn::LowerHalf(raw));
auto v =
hn::Add(hn::WidenMulPairwiseAdd(int32_tag, hn::BitCast(int16_tag, a),
hn::BitCast(int16_tag, src)),
b);
auto w = hn::ShiftRight<kShift>(hn::Add(v, rounding));
hn::StoreU(w, int32_tag, dst + i * dst_stride + j);
}
}
}
// Assumes that C, D are integral images for the original buffer which has been
// extended to have a padding of SGRPROJ_BORDER_VERT/SGRPROJ_BORDER_HORZ pixels
// on the sides. A, B, C, D point at logical position (0, 0).
template <int Step, typename DL>
HWY_ATTR HWY_INLINE void CalcAB(DL int32_tag, int32_t *HWY_RESTRICT A,
int32_t *HWY_RESTRICT B,
const int32_t *HWY_RESTRICT C,
const int32_t *HWY_RESTRICT D, int width,
int height, int buf_stride, int bit_depth,
int sgr_params_idx, int radius_idx) {
constexpr hn::Repartition<int16_t, DL> int16_tag;
constexpr hn::Repartition<uint32_t, DL> uint32_tag;
const sgr_params_type *HWY_RESTRICT const params =
&av1_sgr_params[sgr_params_idx];
const int r = params->r[radius_idx];
const int n = (2 * r + 1) * (2 * r + 1);
const auto s = hn::Set(int32_tag, params->s[radius_idx]);
// one_over_n[n-1] is 2^12/n, so easily fits in an int16
const auto one_over_n =
hn::BitCast(int16_tag, hn::Set(int32_tag, av1_one_by_x[n - 1]));
const auto rnd_z = RoundForShift(int32_tag, SGRPROJ_MTABLE_BITS);
const auto rnd_res = RoundForShift(int32_tag, SGRPROJ_RECIP_BITS);
// Set up masks
const int max_lanes = static_cast<int>(hn::MaxLanes(int32_tag));
HWY_ALIGN hn::Mask<decltype(int32_tag)> mask[max_lanes];
for (int idx = 0; idx < max_lanes; idx++) {
mask[idx] = hn::FirstN(int32_tag, idx);
}
for (int i = -1; i < height + 1; i += Step) {
for (int j = -1; j < width + 1; j += max_lanes) {
const int32_t *HWY_RESTRICT Cij = C + i * buf_stride + j;
const int32_t *HWY_RESTRICT Dij = D + i * buf_stride + j;
auto sum1 = BoxSumFromII(int32_tag, Dij, buf_stride, r);
auto sum2 = BoxSumFromII(int32_tag, Cij, buf_stride, r);
// When width + 2 isn't a multiple of 8, sum1 and sum2 will contain
// some uninitialised data in their upper words. We use a mask to
// ensure that these bits are set to 0.
int idx = AOMMIN(max_lanes, width + 1 - j);
assert(idx >= 1);
if (idx < max_lanes) {
sum1 = hn::IfThenElseZero(mask[idx], sum1);
sum2 = hn::IfThenElseZero(mask[idx], sum2);
}
const auto p = ComputeP(int32_tag, sum1, sum2, bit_depth, n);
const auto z = hn::BitCast(
int32_tag, hn::Min(hn::ShiftRight<SGRPROJ_MTABLE_BITS>(hn::BitCast(
uint32_tag, hn::MulAdd(p, s, rnd_z))),
hn::Set(uint32_tag, 255)));
const auto a_res = hn::GatherIndex(int32_tag, av1_x_by_xplus1, z);
hn::StoreU(a_res, int32_tag, A + i * buf_stride + j);
const auto a_complement = hn::Sub(hn::Set(int32_tag, SGRPROJ_SGR), a_res);
// sum1 might have lanes greater than 2^15, so we can't use madd to do
// multiplication involving sum1. However, a_complement and one_over_n
// are both less than 256, so we can multiply them first.
const auto a_comp_over_n = hn::WidenMulPairwiseAdd(
int32_tag, hn::BitCast(int16_tag, a_complement), one_over_n);
const auto b_int = hn::Mul(a_comp_over_n, sum1);
const auto b_res =
hn::ShiftRight<SGRPROJ_RECIP_BITS>(hn::Add(b_int, rnd_res));
hn::StoreU(b_res, int32_tag, B + i * buf_stride + j);
}
}
}
// Calculate 8 values of the "cross sum" starting at buf.
//
// Pixels are indexed like this:
// xtl xt xtr
// - buf -
// xbl xb xbr
//
// Pixels are weighted like this:
// 5 6 5
// 0 0 0
// 5 6 5
//
// fives = xtl + xtr + xbl + xbr
// sixes = xt + xb
// cross_sum = 6 * sixes + 5 * fives
// = 5 * (fives + sixes) - sixes
// = (fives + sixes) << 2 + (fives + sixes) + sixes
template <typename D>
HWY_ATTR HWY_INLINE hn::VFromD<D> CrossSumFastEvenRow(
D int32_tag, const int32_t *HWY_RESTRICT buf, int stride) {
const auto xtl = hn::LoadU(int32_tag, buf - 1 - stride);
const auto xt = hn::LoadU(int32_tag, buf - stride);
const auto xtr = hn::LoadU(int32_tag, buf + 1 - stride);
const auto xbl = hn::LoadU(int32_tag, buf - 1 + stride);
const auto xb = hn::LoadU(int32_tag, buf + stride);
const auto xbr = hn::LoadU(int32_tag, buf + 1 + stride);
const auto fives = hn::Add(xtl, hn::Add(xtr, hn::Add(xbr, xbl)));
const auto sixes = hn::Add(xt, xb);
const auto fives_plus_sixes = hn::Add(fives, sixes);
return hn::Add(hn::Add(hn::ShiftLeft<2>(fives_plus_sixes), fives_plus_sixes),
sixes);
}
// Calculate 8 values of the "cross sum" starting at buf.
//
// Pixels are indexed like this:
// xl x xr
//
// Pixels are weighted like this:
// 5 6 5
//
// buf points to x
//
// fives = xl + xr
// sixes = x
// cross_sum = 5 * fives + 6 * sixes
// = 4 * (fives + sixes) + (fives + sixes) + sixes
// = (fives + sixes) << 2 + (fives + sixes) + sixes
template <typename D>
HWY_ATTR HWY_INLINE hn::VFromD<D> CrossSumFastOddRow(
D int32_tag, const int32_t *HWY_RESTRICT buf) {
const auto xl = hn::LoadU(int32_tag, buf - 1);
const auto x = hn::LoadU(int32_tag, buf);
const auto xr = hn::LoadU(int32_tag, buf + 1);
const auto fives = hn::Add(xl, xr);
const auto sixes = x;
const auto fives_plus_sixes = hn::Add(fives, sixes);
return hn::Add(hn::Add(hn::ShiftLeft<2>(fives_plus_sixes), fives_plus_sixes),
sixes);
}
// The final filter for the self-guided restoration. Computes a
// weighted average across A, B with "cross sums" (see cross_sum_...
// implementations above).
template <typename DL>
HWY_ATTR HWY_INLINE void FinalFilterFast(
DL int32_tag, int32_t *HWY_RESTRICT dst, int dst_stride,
const int32_t *HWY_RESTRICT A, const int32_t *HWY_RESTRICT B,
int buf_stride, const void *HWY_RESTRICT dgd8, int dgd_stride, int width,
int height, int highbd) {
constexpr hn::Repartition<uint8_t, hn::Half<DL>> uint8_half_tag;
constexpr hn::Repartition<int16_t, DL> int16_tag;
constexpr int nb0 = 5;
constexpr int nb1 = 4;
constexpr int kShift0 = SGRPROJ_SGR_BITS + nb0 - SGRPROJ_RST_BITS;
constexpr int kShift1 = SGRPROJ_SGR_BITS + nb1 - SGRPROJ_RST_BITS;
const auto rounding0 = RoundForShift(int32_tag, kShift0);
const auto rounding1 = RoundForShift(int32_tag, kShift1);
const uint8_t *HWY_RESTRICT dgd_real =
highbd ? reinterpret_cast<const uint8_t *>(CONVERT_TO_SHORTPTR(dgd8))
: reinterpret_cast<const uint8_t *>(dgd8);
for (int i = 0; i < height; ++i) {
if (!(i & 1)) { // even row
for (int j = 0; j < width; j += hn::MaxLanes(int32_tag)) {
const auto a =
CrossSumFastEvenRow(int32_tag, A + i * buf_stride + j, buf_stride);
const auto b =
CrossSumFastEvenRow(int32_tag, B + i * buf_stride + j, buf_stride);
const auto raw = hn::LoadU(uint8_half_tag,
dgd_real + ((i * dgd_stride + j) << highbd));
const auto src =
highbd
? hn::PromoteTo(
int32_tag,
hn::BitCast(
hn::Repartition<int16_t, decltype(uint8_half_tag)>(),
raw))
: hn::PromoteTo(int32_tag, hn::LowerHalf(raw));
auto v = hn::Add(
hn::WidenMulPairwiseAdd(int32_tag, hn::BitCast(int16_tag, a),
hn::BitCast(int16_tag, src)),
b);
auto w = hn::ShiftRight<kShift0>(hn::Add(v, rounding0));
hn::StoreU(w, int32_tag, dst + i * dst_stride + j);
}
} else { // odd row
for (int j = 0; j < width; j += hn::MaxLanes(int32_tag)) {
const auto a = CrossSumFastOddRow(int32_tag, A + i * buf_stride + j);
const auto b = CrossSumFastOddRow(int32_tag, B + i * buf_stride + j);
const auto raw = hn::LoadU(uint8_half_tag,
dgd_real + ((i * dgd_stride + j) << highbd));
const auto src =
highbd
? hn::PromoteTo(
int32_tag,
hn::BitCast(
hn::Repartition<int16_t, decltype(uint8_half_tag)>(),
raw))
: hn::PromoteTo(int32_tag, hn::LowerHalf(raw));
auto v = hn::Add(
hn::WidenMulPairwiseAdd(int32_tag, hn::BitCast(int16_tag, a),
hn::BitCast(int16_tag, src)),
b);
auto w = hn::ShiftRight<kShift1>(hn::Add(v, rounding1));
hn::StoreU(w, int32_tag, dst + i * dst_stride + j);
}
}
}
}
HWY_ATTR HWY_INLINE int SelfGuidedRestoration(
const uint8_t *dgd8, int width, int height, int dgd_stride,
int32_t *HWY_RESTRICT flt0, int32_t *HWY_RESTRICT flt1, int flt_stride,
int sgr_params_idx, int bit_depth, int highbd) {
constexpr hn::ScalableTag<int32_t> int32_tag;
constexpr int kAlignment32Log2 = hwy::CeilLog2(hn::MaxLanes(int32_tag));
// The ALIGN_POWER_OF_TWO macro here ensures that column 1 of Atl, Btl, Ctl
// and Dtl is vector aligned.
const int buf_elts =
ALIGN_POWER_OF_TWO(RESTORATION_PROC_UNIT_PELS, kAlignment32Log2);
int32_t *buf = reinterpret_cast<int32_t *>(
aom_memalign(4 << kAlignment32Log2, 4 * sizeof(*buf) * buf_elts));
if (!buf) {
return -1;
}
const int width_ext = width + 2 * SGRPROJ_BORDER_HORZ;
const int height_ext = height + 2 * SGRPROJ_BORDER_VERT;
// Adjusting the stride of A and B here appears to avoid bad cache effects,
// leading to a significant speed improvement.
// We also align the stride to a multiple of the vector size for efficiency.
int buf_stride =
ALIGN_POWER_OF_TWO(width_ext + (2 << kAlignment32Log2), kAlignment32Log2);
// The "tl" pointers point at the top-left of the initialised data for the
// array.
int32_t *Atl = buf + 0 * buf_elts + (1 << kAlignment32Log2) - 1;
int32_t *Btl = buf + 1 * buf_elts + (1 << kAlignment32Log2) - 1;
int32_t *Ctl = buf + 2 * buf_elts + (1 << kAlignment32Log2) - 1;
int32_t *Dtl = buf + 3 * buf_elts + (1 << kAlignment32Log2) - 1;
// The "0" pointers are (- SGRPROJ_BORDER_VERT, -SGRPROJ_BORDER_HORZ). Note
// there's a zero row and column in A, B (integral images), so we move down
// and right one for them.
const int buf_diag_border =
SGRPROJ_BORDER_HORZ + buf_stride * SGRPROJ_BORDER_VERT;
int32_t *A0 = Atl + 1 + buf_stride;
int32_t *B0 = Btl + 1 + buf_stride;
int32_t *C0 = Ctl + 1 + buf_stride;
int32_t *D0 = Dtl + 1 + buf_stride;
// Finally, A, B, C, D point at position (0, 0).
int32_t *A = A0 + buf_diag_border;
int32_t *B = B0 + buf_diag_border;
int32_t *C = C0 + buf_diag_border;
int32_t *D = D0 + buf_diag_border;
const int dgd_diag_border =
SGRPROJ_BORDER_HORZ + dgd_stride * SGRPROJ_BORDER_VERT;
const uint8_t *dgd0 = dgd8 - dgd_diag_border;
// Generate integral images from the input. C will contain sums of squares; D
// will contain just sums
if (highbd) {
IntegralImages(int32_tag, CONVERT_TO_SHORTPTR(dgd0), dgd_stride, width_ext,
height_ext, Ctl, Dtl, buf_stride);
} else {
IntegralImages(int32_tag, dgd0, dgd_stride, width_ext, height_ext, Ctl, Dtl,
buf_stride);
}
const sgr_params_type *const params = &av1_sgr_params[sgr_params_idx];
// Write to flt0 and flt1
// If params->r == 0 we skip the corresponding filter. We only allow one of
// the radii to be 0, as having both equal to 0 would be equivalent to
// skipping SGR entirely.
assert(!(params->r[0] == 0 && params->r[1] == 0));
assert(params->r[0] < AOMMIN(SGRPROJ_BORDER_VERT, SGRPROJ_BORDER_HORZ));
assert(params->r[1] < AOMMIN(SGRPROJ_BORDER_VERT, SGRPROJ_BORDER_HORZ));
if (params->r[0] > 0) {
CalcAB<2>(int32_tag, A, B, C, D, width, height, buf_stride, bit_depth,
sgr_params_idx, 0);
FinalFilterFast(int32_tag, flt0, flt_stride, A, B, buf_stride, dgd8,
dgd_stride, width, height, highbd);
}
if (params->r[1] > 0) {
CalcAB<1>(int32_tag, A, B, C, D, width, height, buf_stride, bit_depth,
sgr_params_idx, 1);
FinalFilter(int32_tag, flt1, flt_stride, A, B, buf_stride, dgd8, dgd_stride,
width, height, highbd);
}
aom_free(buf);
return 0;
}
HWY_ATTR HWY_INLINE int ApplySelfGuidedRestoration(
const uint8_t *HWY_RESTRICT dat8, int width, int height, int stride,
int eps, const int *HWY_RESTRICT xqd, uint8_t *HWY_RESTRICT dst8,
int dst_stride, int32_t *HWY_RESTRICT tmpbuf, int bit_depth, int highbd) {
constexpr hn::CappedTag<int32_t, 16> int32_tag;
constexpr size_t kBatchSize = hn::MaxLanes(int32_tag) * 2;
int32_t *flt0 = tmpbuf;
int32_t *flt1 = flt0 + RESTORATION_UNITPELS_MAX;
assert(width * height <= RESTORATION_UNITPELS_MAX);
#if HWY_TARGET == HWY_SSE4
const int ret = av1_selfguided_restoration_sse4_1(
dat8, width, height, stride, flt0, flt1, width, eps, bit_depth, highbd);
#elif HWY_TARGET == HWY_AVX2
const int ret = av1_selfguided_restoration_avx2(
dat8, width, height, stride, flt0, flt1, width, eps, bit_depth, highbd);
#elif HWY_TARGET <= HWY_AVX3
const int ret = av1_selfguided_restoration_avx512(
dat8, width, height, stride, flt0, flt1, width, eps, bit_depth, highbd);
#else
#error "HWY_TARGET is not supported."
const int ret = -1;
#endif
if (ret != 0) {
return ret;
}
const sgr_params_type *const params = &av1_sgr_params[eps];
int xq[2];
av1_decode_xq(xqd, xq, params);
auto xq0 = hn::Set(int32_tag, xq[0]);
auto xq1 = hn::Set(int32_tag, xq[1]);
for (int i = 0; i < height; ++i) {
// Calculate output in batches of pixels
for (int j = 0; j < width; j += kBatchSize) {
const int k = i * width + j;
const int m = i * dst_stride + j;
const uint8_t *dat8ij = dat8 + i * stride + j;
auto ep_0 = hn::Undefined(int32_tag);
auto ep_1 = hn::Undefined(int32_tag);
if (highbd) {
constexpr hn::Repartition<uint16_t, hn::Half<decltype(int32_tag)>>
uint16_tag;
const auto src_0 = hn::LoadU(uint16_tag, CONVERT_TO_SHORTPTR(dat8ij));
const auto src_1 = hn::LoadU(
uint16_tag, CONVERT_TO_SHORTPTR(dat8ij) + hn::MaxLanes(int32_tag));
ep_0 = hn::PromoteTo(int32_tag, src_0);
ep_1 = hn::PromoteTo(int32_tag, src_1);
} else {
constexpr hn::Repartition<uint8_t, hn::Half<decltype(int32_tag)>>
uint8_tag;
const auto src_0 = hn::LoadU(uint8_tag, dat8ij);
ep_0 = hn::PromoteLowerTo(int32_tag, src_0);
ep_1 = hn::PromoteUpperTo(int32_tag, src_0);
}
const auto u_0 = hn::ShiftLeft<SGRPROJ_RST_BITS>(ep_0);
const auto u_1 = hn::ShiftLeft<SGRPROJ_RST_BITS>(ep_1);
auto v_0 = hn::ShiftLeft<SGRPROJ_PRJ_BITS>(u_0);
auto v_1 = hn::ShiftLeft<SGRPROJ_PRJ_BITS>(u_1);
if (params->r[0] > 0) {
const auto f1_0 = hn::Sub(hn::LoadU(int32_tag, &flt0[k]), u_0);
v_0 = hn::Add(v_0, hn::Mul(xq0, f1_0));
const auto f1_1 = hn::Sub(
hn::LoadU(int32_tag, &flt0[k + hn::MaxLanes(int32_tag)]), u_1);
v_1 = hn::Add(v_1, hn::Mul(xq0, f1_1));
}
if (params->r[1] > 0) {
const auto f2_0 = hn::Sub(hn::LoadU(int32_tag, &flt1[k]), u_0);
v_0 = hn::Add(v_0, hn::Mul(xq1, f2_0));
const auto f2_1 = hn::Sub(
hn::LoadU(int32_tag, &flt1[k + hn::MaxLanes(int32_tag)]), u_1);
v_1 = hn::Add(v_1, hn::Mul(xq1, f2_1));
}
const auto rounding =
RoundForShift(int32_tag, SGRPROJ_PRJ_BITS + SGRPROJ_RST_BITS);
const auto w_0 = hn::ShiftRight<SGRPROJ_PRJ_BITS + SGRPROJ_RST_BITS>(
hn::Add(v_0, rounding));
const auto w_1 = hn::ShiftRight<SGRPROJ_PRJ_BITS + SGRPROJ_RST_BITS>(
hn::Add(v_1, rounding));
if (highbd) {
// Pack into 16 bits and clamp to [0, 2^bit_depth)
// Note that packing into 16 bits messes up the order of the bits,
// so we use a permute function to correct this
constexpr hn::Repartition<uint16_t, decltype(int32_tag)> uint16_tag;
const auto tmp = hn::OrderedDemote2To(uint16_tag, w_0, w_1);
const auto max = hn::Set(uint16_tag, (1 << bit_depth) - 1);
const auto res = hn::Min(tmp, max);
hn::StoreU(res, uint16_tag, CONVERT_TO_SHORTPTR(dst8 + m));
} else {
// Pack into 8 bits and clamp to [0, 256)
// Note that each pack messes up the order of the bits,
// so we use a permute function to correct this
constexpr hn::Repartition<int16_t, decltype(int32_tag)> int16_tag;
constexpr hn::Repartition<uint8_t, hn::Half<decltype(int32_tag)>>
uint8_tag;
const auto tmp = hn::OrderedDemote2To(int16_tag, w_0, w_1);
const auto res = hn::DemoteTo(uint8_tag, tmp);
hn::StoreU(res, uint8_tag, dst8 + m);
}
}
}
return 0;
}
} // namespace HWY_NAMESPACE
} // namespace
HWY_AFTER_NAMESPACE();
#define MAKE_SELFGUIDED_RESTORATION(suffix) \
extern "C" int av1_selfguided_restoration_##suffix( \
const uint8_t *dgd8, int width, int height, int dgd_stride, \
int32_t *flt0, int32_t *flt1, int flt_stride, int sgr_params_idx, \
int bit_depth, int highbd); \
HWY_ATTR HWY_NOINLINE int av1_selfguided_restoration_##suffix( \
const uint8_t *dgd8, int width, int height, int dgd_stride, \
int32_t *flt0, int32_t *flt1, int flt_stride, int sgr_params_idx, \
int bit_depth, int highbd) { \
return HWY_NAMESPACE::SelfGuidedRestoration( \
dgd8, width, height, dgd_stride, flt0, flt1, flt_stride, \
sgr_params_idx, bit_depth, highbd); \
} \
extern "C" int av1_apply_selfguided_restoration_##suffix( \
const uint8_t *dat8, int width, int height, int stride, int eps, \
const int *xqd, uint8_t *dst8, int dst_stride, int32_t *tmpbuf, \
int bit_depth, int highbd); \
HWY_ATTR int av1_apply_selfguided_restoration_##suffix( \
const uint8_t *dat8, int width, int height, int stride, int eps, \
const int *xqd, uint8_t *dst8, int dst_stride, int32_t *tmpbuf, \
int bit_depth, int highbd) { \
return HWY_NAMESPACE::ApplySelfGuidedRestoration( \
dat8, width, height, stride, eps, xqd, dst8, dst_stride, tmpbuf, \
bit_depth, highbd); \
}
#endif // AV1_COMMON_SELFGUIDED_HWY_H_