케이스 스터디 — 정전 포텐셜 맵

18.1 배경: 분자 시뮬레이션과 ion placement

단백질 분자 동역학(MD, molecular dynamics) 시뮬레이션을 시작하려면, 가장 먼저 시뮬레이션 박스를 물로 가득 채우고 적절한 곳에 이온(Na⁺, Cl⁻ 등)을 배치해 전하 균형을 맞춰야 한다. 어디에 두어야 할까? 답은 의외로 간단하다. 정전 포텐셜이 가장 낮은(또는 가장 높은) 곳에 두면 된다. 양이온은 음의 포텐셜에, 음이온은 양의 포텐셜에 끌린다.

그러려면 박스 안에 3차원 격자(grid)를 깔고, 단백질의 모든 원자가 만들어내는 정전 포텐셜을 격자 위에서 계산해야 한다. 이 격자에 컬러맵을 입혀 시각화하는 도구가 VMD(Visual Molecular Dynamics) 같은 프로그램이다. 격자가 충분히 조밀하면 — 0.5 Å 간격, 100 Å 박스라면 200³ = 8백만 점 — 부드러운 등전위면(isosurface)이 분자 표면 위에 그려진다.

가장 정직한 계산법이 direct Coulomb summation이다. 격자점 g에서의 전위는

V(g) = Σ_a q_a / r_a,g

여기서 a는 원자 인덱스, q_a는 그 전하, r_a,g는 원자와 격자점의 거리. 원자 N_a개, 격자점 N_g개라면 비용은 O(N_a · N_g)다. 단백질 큰 거 잡으면 원자가 수만~수십만, 격자점이 수백만이니 곱하면 10¹¹~10¹² 회 연산. CPU로는 농담이고, GPU에 잘 태우면 분 단위에 끝난다. 잘 태우는 게 우리 일이다.

18.2 scatter vs gather — 또 다시

17장과 같은 분기점이다. 두 매핑 중 어느 쪽?

• scatter(원자 기준): 스레드 하나가 원자 a를 맡아, a가 모든 격자점에 만들어내는 기여를 격자에 누적한다. atomicAdd(V[g], ...)가 매 격자점마다 필요하다. 격자점은 수백만 개라, 같은 격자점에 동시 충돌할 확률은 낮지만 — 그래도 atomic이다. atomic 자체의 latency, 그리고 같은 캐시라인에 떨어지는 격자점에 대한 false sharing이 누적된다.
• gather(격자점 기준): 스레드 하나가 격자점 g를 맡아, 모든 원자 a를 훑으며 자기 V[g]에 누적한다. atomic 없음. 출력은 자기 자신만 쓴다.

GPU에서는 격자점이 원자보다 훨씬 많고(수백만 vs 수만), gather가 atomic 없이 깔끔하다. gather를 디폴트로 잡는다.

struct Atom { float x, y, z, q; };

__constant__ Atom g_atoms[ATOMS_PER_PASS];   // ~4K개 atoms를 한 번에

__global__ void coulomb_gather(
    float* __restrict__ V, int Nx, int Ny, int Nz,
    float h,                       // grid spacing
    int n_atoms_chunk)
{
    int gx = blockIdx.x * blockDim.x + threadIdx.x;
    int gy = blockIdx.y * blockDim.y + threadIdx.y;
    int gz = blockIdx.z;
    if (gx >= Nx || gy >= Ny || gz >= Nz) return;

    float x = gx * h, y = gy * h, z = gz * h;
    float v = V[(gz*Ny + gy)*Nx + gx];   // 이전 청크 누적

    for (int a = 0; a < n_atoms_chunk; ++a) {
        float dx = x - g_atoms[a].x;
        float dy = y - g_atoms[a].y;
        float dz = z - g_atoms[a].z;
        float r2 = dx*dx + dy*dy + dz*dz;
        v += g_atoms[a].q * rsqrtf(r2);
    }
    V[(gz*Ny + gy)*Nx + gx] = v;
}

17장 MRI 커널과 골격이 똑같다. __constant__에 atom을 올리고, 한 번에 못 올리는 양은 chunking한다. 베이스라인이 V100급에서 약 30 GFLOPS 수준. rsqrtf(역제곱근)는 SFU 인트린식이라 한 사이클에 끝난다. 그래서 17장의 sin/cos 단계 같은 별도 인트린식 단계는 여기서는 이미 디폴트로 들어가 있다.

18.3 스레드 코어스닝 — 한 스레드, 여러 격자점

위 커널의 한 가지 비효율은 inner-loop 안에서 g_atoms[a]를 16바이트씩 읽는다는 점이다. 모든 스레드가 같은 a를 동시에 읽으니 broadcast로 해결되어 좋긴 하다. 그러나 같은 atom 데이터를 스레드마다 자기 레지스터로 끌어올리는 비용이 매번 든다. 격자점을 1 스레드에 1점이 아니라 여러 점을 묶어 1 스레드에 줘 보자.

이게 thread coarsening이다. 한 스레드가 인접한 격자점 4개(예: gx, gx+1, gx+2, gx+3)를 동시에 처리하면, 같은 atom 데이터를 한 번 레지스터에 올린 뒤 4번 재사용할 수 있다. 레지스터 데이터 재사용률이 4배가 된다.

__global__ void coulomb_gather_coarse4(
    float* __restrict__ V, int Nx, int Ny, int Nz, float h, int n_atoms_chunk)
{
    int gx0 = (blockIdx.x * blockDim.x + threadIdx.x) * 4;
    int gy  =  blockIdx.y * blockDim.y + threadIdx.y;
    int gz  =  blockIdx.z;
    if (gx0 >= Nx || gy >= Ny || gz >= Nz) return;

    float x0 = gx0 * h, y = gy * h, z = gz * h;
    int base = (gz*Ny + gy)*Nx + gx0;

    float v0 = V[base+0], v1 = V[base+1], v2 = V[base+2], v3 = V[base+3];

    for (int a = 0; a < n_atoms_chunk; ++a) {
        float ax = g_atoms[a].x, ay = g_atoms[a].y, az = g_atoms[a].z, q = g_atoms[a].q;
        float dy = y - ay, dz = z - az;
        float dyz2 = dy*dy + dz*dz;

        float dx0 = x0       - ax;  v0 += q * rsqrtf(dx0*dx0 + dyz2);
        float dx1 = x0 + h   - ax;  v1 += q * rsqrtf(dx1*dx1 + dyz2);
        float dx2 = x0 + 2*h - ax;  v2 += q * rsqrtf(dx2*dx2 + dyz2);
        float dx3 = x0 + 3*h - ax;  v3 += q * rsqrtf(dx3*dx3 + dyz2);
    }
    V[base+0]=v0; V[base+1]=v1; V[base+2]=v2; V[base+3]=v3;
}

이 변환의 효과는 보통 1.4~1.8배다. 4배가 안 나오는 이유는, 메모리 트래픽 절감보다는 ALU 효율(특히 dy, dz 차이를 한 번만 계산해 재사용)의 이득이 크고, 그 이득은 4배가 아니기 때문이다. 코어스닝을 8, 16으로 더 늘리면 레지스터 압력이 올라가 occupancy가 떨어지고, 어느 시점부터 손해다. 보통 4가 sweet spot이고, 카드에 따라 8까지 간다.

18.4 메모리 코얼레싱 — 격자 인덱싱을 어떻게 잡을까

위 코드에서 V[base+0..3]는 X축 따라 인접한 4점이다. 워프의 32스레드는 X 방향으로 연속한 32×4 = 128점에 쓰는 셈이고, 이는 정렬되어 있다면 4개의 128B 코얼레스드 트랜잭션으로 처리된다. 좋다.

주의해야 할 건 격자 인덱싱 순서다. idx = (gz·Ny + gy)·Nx + gx라면 gx가 가장 안쪽 차원, 메모리상 stride 1이다. 워프가 X축에 정렬되어 일하면 쓰기가 코얼레스된다. 만약 인덱싱을 (gx·Ny + gy)·Nz + gz로 잘못 잡으면, 워프 32 스레드가 X 방향으로 32간격 stride로 쓰게 되고 코얼레스가 깨진다. 같은 문제, 같은 알고리즘인데 인덱싱 한 줄에 따라 4~8배 차이가 나는 흔한 함정이다.

또 한 가지: blockDim.x를 32의 배수로, X 방향으로 잡는 것이 자연스럽다. 한 워프가 X축으로 정렬되도록. 코어스닝 인자가 4면 워프가 커버하는 X 범위는 32×4 = 128이다. Nx가 128의 배수가 아니라면 가장자리 처리가 까다로워지는데, 이건 일반적으로 패딩(padding)으로 해결한다.

18.5 컷오프 비닝 — 데이터가 커질 때의 확장성

지금까지의 알고리즘 비용은 O(N_a · N_g). 단백질이 작고 격자도 적당하면 괜찮다. 그런데 시뮬레이션 박스가 커지고 원자가 수십만 개로 늘면 어떻게 될까? GPU가 아무리 빨라도, 박스가 커지면 격자점 수도 함께 커지므로(부피에 비례) 곱셈 비용은 부피의 제곱으로 폭증한다.

물리적 직관으로 돌아가자. 멀리 있는 원자의 기여 q/r은 r이 커지면 작아진다. 일정 거리 r_cut 이상의 원자는 무시해도 영상에 거의 변화가 없다. 이 발상이 컷오프 비닝(cutoff binning)이다.

방법은 이렇다. 박스를 한 변 r_cut인 작은 격자(spatial bins)로 나눈다. 각 bin 안에 포함되는 원자 인덱스 리스트를 만들어 둔다. 격자점 g를 처리할 때, g가 속한 bin과 그 이웃 bin(3D면 27개)만 순회한다. 각 bin 안에서 r < r_cut 조건을 추가로 걸어 자른다.

        bin (i-1, j+1)   bin (i, j+1)   bin (i+1, j+1)
              |               |               |
              +-------+-------+-------+-------+
                      |               |
        bin (i-1, j)  | bin (i, j)    |  bin (i+1, j)
              +-------+-------+-------+-------+
                      |       *g      |
                      |               |
        bin (i-1, j-1)| bin (i, j-1)  | bin (i+1, j-1)
              +-------+-------+-------+-------+

  격자점 *g 는 자기 bin과 8개(3D면 26개) 이웃 bin의
  원자만 본다. 그 외는 r > r_cut 으로 컷.

그림 18.1 — 2D 단순화. 3D에서는 격자점 하나당 27개 bin만 순회.

비닝은 비용 구조를 바꾼다. 각 bin 안의 평균 원자 수가 K개라면, 격자점 1개당 비용은 27·K 정도. 박스가 커져도 이 K는 (밀도가 일정하다면) 일정하다. 즉 비용이 O(N_g)으로 선형화된다. r_cut을 키우면 정확도는 좋아지지만 K가 r_cut³으로 늘어 비용이 다시 커진다. 보통 단백질 환경에서 r_cut = 12 Å 정도가 무난한 경험치다.

// bin 구조: 각 bin에 atom들을 평탄화해 저장
struct BinTable {
    int   bin_start[NX*NY*NZ];   // bin 시작 오프셋
    int   bin_count[NX*NY*NZ];   // bin 안 atom 수
    Atom* atom_pool;             // 모든 atom들을 bin 순서로 정렬
};

__global__ void coulomb_cutoff(
    float* V, int Nx, int Ny, int Nz, float h,
    int Bx, int By, int Bz, float bsize,
    const int* bin_start, const int* bin_count,
    const Atom* atom_pool, float rcut2)
{
    int gx = blockIdx.x*blockDim.x + threadIdx.x;
    int gy = blockIdx.y*blockDim.y + threadIdx.y;
    int gz = blockIdx.z;
    if (gx>=Nx || gy>=Ny || gz>=Nz) return;

    float x = gx*h, y = gy*h, z = gz*h;
    int bi = (int)(x / bsize);
    int bj = (int)(y / bsize);
    int bk = (int)(z / bsize);

    float v = 0.0f;
    for (int dk=-1; dk<=1; ++dk)
    for (int dj=-1; dj<=1; ++dj)
    for (int di=-1; di<=1; ++di) {
        int i=bi+di, j=bj+dj, k=bk+dk;
        if (i<0||j<0||k<0||i>=Bx||j>=By||k>=Bz) continue;
        int b = (k*By + j)*Bx + i;
        int s = bin_start[b], n = bin_count[b];
        for (int idx=0; idx<n; ++idx) {
            Atom a = atom_pool[s+idx];
            float dx=x-a.x, dy=y-a.y, dz=z-a.z;
            float r2 = dx*dx + dy*dy + dz*dz;
            if (r2 < rcut2) v += a.q * rsqrtf(r2);
        }
    }
    V[(gz*Ny + gy)*Nx + gx] = v;
}

비닝의 정량 효과. 100 Å 박스 / 50K atoms / 200³ grid 기준, naive O(N_a·N_g) = 4×10¹¹회. 컷오프 12 Å 비닝은 격자점당 평균 ~600 atom (3D 27 bins × ~22 atoms/bin)이라 총 ~5×10⁹회. 약 80배 감소다. 박스가 두 배 커지면 naive는 8배 비싸지지만 비닝은 8배(부피 비례). 즉 큰 시스템일수록 비닝의 이점이 더 커진다.

18.6 정리

17장과 18장은 같은 사고 흐름을 다른 도메인에 두 번 적용한 셈이다. scatter→gather로 atomic을 없애고, 코어스닝과 코얼레싱으로 메모리/ALU 효율을 끌어올리고, constant memory로 broadcast를 활용한다. 거기에 18장에서는 알고리즘 차원의 무기를 하나 더 꺼냈다 — 비닝이라는 공간적 자료구조다.

핵심 메시지는 두 가지다. 첫째, 같은 도구 박스가 다른 문제에서도 통한다. 둘째, 데이터가 충분히 커지면 마이크로 최적화로는 한계가 있고 알고리즘을 갈아엎어야 한다. 컷오프 비닝은 정확도를 약간 양보하고 확장성을 얻는, 19장에서 다룰 "알고리즘 선택" 의 좋은 예다.