C++のメモ

AVXのベクトル演算機能を用いて高速化を試す．

行列の積を並列に行うことを考えると，答えの1行目は
a00 * (b00, b01, b02, b03) + a01 * ( b10, b11, b12, b13) + ... + a03 * ( b30, b31, b32, b33)
と表せられれる．

要素ごとの積をMatlabに習って.*と表すと

(c00, c01, c02, c03) = (a00, a00, a00, a00) .* (b00, b01, b02, b03) +
(a01, a01, a01, a01) .* (b10, b11, b12, b13) +
(a02, a02, a02, a02) .* (b20, b21, b22, b23) +
(a03, a03, a03, a03) .* (b30, b31, b32, b33)

である．これをaの各行ごとに繰り返すことで答えが得られる．
この要素ごとの積・和はAVX命令にあるのでこれで高速化する．

• _mm256_add_pd: a .* b のベクトル演算
• _mm256_mul_pd: a + b のベクトル演算

#include <intrin.h>
void prod_avx(double const* a, double const* b, double* c)
{
  __declspec(align(32)) double c_[16];
  __m256d b0n, b1n, b2n, b3n;
  b0n = _mm256_load_pd(b);
  b1n = _mm256_load_pd(b + 4);
  b2n = _mm256_load_pd(b + 8);
  b3n = _mm256_load_pd(b + 12);

  for (int i = 0; i < 4; ++i)
  {
    __m256d ai0, ai1, ai2, ai3;
    ai0 = _mm256_broadcast_sd(a + 4 * i);
    ai1 = _mm256_broadcast_sd(a + 4 * i + 1);
    ai2 = _mm256_broadcast_sd(a + 4 * i + 2);
    ai3 = _mm256_broadcast_sd(a + 4 * i + 3);

    ai0 = _mm256_mul_pd(ai0, b0n); 
    ai1 = _mm256_mul_pd(ai1, b1n); 
    ai2 = _mm256_mul_pd(ai2, b2n); 
    ai3 = _mm256_mul_pd(ai3, b3n); 

    ai0 = _mm256_add_pd(ai0, ai1);
    ai2 = _mm256_add_pd(ai2, ai3);
    ai0 = _mm256_add_pd(ai0, ai2);
    _mm256_store_pd(c_ + 4 * i, ai0);
  }
  memcpy(c, c_, 16 * sizeof(double));
}

_mm256_fmadd_pd(a .* b + cのベクトル演算)を用いることでさらに高速化できる．

#include <intrin.h>
void prod_avx2(double const* a, double const* b, double* c)
{
  __declspec(align(32)) double c_[16];
  __m256d b0n, b1n, b2n, b3n;
  b0n = _mm256_load_pd(b);
  b1n = _mm256_load_pd(b + 4);
  b2n = _mm256_load_pd(b + 8);
  b3n = _mm256_load_pd(b + 12);

  for (int i = 0; i < 4; ++i)
  {
    __m256d ai0, ai1, ai2, ai3;
    ai0 = _mm256_broadcast_sd(a + 4 * i);
    ai1 = _mm256_broadcast_sd(a + 4 * i + 1);
    ai2 = _mm256_broadcast_sd(a + 4 * i + 2);
    ai3 = _mm256_broadcast_sd(a + 4 * i + 3);

    ai0 = _mm256_mul_pd(ai0, b0n); 
    ai0 = _mm256_fmadd_pd(ai1, b1n, ai0); // ai0 = ai1 .* b1n + ai0
    ai0 = _mm256_fmadd_pd(ai2, b2n, ai0); 
    ai0 = _mm256_fmadd_pd(ai3, b3n, ai0); 
    _mm256_store_pd(c_ + 4 * i, ai0);
  }
  memcpy(c, c_, 16 * sizeof(double));
}

250msec間の処理回数を計測して算出

Core i7-4800MQ

Normal	0.142 [usec]
AVX	0.091 [usec]
AVX2	0.056 [usec]

AVXにおいて3倍の高速化が得られた．
(といっても0.1usecだが)

ベクトルの要素を入れ替えるshuffle, permuteという命令を繰り返すことで，
ルービックキューブを解くような感じで転置させることができる．

 r0[0] r0[1] r0[2] r0[3]      r0[0] r1[0] r2[0] r3[0]
 r1[0] r1[1] r1[2] r1[3]  ->  r0[1] r1[1] r2[1] r3[1]
 r2[0] r2[1] r2[2] r2[3]      r0[2] r1[2] r2[2] r3[2]
 r3[0] r3[1] r3[2] r3[3]      r0[3] r1[3] r2[3] r3[3]

r0[1]とr1[0]，r0[3]とr1[2]を入れ替える．
同様に，r2[1]とr3[0]，r2[3]とr3[2]を入れ替える．

tmp0 = (r0[0] r1[0] r0[2] r1[2])
tmp1 = (r0[1] r1[1] r0[3] r1[3])
tmp2 = (r2[0] r3[0] r2[2] r3[2])
tmp3 = (r2[1] r3[1] r2[3] r3[3])

この操作は_mm256_shuffle_pdで一回で可能である．

dst[63:0] := (imm[0] == 0) ? a[63:0] : a[127:64] 
dst[127:64] := (imm[1] == 0) ? b[63:0] : b[127:64] 
dst[191:128] := (imm[2] == 0) ? a[191:128] : a[255:192] 
dst[255:192] := (imm[3] == 0) ? b[191:128] : b[255:192]

imm[0]は1bitを表し，上位・下位128bitのどちらをもってくるかを選択できる．

  __mm256d tmp0 = _mm256_shuffle_pd(r0, r1, 0); // b0000
  // -> tmp0 = (r0[0], r1[0], r0[2], r1[2])
  __mm256d tmp1 = _mm256_shuffle_pd(r0, r1, 15); // b1111
  // -> tmp1 = (r0[1], r1[1], r0[3], r1[3])

tmp0の下位128bit(double x2)とtmp2の下位128bitから転置後の1行目が得られる．
tmp0の上位128bitとtmp2の上位128bitから転置後の3行目が得られる．
この操作は_mm256_permute2f128_pdで可能である．

SELECT4(src1, src2, control)
{ 
  CASE(control[1:0]) 
    0: tmp[127:0] := src1[127:0] 
    1: tmp[127:0] := src1[255:128] 
    2: tmp[127:0] := src2[127:0] 
    3: tmp[127:0] := src2[255:128] 
  ESAC 
  IF control[3] 
    tmp[127:0] := 0 
  FI 
  RETURN tmp[127:0] 
} 

dst[127:0] := SELECT4(a[255:0], b[255:0], imm[3:0]) 
dst[255:128] := SELECT4(a[255:0], b[255:0], imm[7:4]) 
dst[MAX:256] := 0

imm[3:0]という表記はimmの0から3までの4bitを表す．

  // shuffleによる結果
  // -> tmp0 = (r0[0], r1[0], r0[2], r1[2])
  // -> tmp2 = (r2[0], r3[0], r2[2], r3[2])
  // t0 = (r0[0], r1[0], r2[0], r3[0])
  // t2 = (r0[2], r1[2], r2[2], r3[2]) を作る
  __m256d t0 = _mm256_permute2f128_pd(tmp0, tmp2, 0x20); // 0010 0000
  /*
    (t0[0] t0[1]) = SELECT4(tmp0, tmp2, 0)
                  = (tmp0[0] tmp0[1])
                  = (r0[0] r1[0])
    (t0[2] t0[3]) = SELECT4(tmp0, tmp2, 2)
                  = (tmp2[0] tmp2[1])
                  = (r2[0] r3[0])
  */
  
  __m256d t2 = _mm256_permute2f128_pd(tmp0, tmp2, 0x31); // 0011 0001
  /*
    (t0[0] t0[1]) = SELECT4(tmp0, tmp2, 1)
                  = (tmp0[2] tmp0[3])
                  = (r0[2] r1[2])
    (t0[2] t0[3]) = SELECT4(tmp0, tmp2, 3)
                  = (tmp2[2] tmp2[3])
                  = (r2[2] r3[2])
  */

1,000,000,000回転置するのに必要な時間を計測し，一回あたりの時間を求めた．
SIMDを使わない実装は下記．

for(int i = 0; i < 4; ++i)
{
  for(int j = 0; j < 4; ++j)
  {
    b[4*j+i] = a[4*i+j];
  }
}

Normal	0.0047 usec
AVX	0.0025 usec

一応，倍速である．