3D Tensor Add

3D tensor는 2D matrix에 z 축이 하나 더 붙은 것으로 볼 수 있다.

int x = blockIdx.x * blockDim.x + threadIdx.x;
int y = blockIdx.y * blockDim.y + threadIdx.y;
int z = blockIdx.z * blockDim.z + threadIdx.z;

보통:

x = width 방향
y = height 방향
z = depth 방향

으로 둔다.

3D에서도 각 z slice 안에서는 같은 2D flattening 규칙을 쓴다.

width 4 height 3

2D logical matrix

1D physical memory

3D launch

dim3 threadsPerBlock(8, 8, 4);
dim3 blocksPerGrid(
    (width  + threadsPerBlock.x - 1) / threadsPerBlock.x,
    (height + threadsPerBlock.y - 1) / threadsPerBlock.y,
    (depth  + threadsPerBlock.z - 1) / threadsPerBlock.z
);

dim3(a, b)처럼 z를 생략하면 z는 1이다. 진짜 3D 데이터를 처리할 때는 z도 지정한다.

3D를 1D로 flatten하기

3D tensor shape이 [depth, height, width]라면 한 z-slice의 원소 수는:

height * width

그래서 index는:

int idx = z * height * width + y * width + x;

뜻은:

z * height * width:
  이전 z slice들이 차지한 원소 수

y * width:
  현재 z slice 안에서 이전 row들이 차지한 원소 수

x:
  현재 row 안에서의 위치

boundary check

if (z < depth && y < height && x < width) {
    C[idx] = A[idx] + B[idx];
}

축이 하나 늘었을 뿐, 패턴은 2D와 같다.

확인

• depth=2, height=3, width=4이면 z-slice 하나의 원소 수는 몇 개인가?
• (z=1, y=2, x=3)의 idx는 몇인가?
• 3D에서 2D와 달라지는 코드는 무엇뿐인가?