3DGS 소스 코드 분석 : Tile-Grid based Rasterization

Tile-Grid based Rasterization

3DGS는 최적의 병렬 처리를 위해 타일 그리드 기반 rasterization을 채택했다. 

#define NUM_CHANNELS 3 // Default 3, RGB
#define BLOCK_X 16
#define BLOCK_Y 16

이는 이미지를 16x16 픽셀 타일로 분할하고 각 타일은 1 CUDA block이고, 각 픽셀은 1 thread이다. 

같은 타일은 같은 Gaussian들을 사용한다.

전체 파이프라인은 다음과 같다.

1. Preprocessing (per-Gaussian)
2. Prefix Sum (offset 계산)
3. Key-Value 쌍 생성
4. Radix Sort (타일별 + 깊이별 정렬)
5. Tile Range 식별
6. 타일별 렌더링 

타일 기반 방식은 다음과 같은 장점들이 있다. 

1. 메모리 효율적이다 : Shared memory로 Gaussian 데이터를 재사용한다. 
2. 병렬성 : 타일 간 프로세스는 독립
3. Early Termination : Transmittance가 0.0001보다 작으면 해당 픽셀은 조기 종료하도록
4. Coalesced Access : 인접 픽셀이 같은 가우시안 데이터 접근 

┌─────────────────────────────────────────────────────┐
│ Coalesced (좋음)                                     │
│                                                     │
│ Thread 0 → Memory[0]                                │
│ Thread 1 → Memory[1]   ──→ 1번의 메모리 트랜잭션  	      │
│ Thread 2 → Memory[2]                                │
│ Thread 3 → Memory[3]                                │
└─────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────┐
│ Non-Coalesced (나쁨)                                 │
│                                                     │
│ Thread 0 → Memory[0]   ──→ 4번의 메모리 트랜잭션         │
│ Thread 1 → Memory[100]                              │
│ Thread 2 → Memory[50]                               │
│ Thread 3 → Memory[200]                              │
└─────────────────────────────────────────────────────┘

point_list 메모리:
┌────┬────┬────┬────┬────┬────┬────┬────┐
│ G3 │ G7 │ G12│ G1 │ G9 │ G4 │ G8 │... │
└────┴────┴────┴────┴────┴────┴────┴────┘
  ↑    ↑    ↑    ↑
  T0   T1   T2   T3  (연속된 thread가 연속된 메모리 접근)
  
  → 1번의 메모리 트랜잭션으로 여러 thread의 데이터를 한꺼번에 가져옴

 

1. Preprocessing

각 가우시안에 대해 다음과 같은 단계를 수행한다. 

1. Frustrum Culling : 카메라 시야 밖이면 제외
2. 3D -> 2D 투영 : means2D -> points_xy_image
3. Covariance 계산 : Scale + Rotation -> 3D Cov -> 2D Cov
4. Bounding Rect 계산 : 2D covariance의 eigenvalue 
5. Tile 개수 카운트 : tiles_touched[idx] = 겹치는 타일 수

오리지널 코드는 이렇게 tiles_touched를 계산함

2. Prefix Sum

// E.g., [2, 3, 0, 2, 1] -> [2, 5, 5, 7, 8]
cub::DeviceScan::InclusiveSum(..., geomState.tiles_touched, geomState.point_offsets, P)

각 가우시안은 여러 타일에 걸쳐 있을 수 있고, 이 정보를 저장할 배열이 필요한데 각 Gaussian마다 겹치는 타일 수가 다르다. 

Gaussian 0: 2개 타일에 걸침
Gaussian 1: 3개 타일에 걸침
Gaussian 2: 0개 타일에 걸침 (화면 밖)
Gaussian 3: 2개 타일에 걸침
Gaussian 4: 1개 타일에 걸침
─────────────────────────────
총 8개의 (Gaussian, Tile) 쌍 저장 필요

Prefix sum은 각 가우시안이 결과 배열에서 어디부터 쓰기 시작해야 하는지 알려준다. 

tiles_touched:  [2,  3,  0,  2,  1]
                 ↓ Prefix Sum (Inclusive)
point_offsets:  [2,  5,  5,  7,  8]

// duplicateWithKeys 커널
if (radii[idx] > 0)
{
    // 이 Gaussian의 쓰기 시작 위치
    uint32_t off = (idx == 0) ? 0 : offsets[idx - 1];
    
    // 겹치는 모든 타일에 대해 key-value 쌍 생성
    for (int y = rect_min.y; y < rect_max.y; y++)
    {
        for (int x = rect_min.x; x < rect_max.x; x++)
        {
            gaussian_keys_unsorted[off] = key;
            gaussian_values_unsorted[off] = idx;
            off++;  // 다음 위치로
        }
    }
}

3. Key-Value 쌍 생성 (Duplicate)

각 가우시안이 겹치는 모든 타일에 대해 key-value 쌍을 생성한다. 나중에 Radix sort로 정렬할 때 1) 타일별로 먼저 그룹화하고 2) 같은 타일 내에서는 Depth 순으로 정렬하기 위함이다. 이를 위해서 64 bit key를 만든다. 

// Key = [tile_id (32bit) | depth (32bit)]
uint64_t key = y * grid.x + x;  // tile ID
key <<= 32;
key |= *((uint32_t*)&depths[idx]);  // depth

gaussian_keys_unsorted[off] = key;
gaussian_values_unsorted[off] = idx;  // Gaussian index

1) Tile id 계산

uint64_t key = y * grid.x + x;

grid.x = 가로 타일 개수 (예: 80)
grid.y = 세로 타일 개수 (예: 45)

타일 좌표 (x=2, y=1)의 ID:
key = 1 * 80 + 2 = 82

┌────┬────┬────┬────┬─...
│ 0  │ 1  │ 2  │ 3  │     ← y=0
├────┼────┼────┼────┼─...
│ 80 │ 81 │ 82 │ 83 │     ← y=1
└────┴────┴────┴────┴─...
              ↑
           tile_id = 82

2) 상위 32비트로 이동

key <<= 32;

Before: key = 0x00000052 (82)
        ┌────────────────┬────────────────┐
        │   상위 32bit   │   하위 32bit   │
        │  0x00000000    │  0x00000052    │
        └────────────────┴────────────────┘

After:  key = 0x0000005200000000
        ┌────────────────┬────────────────┐
        │   상위 32bit   │   하위 32bit   │
        │  0x00000052    │  0x00000000    │
        └────────────────┴────────────────┘

3) Depth를 하위 32비트에 삽입

key |= *((uint32_t*)&depths[idx]);

depths[idx] = 5.3f (float)

float를 비트 그대로 uint32_t로 해석:
*((uint32_t*)&depths[idx]) = 0x40A9999A

OR 연산:
        ┌────────────────┬────────────────┐
key:    │  0x00000052    │  0x00000000    │
        └────────────────┴────────────────┘
                         OR
        ┌────────────────┬────────────────┐
depth:  │  0x00000000    │  0x40A9999A    │
        └────────────────┴────────────────┘
                         =
        ┌────────────────┬────────────────┐
result: │  0x00000052    │  0x40A9999A    │
        └────────────────┴────────────────┘
         (tile_id=82)     (depth=5.3f)

이렇게 하는 이유는 Radix sort는 숫자 크기순으로 정렬하기 때문이다. 

64비트 key 비교 시:
┌────────────────┬────────────────┐
│   Tile ID      │     Depth      │
│  (상위 32bit)  │  (하위 32bit)  │
└────────────────┴────────────────┘
     ↑ 먼저 비교      ↑ 그 다음 비교
     
정렬 결과 :

Tile 0의 Gaussians (depth 순)
Tile 1의 Gaussians (depth 순)
Tile 2의 Gaussians (depth 순)
...

4. Radix Sort (타일별 + 깊이별 정렬)

cub::DeviceRadixSort::SortPairs(
    ...,
    point_list_keys_unsorted, point_list_keys,    // keys
    point_list_unsorted, point_list,              // values (Gaussian indices)
    num_rendered, 0, 32 + bit)  // 정렬 비트 범위

1차 정렬은 Tile ID(상위 32비트)에 대해 이루어지고 2차 정렬(하위 32비트)은 Depth에 대해 이루어진다. 이렇게 정렬하면 같은 타일의 Gaussian들이 Depth 순으로 정렬된다. 

5. Tile Range 식별

정렬된 리스트에서 각 타일의 시작/끝 인덱스를 찾는다. 

uint32_t currtile = key >> 32;
uint32_t prevtile = point_list_keys[idx - 1] >> 32;
if (currtile != prevtile) {
    ranges[prevtile].y = idx;   // 이전 타일 끝
    ranges[currtile].x = idx;   // 현재 타일 시작
}

6. 타일별 렌더링 (Alpha blending)

 각 타일(Block)이 독립적으로 병렬 처리한다. 

// 1 block = 1 tile, 1 thread = 1 pixel
__global__ void renderCUDA(...) {
    // 이 타일이 처리할 Gaussian 범위
    uint2 range = ranges[tile_id];
    int toDo = range.y - range.x;
    
    // Shared memory로 Gaussian 데이터 배치 로드
    __shared__ int collected_id[BLOCK_SIZE];
    __shared__ float2 collected_xy[BLOCK_SIZE];
    __shared__ float4 collected_conic_opacity[BLOCK_SIZE];
    
    float T = 1.0f;  // Transmittance
    float C[3] = {0};  // 누적 색상
    
    for (int i = 0; i < rounds; i++) {
        // Batch로 Gaussian 데이터 로드 (coalesced memory access)
        collected_id[thread_rank] = point_list[range.x + progress];
        collected_xy[thread_rank] = points_xy_image[coll_id];
        ...
        __syncthreads();
        
        // 각 픽셀이 자기 위치에서 Gaussian 기여도 계산
        for (int j = 0; j < BLOCK_SIZE; j++) {
            // Gaussian evaluation (2D conic)
            float power = -0.5f * (con_o.x * d.x * d.x + con_o.z * d.y * d.y) - con_o.y * d.x * d.y;
            float alpha = min(0.99f, con_o.w * exp(power));
            
            // Alpha blending (front-to-back)
            C[ch] += features[...] * alpha * T;
            T *= (1 - alpha);
            
            // Early termination
            if (T < 0.0001f) done = true;
        }
    }
    
    // 최종 색상 = 누적 색상 + 배경 * 남은 투과율
    out_color[pix_id] = C[ch] + T * bg_color[ch];
}