LJの力計算をフルアセンブラで書いてみる

はじめに

Lennard-Jones粒子の力の計算をフルアセンブラで書いてみてインテルコンパイラにケンカを売ってみる。

思想

SIMD化はループ方向ではなく、ループ内で行う。データをAoSで持っておき、3次元ベクトルの掛け算やロードだけSIMDで、肝心の力計算はSIMD化しない。これでどこまでいけるか見てみる。アセンブラ慣れてないのでアホなことをやってる可能性大。

コード

書いたコードはこんな感じ。繰り返すが書き方が正しい自信が無い。calcがSIMD化したもの、checkが素直に書いた二重ループ。2万粒子のO(N^2)計算。カットオフ無し、初期配置は手抜きで一次元的。-DMYSIMDつけてコンパイルすると手でかいたSIMDコードを使う。

//------------------------------------------------------------------------
#include <stdio.h>
//------------------------------------------------------------------------
enum {X, Y, Z};
const int N = 20000;
const double dt = 0.01;
double q[N][4] = {};
double p[N][4] = {};
double c24[4] = {24 * dt, 24 * dt, 24 * dt, 0};
double c48[4] = { -48 * dt, -48 * dt, -48 * dt, 0};
//------------------------------------------------------------------------
static void
calc(void) {
  printf("calc\n");
  __asm__ (
    "xor %r10, %r10\n\t" // r10 = i = 0
    "mov $q, %ecx\n\t"
    "mov $p, %edx\n\t"
    "I.LOOP:\n\t"
    "vmovapd (%ecx), %ymm7\n\t" // qi
    "vmovapd (%edx), %ymm6\n\t" // pi
    "mov %ecx, %ebx\n\t"
    "mov %edx, %eax\n\t"
    "mov %r10, %r11\n\t" // r11 = j = i
    "inc %r11\n\t"
    "J.LOOP:\n\t"
    "add $32,%ebx\n\t"
    "add $32,%eax\n\t"
    "vmovapd (%ebx), %ymm1\n\t" // qj
    "vmovapd c24(%rip), %ymm4\n\t" // ymm4 = (24,24,24,0)
    "vmovapd c48(%rip), %ymm5\n\t" // ymm4 = (48,48,48,0)
    "vsubpd %ymm1, %ymm7, %ymm0\n\t" // ymm0 = (dx, dy, dz,0 )
    "vmulpd %ymm0, %ymm0, %ymm1\n\t" // ymm1 = (dx**2, dy**2, dz**2,0)
    "vpermpd $201, %ymm1, %ymm2\n\t" // ymm2 = (dy**2, dz**2, dx**2,0)
    "vpermpd $210, %ymm1, %ymm3\n\t" // ymm3 = (dz**2, dx**2, dy**2,0)
    "vaddpd %ymm1, %ymm2, %ymm1\n\t" // ymm1 = (xy,yz,zx,0)
    "vaddpd %ymm1, %ymm3, %ymm1\n\t" // ymm1 = (r2,r2,r2,0)
    "vmulpd %ymm1, %ymm1, %ymm2\n\t" // ymm2 = (r4,r4,r4,0)
    "vmulpd %ymm1, %ymm2, %ymm2\n\t" // ymm2 = (r6,r6,r6,0)
    "vmulpd %ymm2, %ymm2, %ymm3\n\t" // ymm3 = (r12,r12,r12,0)
    "vmulpd %ymm1, %ymm3, %ymm3\n\t" // ymm3 = (r14,r14,r14,0)
    "vfmadd132pd %ymm4, %ymm5, %ymm2\n\t" // ymm2 = (24.0*r6-48)*dt
    "vdivpd %ymm3, %ymm2, %ymm1\n\t" //ymm1 = (24.0*r6-48)/r14*dt =df
    "vmulpd %ymm1, %ymm0, %ymm0\n\t" // ymm0 = (df*dx, df*dy, df*dz, 0)
    "vmovapd (%eax), %ymm5\n\t" // pj -> ymm5
    "vaddpd %ymm5, %ymm0, %ymm5\n\t"// pj -= df*dt;
    "vsubpd %ymm0, %ymm6, %ymm6\n\t"// pi += df*dt;
    "vmovapd %ymm5, (%eax)\n\t" // ymm5 -> pj
    "inc %r11\n\t" //j++
    "cmp $20000,%r11\n\t"
    "jl J.LOOP\n\t"
    "vmovapd %ymm6, (%edx)\n\t" // ymm6 -> pi
    "add $32,%ecx\n\t"
    "add $32,%edx\n\t"
    "inc %r10\n\t" //i++
    "cmp $19999,%r10\n\t"
    "jl I.LOOP\n\t"
  );
}
//------------------------------------------------------------------------
void
init(void) {
  for (int i = 0; i < N; i++) {
    q[i][X] = 1.0 + 0.4 * i;
    q[i][Y] = 2.0 + 0.5 * i;
    q[i][Z] = 3.0 + 0.6 * i;
    p[i][X] = 0.0;
    p[i][Y] = 0.0;
    p[i][Z] = 0.0;
  }
}
//------------------------------------------------------------------------
void
check(void) {
  printf("check\n");
  for (int i = 0; i < N; i++) {
    for (int j = i + 1; j < N; j++) {
      const double dx = q[j][X] - q[i][X];
      const double dy = q[j][Y] - q[i][Y];
      const double dz = q[j][Z] - q[i][Z];
      const double r2 = (dx * dx + dy * dy + dz * dz);
      double r6 = r2 * r2 * r2;
      double df = (24.0 * r6 - 48.0) / (r6 * r6 * r2) * dt;
      p[i][X] += df * dx;
      p[i][Y] += df * dy;
      p[i][Z] += df * dz;
      p[j][X] -= df * dx;
      p[j][Y] -= df * dy;
      p[j][Z] -= df * dz;
    }
  }
}
//------------------------------------------------------------------------
int
main(void) {
  init();
#ifdef MYSIMD
  calc();
#else
  check();
#endif

  for (int i = 0; i < 10; i++) {
    printf("%f %f %f\n", p[i][X], p[i][Y], p[i][Z]);
  }
}
//------------------------------------------------------------------------

結果

計算環境

Intel(R) Xeon(R) CPU E5-2680 v3 @ 2.50GHz
icpc (ICC) 16.0.3
g++ (GCC) 4.8.5

結果

1	2	3
コンパイルオプション	速度 [s]	備考
g++ -O3 -DMYSIMD	2.015	自分で書いた奴
icpc -O3 -xHOST	0.942	インテルコンパイラに任せた奴
g++ -O3	2.054	g++に任せた奴

一応GCCにはほんのすこしだけ勝ったが、インテルコンパイラには倍近い差をつけられた。

まとめ

インテルコンパイラにケンカを売ったが返り討ちにあった。GCCにはかろうじて勝ったが、元コードを工夫すればもっと性能出るだろうから、こっちにも実質的には負けているものと思われる。でもまだナイーブに組んだだけなので工夫の余地はあるし、実コードではif分岐や間接参照が入るから、この方向でもう少しいけるかも？