AVX2 中 log2(__m256d) 的高效实现

Efficient implementation of log2(__m256d) in AVX2

SVML 的 __m256d _mm256_log2_pd (__m256d a) 在除 Intel 以外的其他编译器上不可用,他们说它的性能在 AMD 处理器上有缺陷。 AVX log intrinsics (_mm256_log_ps) missing in g++-4.8? and SIMD math libraries for SSE and AVX , however they seem to be more SSE than AVX2. There's also Agner Fog's vector library 中提到了 Internet 上的一些实现,但是它是一个大型库,除了 vector log2 之外还有更多的东西,因此从其中的实现很难找出仅 vector log2 操作的基本部分。

那么有人可以解释一下如何有效地为 4 double 个数字的向量实现 log2() 操作吗? IE。就像 __m256d _mm256_log2_pd (__m256d a) 所做的那样,但可用于其他编译器并且对 AMD 和 Intel 处理器都相当有效。

编辑:在我当前的具体情况下,数字是 0 到 1 之间的概率,对数用于熵计算:P[i]*log(P[i]) 的所有 i 的总和的否定。 P[i] 的浮点指数范围很大,所以数字可以接近于 0。我不确定准确性,所以会考虑以 30 位尾数开头的任何解决方案,尤其是可调解决方案是首选。

EDIT2:到目前为止,这是我的实现,基于 https://en.wikipedia.org/wiki/Logarithm#Power_series 中的 "More efficient series"。如何改进? (性能和准确性都需要改进)

namespace {
  const __m256i gDoubleExpMask = _mm256_set1_epi64x(0x7ffULL << 52);
  const __m256i gDoubleExp0 = _mm256_set1_epi64x(1023ULL << 52);
  const __m256i gTo32bitExp = _mm256_set_epi32(0, 0, 0, 0, 6, 4, 2, 0);
  const __m128i gExpNormalizer = _mm_set1_epi32(1023);
  //TODO: some 128-bit variable or two 64-bit variables here?
  const __m256d gCommMul = _mm256_set1_pd(2.0 / 0.693147180559945309417); // 2.0/ln(2)
  const __m256d gCoeff1 = _mm256_set1_pd(1.0 / 3);
  const __m256d gCoeff2 = _mm256_set1_pd(1.0 / 5);
  const __m256d gCoeff3 = _mm256_set1_pd(1.0 / 7);
  const __m256d gCoeff4 = _mm256_set1_pd(1.0 / 9);
  const __m256d gVect1 = _mm256_set1_pd(1.0);
}

__m256d __vectorcall Log2(__m256d x) {
  const __m256i exps64 = _mm256_srli_epi64(_mm256_and_si256(gDoubleExpMask, _mm256_castpd_si256(x)), 52);
  const __m256i exps32_avx = _mm256_permutevar8x32_epi32(exps64, gTo32bitExp);
  const __m128i exps32_sse = _mm256_castsi256_si128(exps32_avx);
  const __m128i normExps = _mm_sub_epi32(exps32_sse, gExpNormalizer);
  const __m256d expsPD = _mm256_cvtepi32_pd(normExps);
  const __m256d y = _mm256_or_pd(_mm256_castsi256_pd(gDoubleExp0),
    _mm256_andnot_pd(_mm256_castsi256_pd(gDoubleExpMask), x));

  // Calculate t=(y-1)/(y+1) and t**2
  const __m256d tNum = _mm256_sub_pd(y, gVect1);
  const __m256d tDen = _mm256_add_pd(y, gVect1);
  const __m256d t = _mm256_div_pd(tNum, tDen);
  const __m256d t2 = _mm256_mul_pd(t, t); // t**2

  const __m256d t3 = _mm256_mul_pd(t, t2); // t**3
  const __m256d terms01 = _mm256_fmadd_pd(gCoeff1, t3, t);
  const __m256d t5 = _mm256_mul_pd(t3, t2); // t**5
  const __m256d terms012 = _mm256_fmadd_pd(gCoeff2, t5, terms01);
  const __m256d t7 = _mm256_mul_pd(t5, t2); // t**7
  const __m256d terms0123 = _mm256_fmadd_pd(gCoeff3, t7, terms012);
  const __m256d t9 = _mm256_mul_pd(t7, t2); // t**9
  const __m256d terms01234 = _mm256_fmadd_pd(gCoeff4, t9, terms0123);

  const __m256d log2_y = _mm256_mul_pd(terms01234, gCommMul);
  const __m256d log2_x = _mm256_add_pd(log2_y, expsPD);

  return log2_x;
}

到目前为止,我的实现给出了每秒 405 268 490 次操作,而且它看起来精确到第 8 位。使用以下函数测量性能:

#include <chrono>
#include <cmath>
#include <cstdio>
#include <immintrin.h>

// ... Log2() implementation here

const int64_t cnLogs = 100 * 1000 * 1000;

void BenchmarkLog2Vect() {
  __m256d sums = _mm256_setzero_pd();
  auto start = std::chrono::high_resolution_clock::now();
  for (int64_t i = 1; i <= cnLogs; i += 4) {
    const __m256d x = _mm256_set_pd(double(i+3), double(i+2), double(i+1), double(i));
    const __m256d logs = Log2(x);
    sums = _mm256_add_pd(sums, logs);
  }
  auto elapsed = std::chrono::high_resolution_clock::now() - start;
  double nSec = 1e-6 * std::chrono::duration_cast<std::chrono::microseconds>(elapsed).count();
  double sum = sums.m256d_f64[0] + sums.m256d_f64[1] + sums.m256d_f64[2] + sums.m256d_f64[3];
  printf("Vect Log2: %.3lf Ops/sec calculated %.3lf\n", cnLogs / nSec, sum);
}

的结果相比,当前向量实现比 std::log2() 快 4 倍,比 std::log().

快 2.5 倍

具体使用如下近似公式:

通常的策略是基于身份log(a*b) = log(a) + log(b),或者在本例中是log2( 2^exponent * mantissa) ) = log2( 2^exponent ) + log2(mantissa)。或者简化,exponent + log2(mantissa)。尾数的范围非常有限,从 1.0 到 2.0,因此 log2(mantissa) 的多项式只需要适合这个非常有限的范围。 (或者等效地,尾数 = 0.5 到 1.0,并将指数 bias-correction 常数更改为 1)。

泰勒级数展开式是系数的良好起点,但您通常希望在该特定范围内最小化 max-absolute-error(或相对误差),并且泰勒级数系数可能具有较低或该范围内的异常值更高,而不是使最大正误差几乎与最大负误差匹配。因此,您可以执行所谓的系数极小极大拟合。

如果您的函数将 log2(1.0) 精确计算为 0.0 很重要,您可以通过实际使用 mantissa-1.0 作为您的多项式来安排它发生,并且没有常数系数。 0.0 ^ n = 0.0。这也大大改善了接近 1.0 的输入的 相对 误差,即使绝对误差仍然很小。


您需要它有多准确,在什么输入范围内?像往常一样,在准确性和速度之间需要权衡,但幸运的是,沿着这个比例移动是很容易的,例如。添加一个多项式项(和 re-fitting 系数),或者通过删除一些 rounding-error 回避。

Agner Fog's VCL implementation of log_d() 的目标是非常高的准确性,使用技巧来避免舍入错误,方法是避免可能导致添加小数和大数的事情。这在一定程度上掩盖了基本设计。


要获得更快更近似的 float log(),请参阅 http://jrfonseca.blogspot.ca/2008/09/fast-sse2-pow-tables-or-polynomials.html 上的多项式实现。它省略了很多 VCL 使用的 precision-gaining 技巧,因此更容易理解。它对 1.0 到 2.0 范围内的尾数使用多项式近似。

(这是 log() 实现的真正诀窍:您只需要一个适用于小范围的多项式。)

它已经只是 log2 而不是 log,这与 VCL 不同,其中 log-base-e 被烘焙到常量中以及它如何使用它们。阅读它可能是理解 exponent + polynomial(mantissa) 实现 log().

的一个很好的起点

即使是它的 highest-precision 版本也不是完整的 float 精度,更不用说 double,但是您可以用更多项来拟合多项式。或者显然两个多项式的比率效果很好;这就是 VCL 用于 double.

我将 JRF 的 SSE2 函数移植到 AVX2 + FMA(尤其是具有 _mm512_getexp_ps_mm512_getmant_ps 的 AVX512),经过仔细调整后,我得到了很好的结果。 (它是商业项目的一部分,所以我认为我无法 post 代码。)float 的快速近似实现正是我想要的。

在我的 use-case 中,每个 jrf_fastlog() 都是独立的,因此 OOO 执行很好地隐藏了 FMA 延迟,甚至不值得使用 higher-ILP shorter-latency VCL's polynomial_5() function uses ("Estrin's scheme" 的多项式评估方法,它在 FMA 之前执行一些 non-FMA 乘法运算,从而产生更多的总指令)。


Agner Fog 的 VCL 现在是 Apache-licensed,因此任何项目都可以直接包含它。如果你想要高精度,你应该直接使用VCL。它是 header-only,只是内联函数,因此不会使您的二进制文件膨胀。

VCL 的 log float 和 double 函数在 vectormath_exp.h 中。该算法有两个主要部分:

  • 提取指数位并将该整数转换回浮点数(在调整 IEEE FP 使用的偏差后)。

  • 提取某些指数位中的尾数和 OR 以获得 [0.5, 1.0) 范围内 double 值的向量。 (或者(0.5, 1.0],我忘记了)。

    进一步调整if(mantissa <= SQRT2*0.5) { mantissa += mantissa; exponent++;},然后mantissa -= 1.0

    使用多项式逼近 log(x),它在 x=1.0 附近是准确的。 (对于 double,VCL 的 log_d() 使用两个 5 阶多项式的比率。@harold says this is often good for precision。与大量 FMA 混合的一个除法通常不会影响吞吐量,但它确实有比 FMA 延迟更高。使用 vrcpps + Newton-Raphson 迭代通常比在现代硬件上仅使用 vdivps 慢。使用比率还可以通过评估两个 lower-order 来创建更多的 ILP并行多项式,而不是一个 high-order 多项式,与 high-order 多项式的一个长 dep 链相比,整体延迟可能更低(这也会沿着那个长链累积显着的舍入误差)。

然后加上exponent + polynomial_approx_log(mantissa)得到最后的log()结果。 VCL 分多个步骤执行此操作以减少舍入误差。 ln2_lo + ln2_hi = ln(2)。它分为一个小常量和一个大常量以减少舍入误差。

// res is the polynomial(adjusted_mantissa) result
// fe is the float exponent
// x is the adjusted_mantissa.  x2 = x*x;
res  = mul_add(fe, ln2_lo, res);             // res += fe * ln2_lo;
res += nmul_add(x2, 0.5, x);                 // res += x  - 0.5 * x2;
res  = mul_add(fe, ln2_hi, res);             // res += fe * ln2_hi;

如果您的目标不是 0.5 或 1 ulp 精度(或此函数实际提供的任何内容;IDK),您可以放弃 2 步 ln2 内容而只使用 VM_LN2。 =69=]

我猜 x - 0.5*x2 部分实际上是一个额外的多项式项。这就是我所说的 log base e 是 baked-in 的意思:你'需要这些项的系数,或者摆脱那条线和 re-fit log2 的多项式系数。您不能只将所有多项式系数乘以一个常数。

之后,它会检查是否存在下溢、溢出或异常,如果向量中的任何元素需要特殊处理以生成适当的 NaN 或 -Inf 而不是我们从多项式 + 指数中得到的任何垃圾,则进行分支。 如果已知您的值是有限且正的,则可以注释掉这部分并获得显着的加速(甚至分支之前的检查也需要几条指令)。


延伸阅读:

  • http://gallium.inria.fr/blog/fast-vectorizable-math-approx/ 一些关于如何评估多项式近似中的相对和绝对误差,以及对系数进行极小极大修正而不是仅仅使用泰勒级数展开的一些内容。

  • http://www.machinedlearnings.com/2011/06/fast-approximate-logarithm-exponential.html 一个有趣的方法:它 type-puns 一个 floatuint32_t,并且 将该整数转换为 float。由于 IEEE binary32 浮点数以比尾数更高的位存储指数,因此生成的 float 主要表示指数的值,按 1 << 23 缩放,但也包含来自尾数的信息。

    然后它使用带有几个系数的表达式来解决问题并获得 log() 近似值。它包括除以 (constant + mantissa) 以纠正将浮点数 bit-pattern 转换为 float 时的尾数污染。我发现在 HSW 和 SKL 上使用 AVX2 的矢量化版本比使用四阶多项式的 JRF fastlog 更慢且更不准确。 (特别是将它用作快速 arcsinh 的一部分时,它也使用 vsqrtps 的除法单元。)

最后这是我最好的结果,它在 Ryzen 1800X @3.6GHz 上在单个线程中每秒给出大约 8 亿个对数(每个 4 个对数的 2 亿个向量),并且直到最后几位都是准确的尾数。 剧透:看最后如何把性能提升到8.7亿对数/秒

特殊情况: 负数、负无穷大和带负号位的 NaN 被视为非常接近 0(导致一些垃圾大负 "logarithm" 值)。正无穷大和带正号位的 NaNs 产生 1024 左右的对数。如果您不喜欢处理特殊情况的方式,一种选择是添加代码来检查它们并做更适合您的事情。这会使计算变慢。

namespace {
  // The limit is 19 because we process only high 32 bits of doubles, and out of
  //   20 bits of mantissa there, 1 bit is used for rounding.
  constexpr uint8_t cnLog2TblBits = 10; // 1024 numbers times 8 bytes = 8KB.
  constexpr uint16_t cZeroExp = 1023;
  const __m256i gDoubleNotExp = _mm256_set1_epi64x(~(0x7ffULL << 52));
  const __m256d gDoubleExp0 = _mm256_castsi256_pd(_mm256_set1_epi64x(1023ULL << 52));
  const __m256i cAvxExp2YMask = _mm256_set1_epi64x(
    ~((1ULL << (52-cnLog2TblBits)) - 1) );
  const __m256d cPlusBit = _mm256_castsi256_pd(_mm256_set1_epi64x(
    1ULL << (52 - cnLog2TblBits - 1)));
  const __m256d gCommMul1 = _mm256_set1_pd(2.0 / 0.693147180559945309417); // 2.0/ln(2)
  const __m256i gHigh32Permute = _mm256_set_epi32(0, 0, 0, 0, 7, 5, 3, 1);
  const __m128i cSseMantTblMask = _mm_set1_epi32((1 << cnLog2TblBits) - 1);
  const __m128i gExpNorm0 = _mm_set1_epi32(1023);
  // plus |cnLog2TblBits|th highest mantissa bit
  double gPlusLog2Table[1 << cnLog2TblBits];
} // anonymous namespace

void InitLog2Table() {
  for(uint32_t i=0; i<(1<<cnLog2TblBits); i++) {
    const uint64_t iZp = (uint64_t(cZeroExp) << 52)
      | (uint64_t(i) << (52 - cnLog2TblBits)) | (1ULL << (52 - cnLog2TblBits - 1));
    const double zp = *reinterpret_cast<const double*>(&iZp);
    const double l2zp = std::log2(zp);
    gPlusLog2Table[i] = l2zp;
  }
}

__m256d __vectorcall Log2TblPlus(__m256d x) {
  const __m256d zClearExp = _mm256_and_pd(_mm256_castsi256_pd(gDoubleNotExp), x);
  const __m256d z = _mm256_or_pd(zClearExp, gDoubleExp0);

  const __m128i high32 = _mm256_castsi256_si128(_mm256_permutevar8x32_epi32(
    _mm256_castpd_si256(x), gHigh32Permute));
  // This requires that x is non-negative, because the sign bit is not cleared before
  //   computing the exponent.
  const __m128i exps32 = _mm_srai_epi32(high32, 20);
  const __m128i normExps = _mm_sub_epi32(exps32, gExpNorm0);

  // Compute y as approximately equal to log2(z)
  const __m128i indexes = _mm_and_si128(cSseMantTblMask,
    _mm_srai_epi32(high32, 20 - cnLog2TblBits));
  const __m256d y = _mm256_i32gather_pd(gPlusLog2Table, indexes,
    /*number of bytes per item*/ 8);
  // Compute A as z/exp2(y)
  const __m256d exp2_Y = _mm256_or_pd(
    cPlusBit, _mm256_and_pd(z, _mm256_castsi256_pd(cAvxExp2YMask)));

  // Calculate t=(A-1)/(A+1). Both numerator and denominator would be divided by exp2_Y
  const __m256d tNum = _mm256_sub_pd(z, exp2_Y);
  const __m256d tDen = _mm256_add_pd(z, exp2_Y);

  // Compute the first polynomial term from "More efficient series" of https://en.wikipedia.org/wiki/Logarithm#Power_series
  const __m256d t = _mm256_div_pd(tNum, tDen);

  const __m256d log2_z = _mm256_fmadd_pd(t, gCommMul1, y);

  // Leading integer part for the logarithm
  const __m256d leading = _mm256_cvtepi32_pd(normExps);

  const __m256d log2_x = _mm256_add_pd(log2_z, leading);
  return log2_x;
}

它结合了查找 table 方法和一次多项式,主要在维基百科上进行了描述(link 在代码注释中)。我可以在这里分配 8KB 的 L1 缓存(这是每个逻辑核心可用的 16KB L1 缓存的一半),因为对数计算对我来说确实是瓶颈,而且没有更多的东西需要 L1 缓存。

但是,如果您需要更多的 L1 缓存来满足其他需求,您可以通过减少 cnLog2TblBits 来减少对数算法使用的缓存量,例如5 以降低对数计算的准确性为代价。

或者为了保持较高的准确性,您可以通过添加以下内容来增加多项式项的数量:

namespace {
  // ...
  const __m256d gCoeff1 = _mm256_set1_pd(1.0 / 3);
  const __m256d gCoeff2 = _mm256_set1_pd(1.0 / 5);
  const __m256d gCoeff3 = _mm256_set1_pd(1.0 / 7);
  const __m256d gCoeff4 = _mm256_set1_pd(1.0 / 9);
  const __m256d gCoeff5 = _mm256_set1_pd(1.0 / 11);
}

然后在 const __m256d t = _mm256_div_pd(tNum, tDen); 行之后更改 Log2TblPlus() 的尾部:

  const __m256d t2 = _mm256_mul_pd(t, t); // t**2

  const __m256d t3 = _mm256_mul_pd(t, t2); // t**3
  const __m256d terms01 = _mm256_fmadd_pd(gCoeff1, t3, t);
  const __m256d t5 = _mm256_mul_pd(t3, t2); // t**5
  const __m256d terms012 = _mm256_fmadd_pd(gCoeff2, t5, terms01);
  const __m256d t7 = _mm256_mul_pd(t5, t2); // t**7
  const __m256d terms0123 = _mm256_fmadd_pd(gCoeff3, t7, terms012);
  const __m256d t9 = _mm256_mul_pd(t7, t2); // t**9
  const __m256d terms01234 = _mm256_fmadd_pd(gCoeff4, t9, terms0123);
  const __m256d t11 = _mm256_mul_pd(t9, t2); // t**11
  const __m256d terms012345 = _mm256_fmadd_pd(gCoeff5, t11, terms01234);

  const __m256d log2_z = _mm256_fmadd_pd(terms012345, gCommMul1, y);

然后评论// Leading integer part for the logarithm,其余不变。

一般情况下不需要那么多的项,即使是几位table,我只是提供了系数和计算以供参考。如果 cnLog2TblBits==5,您很可能不需要 terms012 以外的任何东西。但是我没有做过这样的测量,你需要自己试验一下。

您计算的多项式项越少,显然,计算速度越快。


编辑:这个问题In what situation would the AVX2 gather instructions be faster than individually loading the data?表明如果

你可能会得到性能提升
const __m256d y = _mm256_i32gather_pd(gPlusLog2Table, indexes,
  /*number of bytes per item*/ 8);

取代
const __m256d y = _mm256_set_pd(gPlusLog2Table[indexes.m128i_u32[3]],
  gPlusLog2Table[indexes.m128i_u32[2]],
  gPlusLog2Table[indexes.m128i_u32[1]],
  gPlusLog2Table[indexes.m128i_u32[0]]);

对于我的实现,它节省了大约 1.5 个周期,将计算 4 个对数的总周期数从 18 减少到 16.5,因此性能提高到每秒 8.7 亿个对数。我将按原样离开当前的实现,因为它更惯用,并且一旦 CPU 开始正确地执行 gather 操作(像 GPU 那样合并)应该会更快。

EDIT2: 你可以通过替换

获得更多的加速(大约 0.5 个周期)
const __m128i high32 = _mm256_castsi256_si128(_mm256_permutevar8x32_epi32(
  _mm256_castpd_si256(x), gHigh32Permute));

  const __m128 hiLane = _mm_castpd_ps(_mm256_extractf128_pd(x, 1));
  const __m128 loLane = _mm_castpd_ps(_mm256_castpd256_pd128(x));
  const __m128i high32 = _mm_castps_si128(_mm_shuffle_ps(loLane, hiLane,
    _MM_SHUFFLE(3, 1, 3, 1)));