web-dev-qa-db-ja.com

AVX2マスクに基づいて左にパックする最も効率的な方法は何ですか?

入力配列と出力配列があり、特定の条件を渡す要素のみを書き込みたい場合、AVX2でこれを行う最も効率的な方法は何ですか?

私はSSEのように行われた場所で見ました:(From: https://deplinenoise.files.wordpress.com/2015/03/gdc2015_afredriksson_simd.pdf =)

__m128i LeftPack_SSSE3(__m128 mask, __m128 val)
{
 // Move 4 sign bits of mask to 4-bit integer value.
 int mask = _mm_movemask_ps(mask);
 // Select shuffle control data
 __m128i shuf_ctrl = _mm_load_si128(&shufmasks[mask]);
 // Permute to move valid values to front of SIMD register
 __m128i packed = _mm_shuffle_epi8(_mm_castps_si128(val), shuf_ctrl);
 return packed;
}

これは、SSEは4幅なので16エントリのLUTで十分ですが、AVXが8幅の場合、LUTは非常に大きくなります(256エントリ、各32バイト、または8k)。

AVXに、パッキング付きのマスクされたストアなど、このプロセスを簡略化するための指示がないように思われるのには驚きです。

左に設定された符号ビットの数をカウントするために少しシャッフルすると、必要な置換テーブルを生成してから、_mm256_permutevar8x32_psを呼び出すことができると思います。しかし、これは私が思うかなりの数の指示でもあります。

AVX2でこれを行うためのトリックを知っている人はいますか?または最も効率的な方法は何ですか?

上記のドキュメントの左パッキング問題の図は次のとおりです。

Left.Packing.Problem

ありがとう

31
Froglegs

AVX2 + BMI2。 AVX512に関する他の回答を参照してください。 (更新:64ビットのビルドでpdepを保存しました。)

AVX2 vpermps_mm256_permutevar8x32_ps (または同等の整数、vpermd)を使用して、車線横断の変数シャッフルを実行できます。

BMI2 pext(Parallel Bits Extract) は操作のビット単位のバージョンを提供するため、オンザフライでマスクを生成できます必要です。

pdep/pextvery AMD CPUでは遅い、たとえば6 uops/18サイクルのレイテンシとRyzenでのスループットであることに注意してください。この実装はAMDではひどく実行されます。 AMDの場合、pshufbまたはvpermilps LUTを使用する128ビットのベクトル、またはマスク入力がベクトルマスク(メモリからの計算済みのビットマスクではない)の場合、コメントで説明されているAVX2変数シフトの提案のいくつかが最適です。いずれにしても、Zen2より前のAMDには128ビットのベクトル実行ユニットしかなく、256ビットのレーン横断シャッフルは低速です。したがって、現在のAMDでは、128ビットのベクトルが非常に魅力的です。


32ビット以上の要素を持つ整数ベクトルの場合:いずれか1)_mm256_movemask_ps(_mm256_castsi256_ps(compare_mask))
または2)_mm256_movemask_epi8を使用し、最初のPDEP定数を0x0101010101010101から0x0F0F0F0F0F0F0F0Fに変更して、連続する4ビットのブロックを分散します。 0xFFUによる乗算をexpanded_mask |= expanded_mask<<4;またはexpanded_mask *= 0x11;に変更します(テストされていません)。どちらの方法でも、VPERMPSではなくVPERMDでシャッフルマスクを使用します。

64ビット整数またはdouble要素の場合、すべてが正常に機能します;比較マスクは常に同じである32ビット要素のペアを持っているため、結果のシャッフルは各64ビット要素の両方の半分を正しい場所に配置します。 (したがって、VPERMPDおよびVPERMQは即時制御オペランドでのみ使用できるため、引き続きVPERMPSまたはVPERMDを使用します。)

16ビット要素の場合、128ビットのベクトルを使用してこれを適応できる場合があります。


アルゴリズム:

パックされた3ビットインデックスの定数から始め、各位置は独自のインデックスを保持します。つまり、[ 7 6 5 4 3 2 1 0 ]各要素の幅は3ビットです。 0b111'110'101'...'010'001'000

pextを使用して、必要なインデックスを整数レジスタの下部にある連続したシーケンスに抽出します。例えばインデックス0と2が必要な場合、pextのコントロールマスクは0b000'...'111'000'111である必要があります。 pextは、セレクターの1ビットと一致する010および000インデックスグループを取得します。選択されたグループは出力の下位ビットにパックされるため、出力は0b000'...'010'000になります。 (つまり、[ ... 2 0 ]

入力ベクトルマスクからpext0b111000111入力を生成する方法については、コメント付きのコードを参照してください。

これで、compressed-LUTと同じボートに入りました。最大8つのパックされたインデックスをアンパックします。

すべてのピースをまとめると、合計で3つのpext/pdepsになります。自分がやりたかったこととは逆に取り組んだので、その方向で理解するのもおそらく最も簡単でしょう。 (つまり、シャッフルラインから開始し、そこから逆方向に作業します)。

パックされた3ビットのグループではなく、バイトごとに1つのインデックスを処理する場合、アンパックを簡略化できます。 8つのインデックスがあるので、これは64ビットコードでのみ可能です。

これとGodbolt Compiler Explorerの32ビットのみのバージョン を参照してください。 #ifdefsを使用したので、-m64または-m32で最適にコンパイルされます。 gccはいくつかの命令を浪費しますが、clangは本当に素晴らしいコードを作成します。

#include <stdint.h>
#include <immintrin.h>

// Uses 64bit pdep / pext to save a step in unpacking.
__m256 compress256(__m256 src, unsigned int mask /* from movmskps */)
{
  uint64_t expanded_mask = _pdep_u64(mask, 0x0101010101010101);  // unpack each bit to a byte
  expanded_mask *= 0xFF;    // mask |= mask<<1 | mask<<2 | ... | mask<<7;
  // ABC... -> AAAAAAAABBBBBBBBCCCCCCCC...: replicate each bit to fill its byte

  const uint64_t identity_indices = 0x0706050403020100;    // the identity shuffle for vpermps, packed to one index per byte
  uint64_t wanted_indices = _pext_u64(identity_indices, expanded_mask);

  __m128i bytevec = _mm_cvtsi64_si128(wanted_indices);
  __m256i shufmask = _mm256_cvtepu8_epi32(bytevec);

  return _mm256_permutevar8x32_ps(src, shufmask);
}

これは、メモリからのロードなしでコードにコンパイルされ、即時定数のみです。 (これと32ビットバージョンのgodboltリンクを参照してください)。

    # clang 3.7.1 -std=gnu++14 -O3 -march=haswell
    mov     eax, edi                   # just to zero extend: goes away when inlining
    movabs  rcx, 72340172838076673     # The constants are hoisted after inlining into a loop
    pdep    rax, rax, rcx              # ABC       -> 0000000A0000000B....
    imul    rax, rax, 255              # 0000000A0000000B.. -> AAAAAAAABBBBBBBB..
    movabs  rcx, 506097522914230528
    pext    rax, rcx, rax
    vmovq   xmm1, rax
    vpmovzxbd       ymm1, xmm1         # 3c latency since this is lane-crossing
    vpermps ymm0, ymm1, ymm0
    ret

したがって、 Agner Fogの数値 によると、これは6 uopsです(定数、またはインライン化すると消えるゼロ拡張movは数えません)。 Intel Haswellでは、待ち時間は16cです(vmovqに1つ、各pdep/imul/pext/vpmovzx/vpermpsに3つ)。命令レベルの並列処理はありません。ただし、これがループキャリー依存関係の一部ではないループでは(ゴッドボルトリンクに含めたように)、ボトルネックはスループットだけであり、一度に複数の反復を実行し続けることが期待されます。

これはおそらく、pdep/pext/imulのポート1でボトルネックとなっている3サイクルに1つのスループットを管理できます。もちろん、ロード/ストアとループのオーバーヘッド(compare、movmsk、popcntを含む)では、uopの合計スループットが問題になる可能性があります。 (たとえば、私のgodboltリンクのフィルターループは、読みやすくするために-fno-unroll-loopsを使用して、clangで14 uopsです。運が良ければ、4cごとに1回の反復を維持し、フロントエンドに追いつくことができます。しかし、clangは、popcntの出力への誤った依存関係を説明できなかったので、compress256関数のレイテンシの3/5でボトルネックになります。)

gccは、8の左シフトとsubを使用して、複数の命令で0xFFを乗算します。これには追加のmov命令が必要ですが、最終結果はレイテンシが2の乗算になります(Haswellはレジスタ名変更ステージでレイテンシがゼロのmovを処理します)。


AVX2をサポートするすべてのハードウェアがBMI2もサポートしているため、BMI2なしのAVX2のバージョンを提供しても意味がありません。

非常に長いループでこれを行う必要がある場合、LUTエントリをアンパックするだけのオーバーヘッドが少ない十分な反復で初期キャッシュミスが償却される場合、LUTはおそらく価値があります。まだmovmskpsが必要なので、マスクをポップしてLUTインデックスとして使用できますが、pdep/imul/pexpは保存します。

LUTエントリは、私が使用したものと同じ整数シーケンスでアンパックできますが、@ Froglegsのset1()/vpsrlvd/vpandは、LUTエントリがメモリで開始され、最初に整数レジスタに入る必要がない場合におそらくより優れています。 。 (32ビットのブロードキャストロードには、Intel CPUのALU uopは必要ありません)。ただし、可変シフトはHaswellでは3 uopsです(Skylakeでは1つのみ)。

25
Peter Cordes

AMD Zenをターゲットにしている場合は、ryzenでのpdepand pextが非常に遅いため(この場合、それぞれ18サイクル)、この方法が推奨されます。

私はこの方法を思いつきました。これは8kではなく768(+1パディング)バイトの圧縮LUTを使用します。単一のスカラー値のブロードキャストが必要で、各レーンで異なる量だけシフトされ、下位3ビットにマスクされ、0〜7のLUTを提供します。

これは、LUTをビルドするためのコードとともに、組み込みバージョンです。

//Generate Move mask via: _mm256_movemask_ps(_mm256_castsi256_ps(mask)); etc
__m256i MoveMaskToIndices(u32 moveMask) {
    u8 *adr = g_pack_left_table_u8x3 + moveMask * 3;
    __m256i indices = _mm256_set1_epi32(*reinterpret_cast<u32*>(adr));//lower 24 bits has our LUT

   // __m256i m = _mm256_sllv_epi32(indices, _mm256_setr_epi32(29, 26, 23, 20, 17, 14, 11, 8));

    //now shift it right to get 3 bits at bottom
    //__m256i shufmask = _mm256_srli_epi32(m, 29);

    //Simplified version suggested by wim
    //shift each lane so desired 3 bits are a bottom
    //There is leftover data in the lane, but _mm256_permutevar8x32_ps  only examines the first 3 bits so this is ok
    __m256i shufmask = _mm256_srlv_epi32 (indices, _mm256_setr_epi32(0, 3, 6, 9, 12, 15, 18, 21));
    return shufmask;
}

u32 get_nth_bits(int a) {
    u32 out = 0;
    int c = 0;
    for (int i = 0; i < 8; ++i) {
        auto set = (a >> i) & 1;
        if (set) {
            out |= (i << (c * 3));
            c++;
        }
    }
    return out;
}
u8 g_pack_left_table_u8x3[256 * 3 + 1];

void BuildPackMask() {
    for (int i = 0; i < 256; ++i) {
        *reinterpret_cast<u32*>(&g_pack_left_table_u8x3[i * 3]) = get_nth_bits(i);
    }
}

MSVCによって生成されたアセンブリは次のとおりです。

  lea ecx, DWORD PTR [rcx+rcx*2]
  lea rax, OFFSET FLAT:unsigned char * g_pack_left_table_u8x3 ; g_pack_left_table_u8x3
  vpbroadcastd ymm0, DWORD PTR [rcx+rax]
  vpsrlvd ymm0, ymm0, YMMWORD PTR __ymm@00000015000000120000000f0000000c00000009000000060000000300000000
7
Froglegs

LUTなしのAVX2 + BMI2については、他の回答を参照してください。

AVX512へのスケーラビリティに関する懸念について言及しているので、心配しないでくださいこれには、AVX512F命令があります

VCOMPRESSPS —スパースパックされた単精度浮動小数点値を高密度メモリに保存 。 (double、32または64ビット整数要素(vpcompressq)のバージョンもありますが、バイトまたはワード(16ビット)はありません)。これはBMI2 pdep/pextに似ていますが、整数regのビットではなくベクトル要素用です。

デスティネーションはベクトルレジスタまたはメモリオペランドにすることができますが、ソースはベクトルおよびマスクレジスタです。レジスタdestを使用すると、上位ビットをマージまたはゼロ化できます。メモリーdestでは、「隣接するベクターのみが宛先メモリー位置に書き込まれます」。

次のベクトルに向けてポインタをどれだけ進めるかを判断するには、マスクをポップします。

配列から0以上の値以外をすべて除外したいとします。

#include <stdint.h>
#include <immintrin.h>
size_t filter_non_negative(float *__restrict__ dst, const float *__restrict__ src, size_t len) {
    const float *endp = src+len;
    float *dst_start = dst;
    do {
        __m512      sv  = _mm512_loadu_ps(src);
        __mmask16 keep = _mm512_cmp_ps_mask(sv, _mm512_setzero_ps(), _CMP_GE_OQ);  // true for src >= 0.0, false for unordered and src < 0.0
        _mm512_mask_compressstoreu_ps(dst, keep, sv);   // clang is missing this intrinsic, which can't be emulated with a separate store

        src += 16;
        dst += _mm_popcnt_u64(keep);   // popcnt_u64 instead of u32 helps gcc avoid a wasted movsx, but is potentially slower on some CPUs
    } while (src < endp);
    return dst - dst_start;
}

これは、(gcc4.9以降で)( Godbolt Compiler Explorer )にコンパイルされます。

 # Output from gcc6.1, with -O3 -march=haswell -mavx512f.  Same with other gcc versions
    lea     rcx, [rsi+rdx*4]             # endp
    mov     rax, rdi
    vpxord  zmm1, zmm1, zmm1             # vpxor  xmm1, xmm1,xmm1 would save a byte, using VEX instead of EVEX
.L2:
    vmovups zmm0, ZMMWORD PTR [rsi]
    add     rsi, 64
    vcmpps  k1, zmm0, zmm1, 29           # AVX512 compares have mask regs as a destination
    kmovw   edx, k1                      # There are some insns to add/or/and mask regs, but not popcnt
    movzx   edx, dx                      # gcc is dumb and doesn't know that kmovw already zero-extends to fill the destination.
    vcompressps     ZMMWORD PTR [rax]{k1}, zmm0
    popcnt  rdx, rdx
    ## movsx   rdx, edx         # with _popcnt_u32, gcc is dumb.  No casting can get gcc to do anything but sign-extend.  You'd expect (unsigned) would mov to zero-extend, but no.
    lea     rax, [rax+rdx*4]             # dst += ...
    cmp     rcx, rsi
    ja      .L2

    sub     rax, rdi
    sar     rax, 2                       # address math -> element count
    ret

パフォーマンス:Skylake-X/Cascade Lakeでは256ビットのベクターの方が高速かもしれません

理論的には、ビットマップを読み込み、1つの配列を別の配列にフィルターするループは、ベクター幅に関係なく、SKX/CSLXで3クロックごとに1つのベクターで実行され、ポート5でボトルネックになります(kmovb/w/d/q k1, eaxはp5で実行され、vcompressps IACAおよび http://uops.info/ によるテストによると、メモリへの格納は2p5 +ストアです。

@ZachBは、実際には、ZMM _mm512_mask_compressstoreu_psを使用したループが実際のCSLXハードウェアの_mm256_mask_compressstoreu_psよりも少し遅いことをコメントで報告します。(256ビットバージョンが「512ビットベクトルモード」から抜け出し、クロックが高くなるようにするマイクロベンチマークなのか、それとも周囲の512ビットコードがあったのかはわかりません。)

ミスアラインメントされたストアが512ビットバージョンに悪影響を与えているのではないかと思います。 vcompresspsは、マスクされた256ビットまたは512ビットのベクトルストアを効果的に実行します。これがキャッシュラインの境界を超える場合、追加の作業を行う必要があります。通常、出力ポインターは16要素の倍数ではないため、フルラインの512ビットストアはほとんどの場合ミスアライメントされます。

512ビットのストアが正しく整列されていない場合、キャッシュライン分割された256ビットのストアよりも、何らかの理由で頻繁に発生するだけでなく、パフォーマンスが低下する可能性があります。他のものの512ビットのベクトル化は、位置合わせに敏感であるように思われます。毎回スプリットロードバッファが不足することが原因である可能性があります。または、キャッシュラインスプリットを処理するフォールバックメカニズムは、512ビットのベクトルでは効率が悪い可能性があります。

vcompresspsをレジスターにベンチマークし、完全なベクトルが重複する個別のストアを使用すると興味深いでしょう。それはおそらく同じuopsですが、別の命令の場合、ストアはマイクロヒューズを使用できます。そして、マスクされたストアとオーバーラップしているストアの間に何らかの違いがある場合、これはそれを明らかにします。


以下のコメントで説明されている別のアイデアは、vpermt2psを使用して、整列した店舗の完全なベクトルを構築することでした。これは ブランチレスで行うのは難しいでしょう であり、ビットマスクにかなり規則的なパターンがないか、all-0とall-1の大きなランがない限り、ベクターを埋めるときの分岐はおそらく誤った予測になります。

vpermt2psと、「フル」のときにそれを置き換えるブレンドまたは何かを使用して、構築中のベクトルを4サイクルまたは6サイクルのループキャリー依存チェーンでブランチなしに実装できる可能性があります。整列されたベクトルストアでは、すべての反復を格納しますが、ベクトルがいっぱいの場合にのみ出力ポインターを移動します。

これは、現在のIntel CPUにアラインされていないストアを持つvcompresspsよりも遅い可能性があります。

7
Peter Cordes

誰かがここで興味がある場合は、ジャンプテーブルという別名のデータLUTの代わりに命令LUTを使用するSSE2のソリューションがあります。ただし、AVXでは256ケースが必要です。

以下のLeftPack_SSE2を呼び出すたびに、jmp、shufps、jmpの3つの命令が使用されます。 16ケースのうち5ケースは、ベクトルを変更する必要はありません。

static inline __m128 LeftPack_SSE2(__m128 val, int mask)  {
  switch(mask) {
  case  0:
  case  1: return val;
  case  2: return _mm_shuffle_ps(val,val,0x01);
  case  3: return val;
  case  4: return _mm_shuffle_ps(val,val,0x02);
  case  5: return _mm_shuffle_ps(val,val,0x08);
  case  6: return _mm_shuffle_ps(val,val,0x09);
  case  7: return val;
  case  8: return _mm_shuffle_ps(val,val,0x03);
  case  9: return _mm_shuffle_ps(val,val,0x0c);
  case 10: return _mm_shuffle_ps(val,val,0x0d);
  case 11: return _mm_shuffle_ps(val,val,0x34);
  case 12: return _mm_shuffle_ps(val,val,0x0e);
  case 13: return _mm_shuffle_ps(val,val,0x38);
  case 14: return _mm_shuffle_ps(val,val,0x39);
  case 15: return val;
  }
}

__m128 foo(__m128 val, __m128 maskv) {
  int mask = _mm_movemask_ps(maskv);
  return LeftPack_SSE2(val, mask);
}
6
Z boson