블록 행렬 연산

matmultilingperformance

블록 행렬 연산은 큰 행렬곱을 작은 행렬곱 여러 개로 나누는 방식이다. GPU 커널에서는 보통 타일링(tiling) 이라고 부른다.

핵심 목적은 단순하다.

AB의 원소를 메모리에서 한 번 가져온 뒤, 가능한 한 많은 C 원소 계산에 재사용한다.

왜 필요한가

matrix-multiplication에서 naive 구현의 입력 읽기량은 원소 기준으로 대략 다음과 같았다.

naive input reads2mnk\text{naive input reads} \approx 2mnk

이 값은 Cmn개 원소마다 A에서 k개, B에서 k개를 다시 읽는다고 봤기 때문에 나온다.

하지만 실제로 A의 서로 다른 원소는 mk개, B의 서로 다른 원소는 kn개뿐이다. 같은 원소를 여러 번 다시 읽고 있는 것이다.

타일로 계산하기

CBM x BN 크기의 출력 타일로 나누어 계산한다고 하자.

하나의 C 타일을 만들 때는:

  • A에서 BM개의 행을 가져온다.
  • B에서 BN개의 열을 가져온다.
  • 가져온 값을 BM x BN개의 출력 원소 계산에 재사용한다.
A [m x k]
x
B [k x n]
=
C [m x n]

읽기량이 어떻게 줄어드는가

타일링 후에도 같은 원소를 여러 번 읽는 일은 남아 있다. 다만 반복 횟수가 줄어든다.

A의 원소 하나는 naive에서는 n개의 출력 열에 쓰이므로 최대 n번 읽힐 수 있다. 출력 타일의 열 크기가 BN이면, 한 번 읽은 A 값을 BN개 열에 재사용할 수 있다.

그래서 A의 입력 읽기량은 대략:

nBNmk\left\lceil \frac{n}{BN} \right\rceil mk

B의 원소 하나는 naive에서는 m개의 출력 행에 쓰이므로 최대 m번 읽힐 수 있다. 출력 타일의 행 크기가 BM이면, 한 번 읽은 B 값을 BM개 행에 재사용할 수 있다.

그래서 B의 입력 읽기량은 대략:

mBMkn\left\lceil \frac{m}{BM} \right\rceil kn

따라서 타일링된 입력 읽기량은:

tiled input readsnBNmk+mBMkn\text{tiled input reads} \approx \left\lceil \frac{n}{BN} \right\rceil mk + \left\lceil \frac{m}{BM} \right\rceil kn

출력 C를 한 번 쓰는 비용은 여전히 mn개 원소다.

이상적인 하한과 실제 구현

메모리 이동량의 이상적인 하한은:

mk+kn+mnmk + kn + mn

하지만 이 하한은 A, B, C의 각 원소를 한 번씩만 메모리에서 움직인다는 가정이다. 실제 GPU에서는 shared memory, register, cache 크기가 제한되어 있어서 이 하한에 정확히 도달하기 어렵다.

그래도 타일링은 naive 구현을 하한에 더 가깝게 만든다.

naive:
2mnk + mn

tiled:
ceil(n / BN)mk + ceil(m / BM)kn + mn

ideal lower bound:
mk + kn + mn

확인

  • 왜 naive 입력 읽기량은 2mnk인가?
  • BM을 키우면 B의 반복 읽기량이 왜 줄어드는가?
  • BN을 키우면 A의 반복 읽기량이 왜 줄어드는가?
  • 타일 크기를 무한정 키울 수 없는 이유는 무엇인가?