任意長の配列処理

以前紹介した、一連の「一次元配列同士の演算」などを任意長を扱えるように拡張します。
arques.hatenablog.com
機能は、配列同士の各要素を演算し、別の配列へ格納するするプログラムです。処理概要を図で示します。

ベクトルで処理する場合、SIMDレジスターが処理できる要素数と、実際の要素数に相関があります。例えば、SIMDレジスターであるzmmレジスターが保持できる要素数がα、実際のデータ長をβとすると、βはαの整数倍でなければなりません。
このルールに外れ、任意長を処理する方法を本章で説明します。例えば、zmmレジスターを利用し64ビット浮動小数点を配列の演算を行う例を考えてみましょう。zmmレジスターは64ビット浮動小数点を8要素処理できます。処理する配列の要素数が259の例を考えてみましょう。

スカラーで処理

まず、スカラーで処理する例を示します。

スカラーですので、ひとつずつ処理するため何も問題ありません。

ベクトルとスカラーの組み合わせ

もっともオーソドックスと思われるベクトル処理とスカラー処理を組み合わせる方法を考えてみましょう。

ベクトル処理できる範囲を8要素単位でベクトル処理し、余りの8要素に満たない部分を「スカラーで処理」と同じ方法で処理します。

マスクレジスターを用いたベクトル処理

すべてをベクトル処理しますが、最後の部分をマスクレジスターで処理する方法を紹介します。

すべてを8要素単位でベクトル処理します。ただし、最後の8要素に満たない部分はマスクレジスターを使用します。マスクレジスターを使用することによって、処理対象外をアクセスしないようにします。

64ビット整数型

呼び出し側・符号付

以降に、呼び出し側のC++コードを示します。一次元配列から指定された値のインデックスと、その数を取得します。
common.hをincludeしていますがSIMDで記述したアセンブリ言語の関数と等価な機能を持つC++言語で記述した共通関数のヘッダです。最下部にソースコードを示します。条件は関数名に含まれます。

#include "..\common.h"  // TEMPLATES

// asmbler関数名: 処理+型
#define T long long
extern "C"
{
    void addq(const T*, const T*, T*, const size_t);
    void subq(const T*, const T*, T*, const size_t);
    void mulq(const T*, const T*, T*, const size_t);
}
typedef void (*Dfunc)(const T*, const T*, T*, const size_t);

extern "C"
Dfunc afunc[] = { addq, subq, mulq };
Dfunc cfunc[] = { cadd, csub, cmul };

// main
int main(void)
{
    const size_t ArrLenReal = 4098;
    T a[ArrLenReal], b[ArrLenReal], c[ArrLenReal], v[ArrLenReal];

    for (size_t ArrLen = 0; ArrLen < ArrLenReal - 1; ArrLen++)
    {
        init(a, b, ArrLen);
        for (int i = 0; i < sizeof(cfunc) / sizeof(*cfunc);i++)
        {
            cfunc[i](a, b, c, ArrLen);
            afunc[i](a, b, v, ArrLen);
            verify(c, v, ArrLen);
        }
    }
    return 0;
}

素数を変化させるだけで、以前紹介したものと同様ですので説明は省きます。

呼び出され側

任意長の64ビット整数型配列同士を演算し、結果を出力用の配列に格納する関数を紹介します。64ビット整数型を処理する関数を示します。
SIMD命令を使用しベクトル処理すると高速化できますが、要素長がSIMD命令のレジスター長に揃っていないと面倒が起きます。一般的には、最後の要素をベクトル処理できない場合、その部分はスカラー処理します。AVX-512命令セットはマスクレジスターをサポートしているため、データ長がレジスター長に揃っていない場合でもベクトル命令を使用できます。ここでは、そのようなマスクレジスターを使用する例を紹介します。
本関数へは、入力配列 a のアドレスが rcx レジスターへ、入力配列 b のアドレスが rdx レジスターへ渡されます。これらの配列要素を演算した結果を格納する格納用配列 c のアドレスが r8 レジスターへ渡されます。配列の要素数はr9 レジスターへ渡されます。渡される内容とレジスターの対応を表で示します。

レジスタ 内容
rcx 配列 a の先頭アドレス
rdx 配列 b の先頭アドレス
r8 配列 c の先頭アドレス
r9 配列の要素数

以降に、64ビット整数型配列の要素を演算する関数をマクロで記述したものを示します。

mymacro macro       MNAME, MINST

        public      MNAME
        align       16
MNAME   proc
        mov         r10, rcx                ; r10 = a[]
        mov         rcx, r9                 ; rcx = length

        xor         rax, rax                ; clrar remainder mask
        and         rcx, 111b               ; rcx = remainder, 0-7
        jz          begin                   ; no remainder

        mov         r11, -1                 ; make remainder mask
        shld        rax, r11, cl            ; rax = remainder mask

begin:
        mov         rcx, r9                 ; rcx = length
        xor         r9, r9                  ; clear index
        sar         rcx, 3                  ; length/=8
        jz          do_remainder

loop_f:
        vmovdqu64   zmm0, [r10+r9]          ; load a[]
        MINST       zmm0, zmm0, [rdx+r9]    ; c[] = a[] <op> b[]
        vmovdqu64   [r8+r9], zmm0           ; stote c[]

        lea         r9, qword ptr[r9+64]    ; next offset
        dec         rcx
        jnz         loop_f

do_remainder:
        or          rax, rax                ; test remainder mask
        jz          exit_f                  ; no remainder

        kmovq       k1, rax
        vmovdqu64   zmm0{k1}, [r10+r9]      ; load a[]
        MINST       zmm0{k1}, zmm0, [rdx+r9]; c[] = a[] <op> b[]
        vmovdqu64   [r8+r9]{k1}, zmm0       ; stote c[]

exit_f:
        ret
MNAME   endp

        endm

;-------------------------------------------------------------------
; code
_TEXT   segment

mymacro addq, vpaddq
mymacro subq, vpsubq
mymacro mulq, vpmuldq

_TEXT   ends
        end

mymacro マクロを定義し、これを使ってaddq、subq、mulq関数を作ります。それぞれ符号付64ビット整数配列同士の加算、減算、乗算を行う関数です。

分かり易くするため、addqが展開されたソースリストを順次示し、その説明を行います。まず、zmmレジスターで処理できない余りの処理を行うためのマスク値をraxレジスターへ求める先頭部を示します。

        public      addq
        align       16
addq    proc
        mov         r10, rcx                ; r10 = a[]
        mov         rcx, r9                 ; rcx = length

        xor         rax, rax                ; clrar remainder mask
        and         rcx, 111b               ; rcx = remainder, 0-7
        jz          begin                   ; no remainder

        mov         r11, -1                 ; make remainder mask
        shld        rax, r11, cl            ; rax = remainder mask

begin:

r10レジスターへ配列 a の先頭アドレスを保存します。次に、r10レジスターへ配列の要素数を設定します。そして、raxを0へ初期化します。raxレジスターは、配列の最後に残る余りの部分を処理するためのマスク値を保持します。要素数をrcxへ移動したのち、下位の3ビットを調べzmmレジスターで処理できない余りの要素数をrcxへ求めます。余りがない(=全要素をベクトル処理できる)場合は、ベクトル処理を行うためにbeginラベルまでジャンプします。そうでないときは、これをshld 命令などを利用し、マスク値をraxレジスターへ求めます。マスク値を64ビットで求めますが、64ビット整数型配列の余りは0~7ですので、1バイトで十分です。64ビットで処理しているのは、32ビット型や8ビット型などのすべてのデータ型と共通コードとしたためです。
beginラベルの手前で、r10 レジスターに配列aの先頭アドレス、rdx レジスターに配列bの先頭アドレス、r8 レジスターに配列cの先頭アドレス、r9 レジスターに要素数、eax レジスターに余りを処理するためのマスク値が求まります。

        mov         rcx, r9                 ; rcx = length
        xor         r9, r9                  ; clear index
        sar         rcx, 3                  ; length/=8
        jz          do_remainder

配列のインデックスを管理するr9レジスターをクリアし、初期化します。次に、ベクトル処理のループ回数をrcx レジスターへ求めます。もし、一回もループ処理が必要ない場合は、do_remainderラベル以降に制御を移します。

loop_f:
        vmovdqu64   zmm0, [r10+r9]          ; load a[]
        vpaddq      zmm0, zmm0, [rdx+r9]    ; c[] = a[] <op> b[]
        vmovdqu64   [r8+r9], zmm0           ; stote c[]

        lea         r9, qword ptr[r9+64]    ; next offset
        dec         rcx
        jnz         loop_f

ベクトル処理で配列の演算を行います。ループ内の処理は「配列同士の演算」で紹介したものと同様です。

do_remainder:
        or          rax, rax                ; test remainder mask
        jz          exit_f                  ; no remainder

        kmovq       k1, rax
        vmovdqu64   zmm0{k1}, [r10+r9]      ; load a[]
        vpaddq      zmm0{k1}, zmm0, [rdx+r9]; c[] = a[] <op> b[]
        vmovdqu64   [r8+r9]{k1}, zmm0       ; stote c[]

exit_f:
        ret
addq    endp

do_remainderラベル以降は、ベクトル処理できなかった余りの処理です。raxレジスターに含まれるマスク値をテストし、余りがない場合は終了に向かいます。余りの処理が必要なら、raxレジスターの値をk1マスクレジスターへ移し、vmovdqu64命令やvpaddq命令へk1マスクレジスターを指定し、余りの分だけ処理します。raxレジスターの値をk1マスクレジスターへ移す際にkmovq命令を使用していますが、マスクレジスターの使用される範囲は、続く命令によって変わります。この例なら、マスクは8ビットで十分ですが、データ型に依存しないように最大幅で処理します。
このようにAVX-512では、データ長がSIMDレジスター長に揃っていなくても、ベクトル命令で全データを処理できます。

なお、

do_remainder:
        or          rax, rax                ; test remainder mask
        jz          exit_f                  ; no remainder

        kmovq       k1, rax
        vmovdqu64   zmm0{k1}, [r10+r9]      ; load a[]
        vpaddq      zmm0{k1}, zmm0, [rdx+r9]; c[] = a[] <op> b[]
        vmovdqu64   [r8+r9]{k1}, zmm0       ; stote c[]

exit_f:
        ret
addq    endp

の部分は;

    :
do_remainder:
        kmovq       k1, rax
        vmovdqu64   zmm0{k1}, [r10+r9]      ; load a[]
        vpaddq      zmm0{k1}, zmm0, [rdx+r9]; c[] = a[] <op> b[]
        vmovdqu64   [r8+r9]{k1}, zmm0       ; stote c[]

        ret
addq    endp

と書き換えても良いでしょう。
余りがなくても(マスクレジスターの値が0の場合)、以降の命令でメモリーアクセスや演算は行われません。このため余りの処理が必要か不要か判断する必要はありません。
ここでは、raxレジスターに含まれるマスク値のテストは一回しか実行されず速度への影響が少ないと考え、余りの有無をチェックします。

common.h、共通関数

ヘッダーファイルの一部を示します。簡単な内容なので、説明は省きします。

    :
// add
template <typename T>
void cadd(const T* a, const T* b, T* c, const size_t length)
{
    for (size_t i = 0; i < length; i++)
    {
        c[i] = a[i] + b[i];
    }
}

// sub
template <typename T>
void csub(const T* a, const T* b, T* c, const size_t length)
{
    for (size_t i = 0; i < length; i++)
    {
        c[i] = a[i] - b[i];
    }
}

// mul
template <typename T>
void cmul(const T* a, const T* b, T* c, const size_t length)
{
    for (size_t i = 0; i < length; i++)
    {
        c[i] = a[i] * b[i];
    }
}
    :

すべての型に対応させるため、テンプレート関数とします。cadd 関数は、2つの一次元配列の各要素を加算します。csub 関数は、2つの一次元配列の各要素を減算します。最後の、cmul 関数は、2つの一次元配列の各要素を乗算します。
このプログラムのビルドし、実行させるとエラーが発生したときのみメッセージが表示され、正常なら何も表示されません。


説明は行いませんが、参考のため8ビット整数型配列の要素を演算する関数をマクロで記述したものを示します。

mymacro macro       MNAME, MINST

        public      MNAME
        align       16
MNAME   proc
        mov         r10, rcx                ; r10 = a[]
        mov         rcx, r9                 ; rcx = length

        xor         rax, rax                ; clrar remainder mask
        and         rcx, 111111b            ; rcx = remainder, 0-63
        jz          begin                   ; no remainder

        mov         r11, -1                 ; make remainder mask
        shld        rax, r11, cl            ; rax = remainder mask

begin:
        mov         rcx, r9                 ; rcx = length
        xor         r9, r9                  ; clear index
        sar         rcx, 6                  ; length/=64
        jz          do_remainder

loop_f:
        vmovdqu8    zmm0, [r10+r9]          ; load a[]
        MINST       zmm0, zmm0, [rdx+r9]    ; c[] = a[] <op> b[]
        vmovdqu8    [r8+r9], zmm0           ; stote c[]

        lea         r9, qword ptr[r9+64]    ; next offset
        dec         rcx
        jnz         loop_f

do_remainder:
        or          rax, rax                ; test remainder mask
        jz          exit_f                  ; no remainder

        kmovq       k1, rax
        vmovdqu8    zmm0{k1}, [r10+r9]      ; load a[]
        MINST       zmm0{k1}, zmm0, [rdx+r9]; c[] = a[] <op> b[]
        vmovdqu8    [r8+r9]{k1}, zmm0       ; stote c[]

exit_f:
        ret
MNAME   endp

        endm

;-------------------------------------------------------------------
; code
_TEXT   segment

mymacro addb,  vpaddb
mymacro subb,  vpsubb

_TEXT   ends
        end