徒然なる日々を送るソフトウェアデベロッパーの記録（２）

技術上思ったことや感じたことを気ままに記録していくブログです。さくらから移設しました。

惑星の生成をシミュレーションする

C++ CUDA 物理学

CUDA のサンプルプログラムに nbody というのが付属していますが、
ちょっと改造して、微惑星から惑星に成長させるプログラムを書いてみました。

まず、太陽系が太陽と微惑星から成り立っており、太陽が微惑星に比べて非常に
重い（微惑星を全て足しあわせても 2.3 倍以上）とします。運動方程式は
{\displaystyle \frac{d^2 \vec{x_i}}{dt^2} = \sum_j{\frac{Gm_j(\vec{x_j} - \vec{x_i})}{|\vec{x_j} - \vec{x_i}|^3 + \epsilon^3}}-\frac{GM \vec{x_i}}{|\vec{x_i}|^3 + \epsilon^3} }
ただし、太陽は非常に重いので、座標系を太陽を中心にとり、重心系と
ほぼ一致するものとします。\epsilon は太陽または微惑星同士が非常に
接近した場合に発散しないように設けた非常に小さい定数です。
距離を天文単位（AU）で測り、時間を年、質量を太陽質量単位で測るものとすると、
GM=39.2 となります。

微惑星同士が非常に近づいた場合（1e-05 AU 以下）は、完全非弾性衝突させます。
つまり、衝突前の質量および衝突後の質量をm_i, m_j, m_{ij}, 衝突前および
衝突後の速度をv_i, v_j, v_{ij}とすると、
{\displaystyle m_i v_i + m_j v_j = m_{ij} v_{ij} }
となります。（本来は相対速度によってくっついたり、弾いたり、粉々に
なったりするでしょう）

本来は質量、速度の入れ替えにアトミック操作が必要ですが、CUDA のスピンロック
条件が思った以上に難しかったので、アトミック操作を取っていません。

初期条件として、\vec{x_i} は 0.5 AU〜1.0AU に存在するものとし、
微惑星はほぼ定常状態にあるものとします。
すると Kepler の第３法則から角速度\omegaを
{\displaystyle \omega = \sqrt{\frac{GM}{|\vec{x_i(0)}|^3}} }
で求めることができます。
このままだと微惑星同士が衝突しないので、実際には位置と速度にわずかに擾乱を加えます。

メインプログラムは以下のようになりました。

#include <GL/freeglut.h>
#include <GL/gl.h>
#include <GL/glext.h>
#include <cuda_runtime.h>
#include <cuda_gl_interop.h>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <math.h>
#include "CosmicCloud.h"

using namespace std;

const float rmin = 0.5;
const float rmax = 1.0;
const float gm = 39.2;
const int ncloud = 65536;
const float alpha = 0.3;
const float dt = 0.001;

float2 *devPos[2] = { NULL, NULL };
float2 *devVel = NULL;
float *mass = NULL;
int *lockVal = NULL;
int flip = 0;
float mag = 0.5f;
GLuint pbo[2];
cudaGraphicsResource *gres[2];

void init(int *argc, char **argv);
void calc(void);
void update(void);
void show(void);
void getMinMaxDist(float *);

static void updateScale(void) {
  glMatrixMode(GL_PROJECTION);
  glLoadIdentity();
  glOrtho(-36.0 * mag, 36.0 * mag, -24.0 * mag, 24.0 * mag, -1.0, 1.0);
}

static void keyboard(unsigned char key, int, int) {
  switch (key) {
  case 27:
    exit(0);
    break;
  case '+':
    mag /= 10.0f;
    updateScale();
    break;
  case '-':
    mag *= 10.0f;
    updateScale();
    break;
  }
}

static void display(void) {
  show();
}

static void idle(void) {
  calc();
  update();
  glutPostRedisplay();
}

int main(int argc, char **argv) {
  init(&argc, argv);
  glutDisplayFunc(display);
  glutKeyboardFunc(keyboard);
  glutIdleFunc(idle);
  glutMainLoop();
  return 0;
}

// 初期値を生成する
void createInitDist(float2 *vel, float2 *pos, float *m) {
  int i;
  float gmi = gm * alpha / ncloud;
  for (i = 0; i < ncloud; ++i) {
    float r = rmin + (float)random() / (float)RAND_MAX * (rmax - rmin);
    float theta = (float)random() / (float)RAND_MAX * 2. * M_PI;
    float omega = sqrt(gm / (r * r * r));
    float r1 = (float)random() / (float)RAND_MAX * 0.02 - 0.01;
    float r2 = (float)random() / (float)RAND_MAX * 0.02 - 0.01;
    pos[i].x = r * cos(theta);
    pos[i].y = r * sin(theta);
    vel[i].x = -omega * pos[i].y + r1;
    vel[i].y = omega * pos[i].x + r2;
    mass[i] = gmi;
  } 
}

void init(int *argc, char **argv) {
  float2 *hostPos = (float2 *)calloc(sizeof(float2), ncloud);
  float2 *hostVel = (float2 *)calloc(sizeof(float2), ncloud);

  // TODO: OpenGL の初期化
  cudaGLSetGLDevice(0);
  glutInit(argc, argv);
  glutInitDisplayMode(GLUT_RGB | GLUT_DOUBLE);
  glutInitWindowSize(720, 480);
  glutCreateWindow("Planets Evolusion Simulation");
  glClearColor(0.0, 0.0, 0.0, 1.0);
  glMatrixMode(GL_MODELVIEW);
  glViewport(-36.0 * mag, -24.0 * mag, 36.0 * mag, 24.0 * mag);
  glMatrixMode(GL_PROJECTION);
  glLoadIdentity();
  glOrtho(-36.0 * mag, 36.0 * mag, -24.0 * mag, 24.0 * mag, -1.0, 1.0);
  glGenBuffers(2, (GLuint *)&pbo[0]);

  // デバイスメモリ領域の確保
//  cudaMalloc((void **)&devPos[0], sizeof(float2) * ncloud);
//  cudaMalloc((void **)&devPos[1], sizeof(float2) * ncloud);
  cudaMalloc((void **)&devVel, sizeof(float2) * ncloud);
  cudaMalloc((void **)&lockVal, sizeof(int) * ncloud);
  cudaMallocManaged((void **)&mass, sizeof(float) * ncloud);
  
  createInitDist(hostVel, hostPos, mass);
  cudaMemset(lockVal, -1, sizeof(int) * ncloud);
  cudaMemcpy(devVel, hostVel, sizeof(float2) * ncloud, cudaMemcpyDefault);

  glBindBuffer(GL_ARRAY_BUFFER, pbo[0]);
  glBufferData(GL_ARRAY_BUFFER, ncloud * sizeof(float2), hostPos, GL_DYNAMIC_DRAW);
  glBindBuffer(GL_ARRAY_BUFFER, pbo[1]);
  glBufferData(GL_ARRAY_BUFFER, ncloud * sizeof(float2), hostPos, GL_DYNAMIC_DRAW);
  glBindBuffer(GL_ARRAY_BUFFER, 0);

  // TODO: OpenGL のリソースを CUDA に登録する
  cudaGraphicsGLRegisterBuffer(&gres[0], pbo[0], cudaGraphicsMapFlagsNone);
  cudaGraphicsGLRegisterBuffer(&gres[1], pbo[1], cudaGraphicsMapFlagsNone);

  // 初期値の開放
  free(hostPos);
  free(hostVel);
}

void calc(void) {
  // TODO: リソースをメモリ空間にマップする
  float2 *devPtr[2] = { NULL, NULL };
  size_t sizes;
//  cudaGraphicsResourceSetMapFlags(gres[flip], cudaGraphicsMapFlagsReadOnly);
//  cudaGraphicsResourceSetMapFlags(gres[1 - flip], cudaGraphicsMapFlagsNone);
  cudaGraphicsMapResources(2, &gres[0], 0);
  cudaGraphicsResourceGetMappedPointer((void **)&devPtr[0], &sizes, gres[0]);
  cudaGraphicsResourceGetMappedPointer((void **)&devPtr[1], &sizes, gres[1]);

  calcDev(&devPtr[0]);

  // TODO: リソースをメモリ空間からアンマップする
  cudaGraphicsUnmapResources(2, &gres[0], 0);
}

// 次の期間のシミュレーションに備える
void update(void) {
  // TODO: ファイルへの保存処理
  flip = 1 - flip;
}

static int countPlanet(void) {
  int i, np = 0;
  for (i = 0; i < ncloud; ++i) {
    if (mass[i] != 0.0) {
      ++np;
    }
  }
  return np;
}

static int compareMassRev(const void *s1, const void *s2) {
  float f1 = *(const float *)s1;
  float f2 = *(const float *)s2;
  return (f2 - f1 > 0.0) ? +1 : (f2 - f1 < 0.0) ? -1 : 0;
}

// 結果をモニタリングする
void show(void) {
  static int count = 0;

  glClear(GL_COLOR_BUFFER_BIT);
  glEnableClientState(GL_VERTEX_ARRAY);
  glBindBuffer(GL_ARRAY_BUFFER, pbo[1 - flip]);
  glVertexPointer(2, GL_FLOAT, 0, 0);
  glPointSize(1);
  glColor3f(1, 1, 1);
  glDrawArrays(GL_POINTS, 0, ncloud);
  glBindBuffer(GL_ARRAY_BUFFER, 0);
  glDisableClientState(GL_VERTEX_ARRAY);

  // 太陽
  glBegin(GL_POINTS);
  glColor3f(1, 0, 0);
  glPointSize(150);
  glVertex2f(0, 0);
  glEnd();

  glutSwapBuffers();

  if (++count % 100 == 0) {
    int i, nplanet = countPlanet();
    float minmax[2];
    float *massTmp = (float *)calloc(ncloud, sizeof(float));
    memcpy(massTmp, mass, ncloud * sizeof(float));
    qsort(massTmp, ncloud, sizeof(float), compareMassRev);
    printf("%f sec: planets: %d (", dt * count, nplanet);
    for (i = 0; i < 10; ++i) {
      printf("%e ", massTmp[i]);
    }
    getMinMaxDist(&minmax[0]);
    printf("x=%f,y=%f", minmax[0], minmax[1]);
    printf(")\n");
    free(massTmp);
  }
}

ヘッダファイルは以下のようになります。

#ifndef COSMICCLOUD_H_
#define COSMICCLOUD_H_

#define BLOCKSIZE (256)

extern const float gm;
extern const int ncloud;
extern const float dt;

extern float2 *devPos[2];
extern float2 *devVel;
extern float *mass;
extern int *lockVal;
extern int flip;

void calcDev(float2 **devPtr);

#endif // COSMICCLOUD_H_

微惑星の時間発展を計算する CUDA コードは以下のようになります。

#include <cuda_runtime.h>
#include <cuda_gl_interop.h>
#include "CosmicCloud.h"
#include <cstdio>

using namespace std;

const float epsilon = 1e-05;

__device__ __managed__ float minmaxdist[2] = { 1e+5, 0.0f };

// 二地点間の距離を求める
__device__
static float distance(float2 s1, float2 s2) {
  float dx = s1.x - s2.x;
  float dy = s1.y - s2.y;
  return sqrt(dx * dx + dy * dy);
}
// 二地点間の距離の逆数を求める
__device__
static float distinv(float2 s1, float2 s2, float ep) {
  float dx = s1.x - s2.x;
  float dy = s1.y - s2.y;
  return 1/sqrt(dx * dx + dy * dy  + ep);
}

// 加速度成分を計算する
__device__
float2 calcAccel(float2 *pos, float *m, float gm, int ncloud, int ntiles) {
  __shared__ float2 cache[BLOCKSIZE];
  float2 accel = make_float2(0.0f, 0.0f);
  float2 zero = make_float2(0.0f, 0.0f);
  float a;
  int i, n;
  int x = threadIdx.x + blockIdx.x * blockDim.x;

  // 惑星間の引力
  for (n = 0; n < ntiles; ++n) {
    cache[threadIdx.x] = pos[threadIdx.x + n * blockDim.x];
    __syncthreads();

    for (i = 0; i < blockDim.x; ++i) {
      int y = i + n * blockDim.x;
      if (y >= ncloud) break;
      a = m[y] * distinv(cache[i], pos[x], epsilon);
      accel.x += a * (cache[i].x - pos[x].x);
      accel.y += a * (cache[i].y - pos[x].y);
    }
  }

  // 太陽の引力
  a = gm * distinv(pos[x], zero, epsilon);
  accel.x -= a * pos[x].x;
  accel.y -= a * pos[x].y;
  return accel;
}

// シミュレーション時間を１単位進める
__global__
void advance(float2 *posNew, float2 *posOld, float2 *vel, float *m, float gm, float dt, int ncloud, int ntiles) {
  int x = threadIdx.x + blockIdx.x * blockDim.x;
  
  if (x >= ncloud) {
    // count exceeds number limit
    return;
  }

  if (m[x] == 0.0f) {
    // 計算対象外の場合は見えないところに惑星を置く
    posNew[x].x = 1e+5;
    posNew[x].y = 1e+5;
    vel[x].x = vel[x].y = 0.0f;
  } else {
    // calculate accelaration
    float2 accel = calcAccel(posOld, m, gm, ncloud, ntiles);

    // calculate velocity
    float2 v = vel[x];
    float2 dv = make_float2(0.0f, 0.0f);

    dv.x = accel.x * dt;
    dv.y = accel.y * dt;
    v.x += dv.x;
    v.y += dv.y;

    // calculate position
    posNew[x].x += v.x * dt + 0.5f * dv.x * dt;
    posNew[x].y += v.y * dt + 0.5f * dv.y * dt;

    vel[x] = v;
  }
}

// 惑星同士の合体をチェックする
__global__
void merge(float2 *pos, float2 *vel, float *m, int *lockVal, int ncloud, int ntiles) {
  int i, n, x = threadIdx.x + blockIdx.x * blockDim.x;
  int xtiles = (x + blockDim.x - 1) / blockDim.x;
  float2 px, vx;
  __shared__ float2 cache[BLOCKSIZE];

  if (x >= ncloud) {
    return;
  }
  px.x = pos[x].x;
  px.y = pos[x].y; 
  vx.x = vel[x].x;
  vx.y = vel[x].y;

  for (n = 0; n < xtiles; ++n) {
    if (threadIdx.x + n * blockDim.x < ncloud) {
      cache[threadIdx.x] = pos[threadIdx.x + n * blockDim.x];
    }
    __syncthreads();

    for (i = 0; i < blockDim.x; ++i) {
      float ri;
      int y = i + n * blockDim.x;
      if (y >= x) break;
      ri = distance(cache[i], px);
      if (ri < epsilon) {
        // 近接を検出したので、完全非弾性衝突させる
        float mg, mx, mi;
        float2 vn;
/*
        int old;

        do {
          old = atomicCAS(lockVal + y, -1, x);
        } while(old != x && old != -1);
*/

        mx = m[x];
        mi = m[i];
        mg = mi + mx;
        if (mg != 0.0f) {
          vn.x = (vel[i].x * mi + vx.x * mx) / mg;
          vn.y = (vel[i].y * mi + vx.y * mx) / mg;
        } else {
          // 既に除外されている場合は加速しない
          vn.x = 0.0f;
          vn.y = 0.0f;
        }
        m[x] = 0.0f;
        m[i] = mg;
        vel[i] = vn;

        lockVal[y] = -1;
      }
    }
  }
}

__host__
void getMinMaxDist(float *p) {
  p[0] = minmaxdist[0];
  p[1] = minmaxdist[1];
}

__host__
void calcDev(float2 **devPos) {
  int ntiles = (ncloud + BLOCKSIZE - 1) / BLOCKSIZE;
  int sharedSize = 2 * BLOCKSIZE * sizeof(float2) + sizeof(int);

  //printf("calcDev: %llx,%llx\n", devPos[0], devPos[1]);

  // シミュレーションを実行する
  cudaMemcpy(devPos[1 - flip], devPos[flip], sizeof(float2) * ncloud, cudaMemcpyDeviceToDevice);
  advance<<<ntiles, BLOCKSIZE, sharedSize>>>(devPos[1 - flip], devPos[flip], devVel, mass, gm, dt, ncloud, ntiles);
  cudaMemset(lockVal, -1, sizeof(int) * ncloud);
  merge<<<ntiles, BLOCKSIZE, sharedSize>>>(devPos[1 - flip], devVel, mass, lockVal, ncloud, ntiles);
  cudaDeviceSynchronize();
}

最後に Makefile です。
アーキテクチャが GTX 10 シリーズ専用になっていますが、2.5 以上なら動くはずなので、
適宜環境に合わせて書き換えてください。
CXXFLAGS に GL_GLEXT_PROTOTYPES をつけないと Vertex Buffer 回りのコードがエラー
になりますのでご注意を。

CSRC=main.cc
CUSRC=calc.cu
OBJS=$(CSRC:.cc=.o) $(CUSRC:.cu=.o)
TARGET=CosmicCloud
CXXFLAGS=-g -O0 -I/usr/local/cuda/include -DGL_GLEXT_PROTOTYPES
NVFLAGS=-arch=compute_61 -code=sm_61
NVCC=/usr/local/cuda/bin/nvcc
LIBS=-L/usr/lib/nvidia-367 -lGL -lGLU -lX11 -lglut

all: $(TARGET)

$(TARGET): $(OBJS)
        $(NVCC) $(NVFLAGS) -cudart=shared -o $(TARGET) $(OBJS) $(LIBS)

.c.o:
        $(CXX) $(CFLAGS) -c $<

calc.o: calc.cu
        $(NVCC) $(NVFLAGS) -c calc.cu

clean:
        -rm $(TARGET) $(OBJS)