AVX2:计算 512 个浮点数组的点积

AVX2: Computing dot product of 512 float arrays

本文关键字:数组 计算 AVX2      更新时间:2023-10-16

首先,我要说的是,我是 SIMD 内联函数的初学者。

从本质上讲,我有一个支持AVX2内部(Intel(R) Core(TM) i5-7500T CPU @ 2.70GHz(的CPU。我想知道计算大小512的两std::vector<float>的点积的最快方法。

我在网上做了一些挖掘,发现了这个和这个,这个堆栈溢出问题建议使用以下函数__m256 _mm256_dp_ps(__m256 m1, __m256 m2, const int mask);,但是,这些都提出了执行点积的不同方法我不确定什么是正确(和最快(的方法。

特别是,我正在寻找对大小为 512 的向量执行点积的最快方法(因为我知道向量大小会影响实现(。

谢谢你的帮助

编辑 1: 我也对 gcc 标志-mavx2有点困惑。如果我使用这些 AVX2 函数,编译时是否需要添加标志?另外,如果我编写一个朴素的点产品实现,gcc 是否能够为我进行这些优化(例如,如果我使用-OFastgcc 标志(?

编辑 2如果有人有时间和精力,如果您能编写一个完整的实现,我将不胜感激。我相信其他初学者也会重视这些信息。

_mm256_dp_ps仅适用于 2 到 4 个元素的点积;对于较长的矢量,在循环中使用垂直 SIMD,并在末尾简化为标量。 在循环中使用_mm256_dp_ps_mm256_add_ps会慢得多。


GCC 和 clang 要求您启用(使用命令行选项(使用内部函数的 ISA 扩展,这与 MSVC 和 ICC 不同。


下面的代码可能接近 CPU 的理论性能限制。未经测试。

用叮当或gcc -O3 -march=native编译它。(至少需要-mavx -mfma,但-march暗示-mtune选项也很好,其他-mpopcntarch=native启用的其他东西也是如此。 对于大多数具有 FMA 的 CPU,调整选项对于有效编译至关重要,特别是-mno-avx256-split-unaligned-load:为什么 gcc 不将_mm256_loadu_pd解析为单个 vmovupd?

或者使用 MSVC-O2 -arch:AVX2编译它

#include <immintrin.h>
#include <vector>
#include <assert.h>
// CPUs support RAM access like this: "ymmword ptr [rax+64]"
// Using templates with offset int argument to make easier for compiler to emit good code.
// Multiply 8 floats by another 8 floats.
template<int offsetRegs>
inline __m256 mul8( const float* p1, const float* p2 )
{
constexpr int lanes = offsetRegs * 8;
const __m256 a = _mm256_loadu_ps( p1 + lanes );
const __m256 b = _mm256_loadu_ps( p2 + lanes );
return _mm256_mul_ps( a, b );
}
// Returns acc + ( p1 * p2 ), for 8-wide float lanes.
template<int offsetRegs>
inline __m256 fma8( __m256 acc, const float* p1, const float* p2 )
{
constexpr int lanes = offsetRegs * 8;
const __m256 a = _mm256_loadu_ps( p1 + lanes );
const __m256 b = _mm256_loadu_ps( p2 + lanes );
return _mm256_fmadd_ps( a, b, acc );
}
// Compute dot product of float vectors, using 8-wide FMA instructions.
float dotProductFma( const std::vector<float>& a, const std::vector<float>& b )
{
assert( a.size() == b.size() );
assert( 0 == ( a.size() % 32 ) );
if( a.empty() )
return 0.0f;
const float* p1 = a.data();
const float* const p1End = p1 + a.size();
const float* p2 = b.data();
// Process initial 32 values. Nothing to add yet, just multiplying.
__m256 dot0 = mul8<0>( p1, p2 );
__m256 dot1 = mul8<1>( p1, p2 );
__m256 dot2 = mul8<2>( p1, p2 );
__m256 dot3 = mul8<3>( p1, p2 );
p1 += 8 * 4;
p2 += 8 * 4;
// Process the rest of the data.
// The code uses FMA instructions to multiply + accumulate, consuming 32 values per loop iteration.
// Unrolling manually for 2 reasons:
// 1. To reduce data dependencies. With a single register, every loop iteration would depend on the previous result.
// 2. Unrolled code checks for exit condition 4x less often, therefore more CPU cycles spent computing useful stuff.
while( p1 < p1End )
{
dot0 = fma8<0>( dot0, p1, p2 );
dot1 = fma8<1>( dot1, p1, p2 );
dot2 = fma8<2>( dot2, p1, p2 );
dot3 = fma8<3>( dot3, p1, p2 );
p1 += 8 * 4;
p2 += 8 * 4;
}
// Add 32 values into 8
const __m256 dot01 = _mm256_add_ps( dot0, dot1 );
const __m256 dot23 = _mm256_add_ps( dot2, dot3 );
const __m256 dot0123 = _mm256_add_ps( dot01, dot23 );
// Add 8 values into 4
const __m128 r4 = _mm_add_ps( _mm256_castps256_ps128( dot0123 ), _mm256_extractf128_ps( dot0123, 1 ) );
// Add 4 values into 2
const __m128 r2 = _mm_add_ps( r4, _mm_movehl_ps( r4, r4 ) );
// Add 2 lower values into the final result
const __m128 r1 = _mm_add_ss( r2, _mm_movehdup_ps( r2 ) );
// Return the lowest lane of the result vector.
// The intrinsic below compiles into noop, modern compilers return floats in the lowest lane of xmm0 register.
return _mm_cvtss_f32( r1 );
}

可能的进一步改进:

按 8 个向量而不是
  1. 4 个向量展开。我已经检查了 gcc 9.2 asm 输出,编译器只使用了 8 个可用向量寄存器中的 16 个。

  2. 确保两个输入向量对齐,例如,使用自定义分配器在msvc上调用_aligned_malloc/_aligned_free,或在gcc和clang上调用aligned_alloc/free。然后将_mm256_loadu_ps替换为_mm256_load_ps


要自动矢量化简单的标量点积,您还需要 OpenMP SIMD 或-ffast-math(由-Ofast暗示(让编译器将 FP 数学视为关联,即使它不是(因为舍入(。 但是 GCC 在自动矢量化时不会使用多个累加器,即使它确实展开了,所以你会成为 FMA 延迟的瓶颈,而不是负载吞吐量。

(每个 FMA 2 个负载意味着此代码的吞吐量瓶颈是矢量负载,而不是实际的 FMA 操作。