はじめに

16個のintを、マスクに従って選択的にストアしたい。

つまり、

int in[16] = {1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16};
int mask[16] = {1,1,0,0,1,0,1,0,0,0,1,0,1,0,0,0};

みたいなデータから、

int out[16] = {1,2,0,0,5,0,7,0,0,0,11,0,13,0,0,0}

を作りたい。これをAVX-512でintを16個マスク付きでストアで実現しようとして、ちょっとだけハマったのでメモ。

_mm512_store_epi32

まずはマスク無しでデータのコピーができるか試す。単にzmmに配列からデータをロードして、別のところに書き込むサンプル。

store.cpp

# include <immintrin.h>
# include <stdio.h>
int out[16] = {};
int in[16] = {1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16};
void
dump(int a[16]){
  for(int i=0;i<16;i++){
    printf("%02d ",a[i]);
  }
  printf("\n");
}
int
main(void){
  __m512i vin = _mm512_loadu_si512(in);
  _mm512_store_epi32(out,vin);
  dump(in);
  dump(out);
}

$ icpc store.cpp
$ ./a.out
01 02 03 04 05 06 07 08 09 10 11 12 13 14 15 16 
01 02 03 04 05 06 07 08 09 10 11 12 13 14 15 16

問題なくできた。

_mm512_mask_store_epi32

次、マスク付きでやってみる。コードはこんな感じになるだろう。

maskstore.cpp

# include <immintrin.h>
# include <stdio.h>
int out[16] = {};
int in[16] = {1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16};
int mask[16] = {1,1,0,0,1,0,1,0,0,0,1,0,1,0,0,0};
void
dump(int a[16]){
  for(int i=0;i<16;i++){
    printf("%02d ",a[i]);
  }
  printf("\n");
}
int
main(void){
  __m512i vin = _mm512_loadu_si512(in);
  __m512i vmask = _mm512_loadu_si512(mask);
  __mmask16 vm = _mm512_test_epi32_mask(vmask,vmask);
  _mm512_mask_store_epi32(out,vm,vin);
  dump(in);
  dump(mask);
  dump(out);
}

一度zmmにマスクデータをロードし、vptestmdでビットに落とすことでマスクレジスタ(__mmask16 vm)を作っている。¹

さて、コンパイルして実行する。

$ icpc maskstore.cpp
$ ./a.out
Segmentation fault

ありゃ、SIGSEGVった。

原因は、_mm512_test_epi32_mask、というかvmovdqa32が64バイトアラインを要求するから。出力先を

__attribute__((aligned(64))) int out[16] = {};

と64バイトアラインしてやると、

$ ./a.out
01 02 03 04 05 06 07 08 09 10 11 12 13 14 15 16 
01 01 00 00 01 00 01 00 00 00 01 00 01 00 00 00 
01 02 00 00 05 00 07 00 00 00 11 00 13 00 00 00

と、所望の動作になる。

原因というか言い訳というか

インテルのマニュアルには、_mm512_store_epi32も_mm512_test_epi32_maskも、対応する命令はvmovdqa32で、それは64バイトアラインを要求する、と書いてあるように読める。

で、何も指定しないで_mm512_store_epi32がうまく行ったから、「あ、コンパイラが勝手に64バイトアラインしてくれたんだ」と思い込んでて、_mm512_test_epi32_maskを使ったらSIGSEGVって、別の原因を探していた。

で、アセンブリ見ると、_mm512_store_epi32はvmovdqa32ではなくvmovupsを呼んでいる。store.cppのアセンブリはこんな感じ。

        vmovups   in(%rip), %zmm0
        vmovups   %zmm0, out(%rip)

vmovupsは32バイトアラインで良くて、なんか(古いのは知らんけど少なくとも最近の)インテルコンパイラは勝手に32バイトアラインで配列を取ってくれるので、問題なく動作してた。で、もちろん_mm512_test_epi32_maskは、vmovdqa32を呼ぶ。こんな感じ。

        vmovups   mask(%rip), %zmm0
        vmovups   in(%rip), %zmm1
        vptestmd  %zmm0, %zmm0, %k1
        vmovdqa32 %zmm1, out(%rip){%k1}

vmovdqa32は格納先が64バイトアラインされていることを要求するのでSIGSEGVる。

まとめ

マニュアルには_mm512_store_epi32も_mm512_test_epi32_maskも対応する命令がvmovdqa32と書いてあったけど、筆者が試した環境では前者はvmovupsが呼ばれてる。混乱無く使うためには、AVX-512使う際には64バイトアラインする癖つけといたほうが良さそう。

どうでもいいけど、マスクレジスタ便利っぽい。いろいろ使えそう。

二進表現マスク作るといつもMSBと配列の向きが逆なので混乱するので、ここではわかりやすさを優先した。 ↩