はじめに
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に配列からデータをロードして、別のところに書き込むサンプル。
#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
次、マスク付きでやってみる。コードはこんな感じになるだろう。
#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
)を作っている。1
さて、コンパイルして実行する。
$ 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と配列の向きが逆なので混乱するので、ここではわかりやすさを優先した。 ↩