使用内在指令的欧几里得距离

Euclidean distance using intrinsic instruction

本文关键字:几里 距离 指令      更新时间:2023-10-16

对于一个研究项目,我需要计算很多欧几里得距离,其中必须选择某些维度,而放弃其他维度。在程序的当前状态下,所选维度的数组有100个元素,我计算了大约200-300万个距离。我当前的代码如下:

float compute_distance(const float* p1, const float* p2) const
{
__m256 euclidean = _mm256_setzero_ps();
const uint16_t n = nbr_dimensions;
const uint16_t aligend_n = n - n % 16;
const float* local_selected = selected_dimensions;
for (uint16_t i = 0; i < aligend_n; i += 16)
{
const __m256 r1 = _mm256_sub_ps(_mm256_load_ps(&p1[i]), _mm256_load_ps(&p2[i]));
euclidean = _mm256_fmadd_ps(_mm256_mul_ps(r1, r1), _mm256_load_ps(&local_selected[i]), euclidean);
const __m256 r2 = _mm256_sub_ps(_mm256_load_ps(&p1[i + 8]), _mm256_load_ps(&p2[i + 8]));
euclidean = _mm256_fmadd_ps(_mm256_mul_ps(r2, r2), _mm256_load_ps(&local_selected[i + 8]), euclidean);
}
float distance = hsum256_ps_avx(euclidean);
for (uint16_t i = aligend_n; i < n; ++i)
{
const float num = p1[i] - p2[i];
distance += num * num * local_selected[i];
}
return distance;
}

选定的尺寸是预先确定的。因此,我可以预先计算__m256的数组以传递到_mm256_blendv_ps,而不是在行euclidean = _mm256_fmadd_ps(_mm256_mul_ps(r1, r1), _mm256_load_ps(&local_selected[i]), euclidean);中乘以0或1。但我在内在指令方面还是个新手,我还没有找到一个有效的解决方案。

我想知道你们是否可以有一些建议,甚至代码建议,来提高这个函数的运行速度。顺便说一句,我无法访问AVX-512指令。


更新:使用上述第一个解决方案,它来了:

float compute_distance(const float* p1, const float* p2) const
{
const size_t n = nbr_dimensions;
const size_t aligend_n = n - n % 16;
const unsigned int* local_selected = selected_dimensions;
const __m256* local_masks = masks;
__m256 euc1 = _mm256_setzero_ps(), euc2 = _mm256_setzero_ps(),
euc3 = _mm256_setzero_ps(), euc4 = _mm256_setzero_ps();
const size_t n_max = aligend_n/8;
for (size_t i = 0; i < n_max; i += 4)
{       
const __m256 r1 = _mm256_sub_ps(_mm256_load_ps(&p1[i * 8 + 0]), _mm256_load_ps(&p2[i * 8 + 0]));
const __m256 r1_1 = _mm256_and_ps(r1, local_masks[i + 0]);
euc1 = _mm256_fmadd_ps(r1_1, r1_1, euc1);
const __m256 r2 = _mm256_sub_ps(_mm256_load_ps(&p1[i * 8 + 8]), _mm256_load_ps(&p2[i * 8 + 8]));
const __m256 r2_1 = _mm256_and_ps(r2, local_masks[i + 1]);
euc2 = _mm256_fmadd_ps(r2_1, r2_1, euc2);
const __m256 r3 = _mm256_sub_ps(_mm256_load_ps(&p1[i * 8 + 16]), _mm256_load_ps(&p2[i * 8 + 16]));
const __m256 r3_1 = _mm256_and_ps(r3, local_masks[i + 2]);
euc3 = _mm256_fmadd_ps(r3_1, r3_1, euc3);
const __m256 r4 = _mm256_sub_ps(_mm256_load_ps(&p1[i * 8 + 24]), _mm256_load_ps(&p2[i * 8 + 24]));
const __m256 r4_1 = _mm256_and_ps(r4, local_masks[i + 3]);
euc4 = _mm256_fmadd_ps(r4_1, r4_1, euc4);
}
float distance = hsum256_ps_avx(_mm256_add_ps(_mm256_add_ps(euc1, euc2), _mm256_add_ps(euc3, euc4)));
for (size_t i = aligend_n; i < n; ++i)
{
const float num = p1[i] - p2[i];
distance += num * num * local_selected[i];      
}
return distance;
}

基本建议:

不要将uint16_t用于循环计数器,除非您真的想强制编译器每次截断为16位。至少使用unsigned,或者有时使用uintptr_t(或者更传统的size_t)会得到更好的asm。从32位到指针宽度的零扩展在x86-64上是免费的,只需使用32位操作数大小的asm指令,但有时编译器仍然做得不好。

使用五个或多个单独的累加器而不是一个euclidean,因此多个子/FMA指令可以在飞行中,而不会对将FMA转换为一个累加器的循环携带依赖链的延迟造成瓶颈。

在Intel Haswell上,FMA的延迟为5个周期,但吞吐量为每0.5个周期一个。另请参阅intel内部函数中的延迟与吞吐量,以及我对为什么mulss在Haswell上只需要3个周期的回答,这与Agner';s指令表?以获得更高级的版本。


避免通过全局变量传递参数。显然,您的n是一个编译时常数(这很好),但selected_dimensions不是,是吗?如果是的话,那么你在整个程序中只使用了一组掩码,所以不要介意下面关于压缩掩码的内容。

当使用全局变量将函数内联到调用方时,编译器优化可能会失败。调用方在调用全局变量之前设置了全局变量。(通常只有在设置全局变量和使用全局变量之间有非内联函数调用的情况下,但这并不罕见。)


更新:您的数组很小,只有大约100个元素,所以只展开2个可能很好,可以减少启动/清理开销。无序执行可能会在短时间的迭代中隐藏FMA延迟,尤其是在不需要此函数调用的最终结果来决定下一次调用的输入参数的情况下。

总的函数调用开销很重要,而不仅仅是矢量化对大型数组的效率。

正如评论中所讨论的,剥离循环的第一次迭代可以通过初始化euc1 = stuff(p1[0], p2[0]);而不是_mm256_setzero_ps()来避免第一次FMA。

用零将数组填充到一个完整的向量(甚至是一个由2个向量组成的完整展开的循环体),可以完全避免标量清理循环,并使整个函数非常紧凑。

如果你不能只填充,你仍然可以通过加载一个一直到输入末尾的未对齐向量来避免标量清理,并屏蔽它以避免重复计数。(请参阅本问答,了解基于未对准计数生成掩模的方法)。在编写输出数组的其他类型的问题中,可以重做重叠的元素。

您没有显示hsum256_ps_avx代码,但这只是总延迟的一小部分,可能是函数吞吐量的一部分。确保针对吞吐量进行优化:例如,避免haddps/_mm_hadd_ps。请参阅我在x86上进行水平浮点矢量求和的最快方法上的答案。


您的具体案例

因此,我可以预先计算__m256的数组,以传递给_mm256_blendv_ps,而不是在FMA中乘以0或1。

是的,那会更好,尤其是如果它允许您将其他东西折叠到FMAdd/FMSub中。但比这更好的是,使用一个全为零或全为一的布尔_mm256_and_ps。这使值保持不变(1 & x == x)或为零(0 & x == 0,浮点0.0的二进制表示为全零。)

如果你的口罩在缓存中没有丢失,那么把它们完全打开包装保存起来,这样就可以直接加载了。

如果使用具有相同p1p2的不同掩码,则可以预先计算p1-p2的平方,然后对其进行掩码add_ps约简。(但请注意,FMA的吞吐量比英特尔Skylake之前的ADD要好。Haswell/Broadwell有2个FMA单元,但在延迟较低的专用单元上运行ADDPS(3c比5c)。只有一个矢量FP加法单元。Skylake只在FMA单元上运行所有程序,延迟4个周期。)无论如何,这意味着使用FMA作为1.0 * x + y实际上是一个吞吐量的胜利。但您可能很好,因为您仍然需要分别加载掩码和square(p1-p2),因此每个FP添加2个负载,因此每个循环一个负载可以跟上负载吞吐量。除非您(或编译器)在前面剥离一些迭代,并将这些迭代的浮点数据保存在多个不同local_selected掩码的寄存器中。


更新:我写这篇文章时假设数组大小是200-300万,而不是大约100。L1D缓存的配置文件未命中,无法决定是否值得花费更多CPU指令来减少缓存占用。如果您总是对所有300万个调用使用相同的掩码,则可能不值得压缩。

您可以将掩码压缩到每个元素8位,并用pmovsx(_mm256_cvtepi8_epi32)加载它们(扩展全1值的符号会产生更宽的全1,因为2的补码-1就是这样工作的)。不幸的是,将其用作负载是令人讨厌的;编译器有时无法将_mm256_cvtepi8_epi32(_mm_cvtsi64x_si128(foo))优化为vpmovsxbd ymm0, [mem],而是实际使用单独的vmovq指令。

const uint64_t *local_selected = something;  // packed to 1B per element
__m256 euc1 = _mm256_setzero_ps(), euc2 = _mm256_setzero_ps(),
euc3 =  _mm256_setzero_ps(), euc4 =  _mm256_setzero_ps();
for (i = 0 ; i < n ; i += 8*4) {  // 8 floats * an unroll of 4
__m256 mask = _mm256_castsi256_ps( _mm256_cvtepi8_epi32(_mm_cvtsi64x_si128(local_selected[i*1 + 0])) );
// __m256 mask = _mm256_load_ps(local_selected[i*8 + 0]); //  without packing
const __m256 r1 = _mm256_sub_ps(_mm256_load_ps(&p1[i*8 + 0]), _mm256_load_ps(&p2[i*8 + 0]));
r1 = _mm256_and_ps(r1, mask);             // zero r1 or leave it untouched.
euc1 = _mm256_fmadd_ps(r1, r1, euc1);    // euc1 += r1*r1
// ... same for r2 with local_selected[i + 1]
// and p1/p2[i*8 + 8]
// euc2 += (r2*r2) & mask2
// and again for euc3 (local_selected[i + 2], p1/p2[i*8 + 16]
// and again for euc3 (local_selected[i + 3], p1/p2[i*8 + 24]
}
euclidean = hsum (euc1+euc2+euc3+euc4);

我想在没有pmovsx的情况下,您的负载吞吐量会有点瓶颈,因为您有三个负载用于三个矢量ALU操作。(对于微融合,它在英特尔CPU上总共只有4个融合域uop,所以它在前端不会受到瓶颈限制)。三个ALU uop可以在不同的端口上运行(vandps是英特尔Skylake之前的端口5的1个uop。在SKL上,它可以在任何端口上运行)。

在端口5(Haswell/Broadwell)上添加混洗(pmovsx)可能会造成瓶颈。如果您正在为HSW/BDW进行调优,您可能需要使用vpand进行屏蔽,这样它就可以在任何端口上运行,即使它们在整数AND和FP数学指令之间有额外的旁路延迟。有了足够的累加器,您就不会受到延迟限制。(Skylake对VANDPS有额外的旁路延迟,这取决于它恰好在哪个端口上运行)。

blendv比AND慢:总是至少2个uops。


对于大型阵列,进一步压缩掩码

如果您的数组大于二级缓存,并且掩码数组的元素与浮点数组的元素一样多,那么您很可能会在负载带宽上遇到瓶颈(至少在使用多个向量累加器展开时)。这意味着花更多的指令来拆包掩码数据是值得的,以减少这部分带宽需求。

我认为掩码数据的理想格式是32个掩码矢量交错,这使得在飞行中"打开"非常便宜。使用移位将右侧掩码带入每个32位元素的高位,并将其与vblendvps一起使用,通过与零混合来有条件地将元素归零。(或算术右移+布尔AND)

__m256i masks = _mm256_load_si256(...);
// this actually needs a cast to __m256, omitted for readability
r0 = _mm256_blendv_ps(_mm256_setzero_ps(), r0, masks);
...
__m256i mask1 = _mm256_slli_epi32(masks, 1);
r1 = _mm256_blendv_ps(_mm256_setzero_ps(), r1, mask1);
...
__m256i mask2 = _mm256_slli_epi32(masks, 2);
r2 = _mm256_blendv_ps(_mm256_setzero_ps(), r2, mask2);
...
// fully unrolling is overkill; you can set up for a loop back to r0 with
masks = _mm256_slli_epi32(masks, 4);

您也可以在每一步中执行masks = _mm256_slli_epi32(masks, 1);,这可能会更好,因为它少使用1个寄存器。但它可能对导致掩码dep链延迟的资源冲突更敏感,因为每个掩码都依赖于前一个掩码。

Intel Haswell仅在端口5上运行两个vblendvpsuop,因此您可以考虑使用_mm256_srai_epi32+_mm256_and_ps。但是Skylake可以在p015中的任何一个上运行2个uop,所以在那里混合很好(尽管它确实占用了一个保存全零向量的向量寄存器)。


用压缩比较生成交错格式的掩码,然后_mm256_srli_epi32(cmp_result, 31)并将其与您正在构建的向量进行OR运算。然后左移1。重复32次。


如果数组中的数据完整向量少于32,则仍然可以使用此格式。较低的位将不被使用。或者,您可以为每个向量设置2个或多个selected_dimensions的遮罩。例如每个元素的前16位用于一个selected_dimensions,而下16位用于另一个。你可以做一些类似的事情

__m256i masks =  _mm256_load_si256(dimensions[selector/2]);
masks = _mm256_sll_epi32(masks, 16 * (selector % 2));
// or maybe
if (selector % 2) {
masks = _mm256_slli_epi32(masks, 16);
}

AVX512:

AVX512可以直接使用位图掩码,因此效率更高。只需使用const __mmask16 *local_selected = whatever;声明16位掩码的数组(用于16个浮点的512b向量),并使用r0 = _mm512_maskz_sub_ps(p1,p2, local_selected[i]);对减法进行零掩码。

如果你真的在加载端口uop吞吐量上遇到瓶颈(每个时钟2次加载),你可以尝试一次加载64位掩码数据,并使用掩码移位来获得不同的低16位掩码数据。不过,除非你的数据在L1D缓存中很热,否则这可能不会成为问题。

首先通过比较掩码生成掩码数据非常容易,不需要交织。


理想情况下,您可以缓存块调用它的代码,这样您就可以在数据在缓存中处于热状态时重用数据。例如,从p1和p2的前64kiB中获得您想要的所有组合,然后转到后面的元素,并在它们在缓存中处于热状态时进行这些操作。