🦔

xsave

2021/12/12に公開

x86

これは AVX/AVX2/AVX512 Advent Calendar 2021(https://qiita.com/advent-calendar/2021/avx ) 向けに書きました。

(以下は調べながら書いたので、読みづらいです https://zenn.dev/tanakmura/articles/80391e3284c6bb ちゃんと調べてから書いた版があります)

AVXとレジスタのコンテキスト

AVX登場より前は、コンテキストスイッチのFPUレジスタ保存時に保存するレジスタセットにあわせて命令を駆使していた。

最終的にSSEレジスタ保存時には、fxsave/fxrstor という命令を使うようになっている(と思う)。

fxsave/fxrstor は、固定フォーマットで FPU, MMX, SSE の状態を保存、復元する。

(Intel® 64 and IA-32 Architectures Software Developer’s Manual Volume 1: Basic Architecture より)

fxsave/fxrstorは、レジスタの数やサイズが変わるごとに命令を追加しないといけないという問題を抱えている。

実際、これはAVXに対応していないし、さらにその上にAVX-512が乗ると、さらに命令追加しないといけない。さらに言うと、現時点で既にAMXというさらに追加の拡張が入ることも決まっている。

このへんの問題を整理し、FPUレジスタの保存、復元を担う命令が、xsave/xrstor である。

xsave/xrstor は、fxsave の動作に加えて、以下のように拡張されたレジスタを保存する。

何が保存されているかを示す XSAVE Header を保存する
保存するものは固定ではなく、XCR0 というレジスタで指定されたものを保存する

これによって、AVX、AVX-512、AMX (あと今は亡きMPX) のレジスタを保存できる命令になっている。

保存するものを指定するXCR0のビットはまだ余っていて、さらにその先の拡張があったときにもしばらくは対応できる。

(というほど簡単ではなく、sigaltstack を使うときのスタックサイズとしてよく使われていたSIGSTKSZ https://www.gnu.org/software/libc/manual/html_node/Signal-Stack.html は 8KiB しかなく、AMXのレジスタを保存できないのでプロセス起動後はAMXが使えないという問題が発生していたりする。Linux環境での話については次回(いつ？)書く)

xsave/xrstor

まずはxsaveを使ってみる。

xsave で保存すべき情報は、XCR0 というレジスタに入っていて、これを読めば、何を保存しないといけないかがわかる。

XCR0 は、xgetbvに0を指定すれば取れる

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

int main() {
    uint64_t xcr0 = _xgetbv(0);
    printf("xcr0 = %016llx\n", (long long)xcr0);
}

$ gcc -mavx xsave_xcr0.cpp  # -mavx がいる

# AVX-512が無い場合
$ ./a.out 
xcr0 = 0000000000000207

# AVX-512がある場合 (XCR0がAVX-512のレジスタを保存しないといけないことを示している)
$ ./a.out
xcr0 = 00000000000002e7

この、XCR0 の各ビットは、保存すべきレジスタセットを示していて、

bitpos	レジスタ
0	x87 (+MMX)
1	SSE
2	AVX
3	今は亡きMPX
4	今は亡きMPX
5	AVX-512 opmask (k0-k7のこと)
6	AVX-512 ZMM0-15 の上位
7	AVX-512 ZMM16-31
8	PT (Processor Trace の何か)
9	PKRU (知らん)
11	CET
12	CET
13	HDC (知らん)
16	HWp (知らん)
17	AMX の XTILECFG
18	AMX の XTILEDATA

となっている。

AVXよりあとの各領域のサイズは、CPUID で取ってこれる。 (x87 と SSE は固定フォーマットなので、マニュアルを見て調べる)

#include <cpuid.h>
#include <immintrin.h>
#include <stdio.h>
#include <string.h>

#include <string>
#include <vector>

struct ContextEntry {
    ContextEntry(int bitpos, std::string const &name)
        : bitpos(bitpos), name(name) {}

    int bitpos;
    std::string name;

    int bytepos = 0;
    int bytesize = 0;
    bool available = false;
};

std::vector<struct ContextEntry> offset_table = {
    {0, "x87"},
    {1, "SSE"},
    {2, "AVX"},
    {3, "MPX BNDREGS"},
    {4, "MPX BNDCSR"},
    {5, "AVX512 opmask"},
    {6, "AVX512 ZMM_Hi256"},
    {7, "AVX512 Hi16_ZMM"},
    {8, "PT"},
    {9, "PKRU"},
    // {10, 0, 0, "reserved"},
    {11, "CET_U"},
    {12, "CET_S"},
    {13, "HDC"},
    // {14, 0, "??"},
    // {15, 0, "reserved"},
    {16, "HWP"},

    /* amx */
    {17, "XTILECFG"},
    {18, "XTILEDATA"},
};

static int xsave_fullsize;

static void init_offset() {
    uint64_t xcr0 = _xgetbv(0);
    for (auto &&e : offset_table) {
        if (xcr0 & (1ULL << e.bitpos)) {
            unsigned int eax, ebx, ecx, edx;
            __cpuid_count(0x0d, e.bitpos, eax, ebx, ecx, edx);

            if (e.bitpos == 0) {
                // x87
                xsave_fullsize = ebx;
                e.bytepos = 0;
                e.bytesize = 160;
                e.available = true;
            } else if (e.bitpos == 1) {
                // sse
                e.bytepos = 160;
                e.bytesize = 16 * 16;
                e.available = true;

            } else {
                e.bytepos = ebx;
                e.bytesize = eax;
                e.available = true;
            }
        }
    }

    printf("xsave context size = %d\n", xsave_fullsize);

    for (auto &&e : offset_table) {
        if (e.available) {
            printf("%16s (%4d): offset=%d, size=%d\n", e.name.c_str(), e.bitpos,
                   e.bytepos, e.bytesize);
        }
    }
}

static void *alloc(size_t sz) {
    void *p;
    posix_memalign(&p, 64, sz);
    return p;
}

int main() {
    init_offset();

    char *area = (char*)alloc(xsave_fullsize);
    memset(area, 0, xsave_fullsize);

    /* 十分なサイズがある */
    _xsave((void *)area, ~0ULL);
}

# AVX-512 なし (alderlake)
$ ./a.out
xsave context size = 2696
             x87 (   0): offset=0, size=160
             SSE (   1): offset=160, size=256
             AVX (   2): offset=576, size=256
            PKRU (   9): offset=2688, size=8

# AVX-512 ほんとうになし (Kaby Lake)
$ ./a.out
xsave context size = 1088
             x87 (   0): offset=0, size=160
             SSE (   1): offset=160, size=256
             AVX (   2): offset=576, size=256
     MPX BNDREGS (   3): offset=960, size=64
      MPX BNDCSR (   4): offset=1024, size=64


# AVX-512 あり
$ ./a.out
xsave context size = 2696
             x87 (   0): offset=0, size=160
             SSE (   1): offset=160, size=256
             AVX (   2): offset=576, size=256
   AVX512 opmask (   5): offset=1088, size=64
AVX512 ZMM_Hi256 (   6): offset=1152, size=512
 AVX512 Hi16_ZMM (   7): offset=1664, size=1024
            PKRU (   9): offset=2688, size=8

これを見れば、固定でコンテキストサイズを決めなくてよくなるので、将来拡張が入った場合にも適切にコンテキストを保存できるようになる。

これを使ってみよう。gccなら_xsave intrinsicsを使えば、xsaveができる。

#include <cpuid.h>
#include <immintrin.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>

#include <string>
#include <vector>

struct ContextEntry {
    ContextEntry(int bitpos, std::string const &name)
        : bitpos(bitpos), name(name) {}

    int bitpos;
    std::string name;

    int bytepos = 0;
    int bytesize = 0;
    bool available = false;
};

std::vector<struct ContextEntry> offset_table = {
    {0, "x87"},
    {1, "SSE"},
    {2, "AVX"},
    {3, "MPX BNDREGS"},
    {4, "MPX BNDCSR"},
    {5, "AVX512 opmask"},
    {6, "AVX512 ZMM_Hi256"},
    {7, "AVX512 Hi6_ZMM"},
    {8, "PT"},
    {9, "PKRU"},
    // {10, 0, 0, "reserved"},
    {11, "CET_U"},
    {12, "CET_S"},
    {13, "HDC"},
    // {14, 0, "??"},
    // {15, 0, "reserved"},
    {16, "HWP"},

    /* amx */
    {17, "XTILECFG"},
    {18, "XTILEDATA"},
};

ContextEntry &lookup_entry(int bitpos) {
    for (auto &&e : offset_table) {
        if (e.bitpos == bitpos) {
            if (!e.available) {
                puts("not available");
                abort();
            }
            return e;
        }
    }

    abort();
}

static int xsave_fullsize;

static void init_offset() {
    uint64_t xcr0 = _xgetbv(0);
    for (auto &&e : offset_table) {
        if (xcr0 & (1ULL << e.bitpos)) {
            unsigned int eax, ebx, ecx, edx;
            __cpuid_count(0x0d, e.bitpos, eax, ebx, ecx, edx);

            if (e.bitpos == 0) {
                // x87
                xsave_fullsize = ebx;
                e.bytepos = 0;
                e.bytesize = 160;
                e.available = true;
            } else if (e.bitpos == 1) {
                // sse
                e.bytepos = 160;
                e.bytesize = 16 * 16;
                e.available = true;

            } else {
                e.bytepos = ebx;
                e.bytesize = eax;
                e.available = true;
            }
        }
    }

    printf("xsave context size = %d\n", xsave_fullsize);

    for (auto &&e : offset_table) {
        if (e.available) {
            printf("%16s (%4d): offset=%d, size=%d\n", e.name.c_str(), e.bitpos,
                   e.bytepos, e.bytesize);
        }
    }
}

#define SET_YMM(off, v) \
    asm volatile("vmovdqu %0, %%ymm" #off ::"x"(v) : "ymm" #off);

#define GET_YMM(off, v) \
    asm volatile("vmovdqu %%ymm" #off ",%0" :"=x"(v) :: "ymm" #off);

void set_avx_reg(__m256i init) {
    SET_YMM(1, init);
    SET_YMM(2, init);
    SET_YMM(3, init);
    SET_YMM(4, init);
    SET_YMM(5, init);
    SET_YMM(6, init);
    SET_YMM(7, init);
    SET_YMM(8, init);
    SET_YMM(9, init);
    SET_YMM(10, init);
    SET_YMM(11, init);
    SET_YMM(12, init);
    SET_YMM(13, init);
    SET_YMM(14, init);
    SET_YMM(15, init);
}

static void *alloc(size_t sz) {
    void *p;
    posix_memalign(&p, 64, sz);
    return p;
}

static void dump_xsave_header(char *p) {
    uint64_t *header = (uint64_t *)(p + 512);

    printf("XSTATE_BV : %016llx\n", (long long)header[0]);
    printf("XCOMP_BV  : %016llx\n", (long long)header[1]);
}

int main() {
    init_offset();

    char *area = (char *)alloc(xsave_fullsize);
    memset(area, 0, xsave_fullsize);

    set_avx_reg(_mm256_set1_epi8(0xff));
    SET_YMM(14, _mm256_set1_epi8(0xaa));
    SET_YMM(15, _mm256_set_epi32(1, 2, 3, 4, 5, 6, 7, 8));

    _xsave((void *)area, ~0ULL);

    {
        auto &avx_lo_context = lookup_entry(1);  // AVX lo (SSE)
        uint64_t *avx_lo_region = (uint64_t *)(area + avx_lo_context.bytepos);

        auto &avx_hi_context = lookup_entry(2);  // AVX hi
        uint64_t *avx_hi_region = (uint64_t *)(area + avx_hi_context.bytepos);

        for (int i = 0; i < 16; i++) {
            printf("ymm%-2d : %016llx %016llx %016llx %016llx\n", i,
                   (long long)avx_hi_region[i * 2 + 1],
                   (long long)avx_hi_region[i * 2 + 0],
                   (long long)avx_lo_region[i * 2 + 1],
                   (long long)avx_lo_region[i * 2 + 0]);
        }
    }

    /* avx reg clear */
    set_avx_reg(_mm256_set1_epi8(0x0));

    /* 復元 */
    _xrstor((void *)area, ~0ULL);

    __m256i v;
    GET_YMM(15, v);

    printf("%016llx %016llx %016llx %016llx\n",
           _mm256_extract_epi64(v, 3),
           _mm256_extract_epi64(v, 2),
           _mm256_extract_epi64(v, 1),
           _mm256_extract_epi64(v, 0));

    return 0;
}

cpuid で取り出した xsave の状態内のオフセットから値を取り出せば、avx のレジスタが取得できている。(ymm14 と ymm15 に設定した値が見えているはず。これはコンパイラの気持ちによって動作が変わるので、出ないかもしれない)

xsave が保存する値は、オペランドで決められる。_xsave の二番目の引数が、保存するレジスタセットを示していて、これと XCR0 の and を取ったものが保存される。

xrstor は、その逆の処理をする。

xsaveは保存するときに、何を保存したかをXSTATE_BVという領域に一緒に記録している。これは保存領域の先頭から512byteに書かれていて、フォーマットは、XCR0 と同じになる。

xrstor は、このXSTATE_BVを見て復元すべきレジスタを決める。XSTATE_BVに記録されていないレジスタは初期化される。

xsaveopt

AVX が有効になっていると、XCR0 は、AVX がコンテキストに含まれるように設定される。

AVX はまだいいが、これが AVX-512 にもなると、コンテキストのサイズが2KiB を超えるようになって、毎回全部保存するのは、効率がよくない。可能なら、使ったレジスタだけを保存してほしい。

x86 は、使われたレジスタセットを、XINUSEという状態に保存していて、これのビットが立っているかどうかを見て、保存/復元すべきレジスタセットを自動的に削減させる、xsaveopt という命令を用意している。

XINUSEの状態は、xgetbv の 1番から取得できる

#include <cpuid.h>
#include <immintrin.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>

#include <string>
#include <vector>

#define SET_YMM(off, v) \
    asm volatile("vmovdqu %0, %%ymm" #off ::"x"(v) : "ymm" #off);

#define GET_YMM(off, v) \
    asm volatile("vmovdqu %%ymm" #off ",%0" :"=x"(v) :: "ymm" #off);

static void
dump_xcr0() {
    uint64_t xcr0 = _xgetbv(0);
    uint64_t xinuse = _xgetbv(1);

    printf("XCR0   : %016llx\n", xcr0);
    printf("XINUSE : %016llx\n", xinuse);
}

int main()
{
    dump_xcr0();

    SET_YMM(1, _mm256_set1_epi8(1));

    dump_xcr0();
}

$ ./a.out 
XCR0   : 0000000000000207
XINUSE : 0000000000000202
XCR0   : 0000000000000207
XINUSE : 0000000000000206 # AVXが使われた

xsaveopt は、これを見て保存すべきコンテキストを決め、どれを保存したかをレジスタと一緒にメモリに保存する。これは、XSAVE headerと呼ばれて、xsave が保存した状態の

#include <cpuid.h>
#include <immintrin.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>

#include <string>
#include <vector>

struct ContextEntry {
    ContextEntry(int bitpos, std::string const &name)
        : bitpos(bitpos), name(name) {}

    int bitpos;
    std::string name;

    int bytepos = 0;
    int bytesize = 0;
    bool available = false;
};

std::vector<struct ContextEntry> offset_table = {
    {0, "x87"},
    {1, "SSE"},
    {2, "AVX"},
    {3, "MPX BNDREGS"},
    {4, "MPX BNDCSR"},
    {5, "AVX512 opmask"},
    {6, "AVX512 ZMM_Hi256"},
    {7, "AVX512 Hi6_ZMM"},
    {8, "PT"},
    {9, "PKRU"},
    // {10, 0, 0, "reserved"},
    {11, "CET_U"},
    {12, "CET_S"},
    {13, "HDC"},
    // {14, 0, "??"},
    // {15, 0, "reserved"},
    {16, "HWP"},

    /* amx */
    {17, "XTILECFG"},
    {18, "XTILEDATA"},
};

ContextEntry &lookup_entry(int bitpos) {
    for (auto &&e : offset_table) {
        if (e.bitpos == bitpos) {
            if (!e.available) {
                puts("not available");
                abort();
            }
            return e;
        }
    }

    abort();
}

static int xsave_fullsize;

static void init_offset() {
    uint64_t xcr0 = _xgetbv(0);
    for (auto &&e : offset_table) {
        if (xcr0 & (1ULL << e.bitpos)) {
            unsigned int eax, ebx, ecx, edx;
            __cpuid_count(0x0d, e.bitpos, eax, ebx, ecx, edx);

            if (e.bitpos == 0) {
                // x87
                xsave_fullsize = ebx;
                e.bytepos = 0;
                e.bytesize = 160;
                e.available = true;
            } else if (e.bitpos == 1) {
                // sse
                e.bytepos = 160;
                e.bytesize = 16 * 16;
                e.available = true;

            } else {
                e.bytepos = ebx;
                e.bytesize = eax;
                e.available = true;
            }
        }
    }

    printf("xsave context size = %d\n", xsave_fullsize);

    for (auto &&e : offset_table) {
        if (e.available) {
            printf("%16s (%4d): offset=%d, size=%d\n", e.name.c_str(), e.bitpos,
                   e.bytepos, e.bytesize);
        }
    }
}

#define SET_YMM(off, v) \
    asm volatile("vmovdqu %0, %%ymm" #off ::"x"(v) : "ymm" #off);

#define GET_YMM(off, v) \
    asm volatile("vmovdqu %%ymm" #off ",%0" :"=x"(v) :: "ymm" #off);

static void *alloc(size_t sz) {
    void *p;
    posix_memalign(&p, 64, sz);
    return p;
}

static void dump_xsave_header(char *p) {
    uint64_t *header = (uint64_t *)(p + 512);

    printf("XSTATE_BV : %016llx\n", (long long)header[0]);
    //printf("XCOMP_BV  : %016llx\n", (long long)header[1]);
}

int main() {
    init_offset();

    char *area = (char *)alloc(xsave_fullsize);
    memset(area, 0, xsave_fullsize);

    _xsave((void *)area, ~0ULL);
    dump_xsave_header(area);

    SET_YMM(0, _mm256_set1_epi8(1));

    _xsave((void *)area, ~0ULL);
    dump_xsave_header(area);

    return 0;
}

参考

Intel SDM Vol1 Chapter 13 が一番参考になると思います

GitHubで編集を提案

AVXとレジスタのコンテキスト

xsave/xrstor

xsaveopt

参考

Discussion