AVX 或 SSE 上的水平尾随最大值
Posted
技术标签:
【中文标题】AVX 或 SSE 上的水平尾随最大值【英文标题】：Horizontal trailing maximum on AVX or SSE 【发布时间】：2017-04-06 14:45:31 【问题描述】：我有一个由 16 位值组成的 __m256i 寄存器，我想获取每个尾随元素的最大值，它们是零。
举个例子：
input: 1 0 0 3 0 0 4 5 0 0 0 0 4 3 0 2
output: 1 1 1 3 3 3 4 5 5 5 5 5 4 3 3 2
在 AVX 或 AVX 架构上是否有任何有效的方法来执行此操作？也许 log(16) = 4 次迭代？
补充： 任何关于 128 位数字的解决方案，其中包含 8 个 uint_16，也值得赞赏。
【问题讨论】：
嗯 - 水平操作总是很棘手，尤其是 AVX 中的拆分通道。在最坏的情况下，您必须将单个非零元素传播到 15 个其他元素，因此您可能会查看多达 15 次迭代，每次迭代涉及 2 或 3 条指令。我认为这不会非常有效...... 实际上，可以对 8 个 128 位的数字执行此操作。我正在尝试开发一种算法，因此任何解决方案都会受到赞赏。 第一个元素为零时的行为应该是什么？ 好的，所以0 0 0 3 0 0 4 5 0 0 0 0 4 3 0 2 => 0 0 0 3 3 3 4 5 5 5 5 5 4 3 3 2 ?
pmovmskb + pshufb 技巧可以解决这个问题（虽然查找表大得离谱.. 128 位的情况不会那么糟糕）
【参考方案1】：
您确实可以在log_2(SIMD_width) 步骤中执行此操作。这个想法是将输入向量x_vec 移动两个字节。然后我们混合
x_vec 与移位向量使得x_vec 被移位向量替换，但仅在x_vec 的零位置。
以 4、8 和 16 个字节的移位重复此过程。您可以取消注释代码中的printf-s 以查看x_vec 和x_trail 之间发生的情况。
#include <stdio.h>
#include <x86intrin.h>
/* gcc -O3 -Wall -m64 -march=broadwell -falign-loops=16 horz_trail_max.c */
int print_vec_short(__m256i x);
__m256i hor_tr_max(__m256i x_vec)
__m256i zero = _mm256_setzero_si256();
__m256i pshufb_cnst = _mm256_set_epi64x(0x8080808080808080,0x8080808080808080,0x0F0E0F0E0F0E0F0E,0x0F0E0F0E0F0E0F0E);
__m256i mask1 = _mm256_cmpeq_epi16(x_vec,zero);
__m256i t1 = _mm256_slli_si256(x_vec,2); /* _mm256_slli_si256() doesn't cross the 128b lanes */
__m256i t2 = _mm256_blendv_epi8(x_vec,t1,mask1);
__m256i mask3 = _mm256_cmpeq_epi16(t2,zero);
__m256i t3 = _mm256_slli_si256(t2,4);
__m256i t4 = _mm256_blendv_epi8(t2,t3,mask3);
__m256i mask5 = _mm256_cmpeq_epi16(t4,zero);
__m256i t5 = _mm256_slli_si256(t4,8);
__m256i t6 = _mm256_blendv_epi8(t4,t5,mask5);
__m256i mask7 = _mm256_cmpeq_epi16(t6,zero);
__m256i t7_0 = _mm256_shuffle_epi8(t6,pshufb_cnst); /* _mm256_slli_si256() doesn't cross the 128b boundaries. Therefore we need a shuffle and a permute here. */
__m256i t7_1 = _mm256_permute2x128_si256(t7_0,t7_0,0x01); /* t7_1=t6[7], t6[7],...,t6[7], 0,0,0,0, 0,0,0,0 */
__m256i x_trail = _mm256_blendv_epi8(t6,t7_1,mask7);
/* Uncomment the next few lines to print the values of the intermediate variables */
/*
printf("\n15...0 = 15 14 13 12 11 10 9 8 7 6 5 4 3 2 1 0\n");
printf("x_vec = ");print_vec_short(x_vec );printf("mask1 = ");print_vec_short(mask1 );
printf("t1 = ");print_vec_short(t1 );printf("t2 = ");print_vec_short(t2 );
printf("mask3 = ");print_vec_short(mask3 );printf("t3 = ");print_vec_short(t3 );
printf("t4 = ");print_vec_short(t4 );printf("mask5 = ");print_vec_short(mask5 );
printf("t5 = ");print_vec_short(t5 );printf("t6 = ");print_vec_short(t6 );
printf("mask7 = ");print_vec_short(mask7 );printf("t7_0 = ");print_vec_short(t7_0 );
printf("t7_1 = ");print_vec_short(t7_1 );printf("x_trail = ");print_vec_short(x_trail );printf("\n");
*/
return x_trail;
int hor_tr_max_n(short int * x_in, short int * x_out, int n)
__m256i minus_1 = _mm256_set1_epi8(-1);
__m256i zero = _mm256_setzero_si256();
__m256i pshufb_cnst = _mm256_set_epi64x(0x8080808080808080,0x8080808080808080,0x0F0E0F0E0F0E0F0E,0x0F0E0F0E0F0E0F0E);
int indx_last_nz = 0;
for (int i=0;i<n;i=i+16)
__m256i x_vec = _mm256_load_si256((__m256i*)&x_in[i]);
__m256i mask1 = _mm256_cmpeq_epi16(x_vec,zero);
__m256i t1 = _mm256_slli_si256(x_vec,2);
__m256i t2 = _mm256_blendv_epi8(x_vec,t1,mask1);
__m256i mask3 = _mm256_cmpeq_epi16(t2,zero);
__m256i t3 = _mm256_slli_si256(t2,4);
__m256i t4 = _mm256_blendv_epi8(t2,t3,mask3);
__m256i mask5 = _mm256_cmpeq_epi16(t4,zero);
__m256i t5 = _mm256_slli_si256(t4,8);
__m256i t6 = _mm256_blendv_epi8(t4,t5,mask5);
__m256i mask7 = _mm256_cmpeq_epi16(t6,zero);
__m256i t7_0 = _mm256_shuffle_epi8(t6,pshufb_cnst);
__m256i t7_1 = _mm256_permute2x128_si256(t7_0,t7_0,0x01);
__m256i x_trail = _mm256_blendv_epi8(t6,t7_1,mask7);
__m256i isnonzero = _mm256_xor_si256(mask1,minus_1);
int mvmsk_nonz = _mm256_movemask_epi8(isnonzero);
int lz_x_vec = _lzcnt_u32( mvmsk_nonz ) >>1;
__m256i x_last_nz = _mm256_broadcastw_epi16(_mm_load_si128((__m128i*)&x_in[indx_last_nz]));
indx_last_nz = mvmsk_nonz ? (i+15-lz_x_vec) : indx_last_nz;
__m256i x_tr_is_zero = _mm256_cmpeq_epi16(x_trail,zero);
__m256i x_trail_upd = _mm256_blendv_epi8(x_trail,x_last_nz,x_tr_is_zero);
_mm256_store_si256((__m256i*)&x_out[i],x_trail_upd);
return 0;
int main()
#define test 0
#if test == 0
printf("Test 0: test functionality\n");
short x[16] = 1, 0, 0, 3, 0, 0, 4, 5, 0, 0, 0, 0, 4, 3, 0, 2;
// short x[16] = 0, 0, 0, 3, 0, 0, 4, 5, 0, 0, 0, 0, 4, 3, 0, 2;
// short x[16] = 1, 0, 0, 3, 0, 0, 4000, 0, 0, 0, 10, 0, 0, 3, 0, 2;
// short x[16] = 1100, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 5000, 0, 0, 0;
// short x[16] = 1100, 0, 0, 0, 0, 0, 0, 8888, 0, 0, 0, 0, 5000, 0, 0, 0;
printf("\n15...0 = 15 14 13 12 11 10 9 8 7 6 5 4 3 2 1 0\n");
__m256i x_vec = _mm256_loadu_si256((__m256i*)x);
printf("x_vec = ");print_vec_short(x_vec );
__m256i x_trail = hor_tr_max(x_vec);
printf("x_trail = ");print_vec_short(x_trail );
#elif test == 1 || test == 2
int i, i_o, k;
int n = 8000;
int d = 50;
short int *x_in;
short int *x_out;
x_in = _mm_malloc(n*sizeof(short int),32);
x_out = _mm_malloc(n*sizeof(short int),32);
int j = 73659343; /* Generate some a pseudo random array a. */
for (i = 0;i < n;i++)
j = j*653+1;
k = (j & 0x3FF00)>>8; /* k is a pseudo random number between 0 and 1023 */
if (k < d) /* with a small d, x_in has many zeros, try e.g. d=6, d=60 and d=600 */
x_in[i] = (j&0xFFE)+1-2048; /* Set x_in[i] to some nonzero. */
else
x_in[i] = 0;
#endif
#if test == 1
printf("Test 1: test performance for short int arrays of size n. Use: perf stat -d ./a.out \n");
for (i_o=0;i_o<400000;i_o++) /* The compiler should not interchange the inner and outer loop after function inlining, check compiler output (-S). */
hor_tr_max_n(x_in,x_out,n);
#elif test == 2
printf("Test 2: test performance of the unrolled scalar loop for short int arrays of size n. Use: perf stat -d ./a.out\n");
short int prev_x = 0;
for (i_o=0;i_o<400000;i_o++) /* The compiler should not interchange the inner and outer loop, check compiler output (-S). */
for (i=0;i<n;i=i+4)
short int x_in_i0 = x_in[i];
short int x_in_i1 = x_in[i+1];
short int x_in_i2 = x_in[i+2];
short int x_in_i3 = x_in[i+3];
prev_x = (x_in_i0)?(x_in_i0):(prev_x); x_out[i] = prev_x;
prev_x = (x_in_i1)?(x_in_i1):(prev_x); x_out[i+1] = prev_x;
prev_x = (x_in_i2)?(x_in_i2):(prev_x); x_out[i+2] = prev_x;
prev_x = (x_in_i3)?(x_in_i3):(prev_x); x_out[i+3] = prev_x;
#elif test == 3
printf("Test 3: Estimate approximately the latency and throughput of hor_tr_max with: perf stat -d ./a.out \n");
int i;
short x0[16] = 1, 0, 0, 3, 0, 0, 4, 5, 0, 0, 0, 0, 4, 3, 0, 2;
short x1[16] = 0, 0, 0, 3, 0, 12, 4, 5, 0, 0, 0, 0, 4, 3, 0, 2;
short x2[16] = 1, 0, 0, 3, 0, 0, 4, 5, 0, 0, 10, 0, 4, 3, 0, 2;
short x3[16] = 110, 0, 0, 1113, 0, 0, 4, 5, 0, 0, 0, 0, 4000, 3, 0, 2;
short x4[16] = 110, 4, 0, 1113, 0, 0, 4, 5, 0, 7, 0, 0, 4000, 3, 0, 2;
__m256i x_vec0 = _mm256_loadu_si256((__m256i*)x0); printf("x_vec0 = ");print_vec_short(x_vec0); __m256i x_trail0 = hor_tr_max(x_vec0);
__m256i x_vec1 = _mm256_loadu_si256((__m256i*)x1); printf("x_vec1 = ");print_vec_short(x_vec1); __m256i x_trail1 = hor_tr_max(x_vec1);
__m256i x_vec2 = _mm256_loadu_si256((__m256i*)x2); printf("x_vec2 = ");print_vec_short(x_vec2); __m256i x_trail2 = hor_tr_max(x_vec2);
__m256i x_vec3 = _mm256_loadu_si256((__m256i*)x3); printf("x_vec3 = ");print_vec_short(x_vec3); __m256i x_trail3 = hor_tr_max(x_vec3);
__m256i x_vec4 = _mm256_loadu_si256((__m256i*)x4); printf("x_vec4 = ");print_vec_short(x_vec4); __m256i x_trail4 = hor_tr_max(x_vec4);
for(i=0;i<100000000;i++)
x_trail0 = hor_tr_max(x_trail0); /* Use this line for latency testing, uncomment next 4 lines for throughput testing */
// x_trail1 = hor_tr_max(x_trail1);
// x_trail2 = hor_tr_max(x_trail2);
// x_trail3 = hor_tr_max(x_trail3);
// x_trail4 = hor_tr_max(x_trail4);
printf("x_trail0 = ");print_vec_short(x_trail0 );
printf("x_trail1 = ");print_vec_short(x_trail1 );
printf("x_trail2 = ");print_vec_short(x_trail2 );
printf("x_trail3 = ");print_vec_short(x_trail3 );
printf("x_trail4 = ");print_vec_short(x_trail4 );
#endif
#if test == 1 || test == 2
for (i=0;i<400;i++)
printf("%6i %6hi %6hi\n",i,x_in[i],x_out[i]);
#endif
return 0;
int print_vec_short(__m256i x)
short int v[16];
_mm256_storeu_si256((__m256i *)v,x);
printf("%4hi %4hi %4hi %4hi | %4hi %4hi %4hi %4hi | %4hi %4hi %4hi %4hi | %4hi %4hi %4hi %4hi\n",
v[15],v[14],v[13],v[12],v[11],v[10],v[9],v[8],v[7],v[6],v[5],v[4],v[3],v[2],v[1],v[0]);
return 0;
输出是：
15...0 = 15 14 13 12 11 10 9 8 7 6 5 4 3 2 1 0
x_vec = 2 0 3 4 | 0 0 0 0 | 5 4 0 0 | 3 0 0 1
x_trail = 2 3 3 4 | 5 5 5 5 | 5 4 3 3 | 3 1 1 1
此函数hor_tr_max 的延迟和吞吐量约为 14.2 和 6.4 个周期（英特尔 Skylake Core i5-6500）。
请注意，标准的稍微展开的标量循环，例如：
short int prev_x = 0;
for (i=0;i<n;i=i+4)
short int x_in_i0 = x_in[i];
short int x_in_i1 = x_in[i+1];
short int x_in_i2 = x_in[i+2];
short int x_in_i3 = x_in[i+3];
prev_x = (x_in_i0)?(x_in_i0):(prev_x); x_out[i] = prev_x;
prev_x = (x_in_i1)?(x_in_i1):(prev_x); x_out[i+1] = prev_x;
prev_x = (x_in_i2)?(x_in_i2):(prev_x); x_out[i+2] = prev_x;
prev_x = (x_in_i3)?(x_in_i3):(prev_x); x_out[i+3] = prev_x;
每个short int 大约需要 1.26 个周期，即每 16 个short int-s 需要 20.2 个周期。所以，向量化是
在这里有利可图。
大小为 n 的数组的水平尾随最大值
我们也可以使用hor_tr_max 来计算大小为n 的数组的水平尾随最大值，其中n 远大于16。
但是，需要步骤i 的输出来计算下一步。这种循环携带的依赖性导致代码性能低下。
函数hor_tr_max_n，在上面的代码中，实现了一个稍微不同的方法，使依赖链更短，这是有益的，因为
乱序调度。
函数 hor_tr_max_n 每 16 个 short ints 花费 12.2 个周期，这比展开的函数少 40%
标量循环。
使用即将推出的 Skylake-SP 处理器，“水平尾随最大值”的矢量化可能会 由于更宽的向量寄存器，利润更高。
【讨论】：
以上是关于AVX 或 SSE 上的水平尾随最大值的主要内容，如果未能解决你的问题，请参考以下文章