같은 MAC 수의 GEMM이 서로 다른 시간에 끝나는 이유

GEMM 두 개의 MAC 수가 같다면 같은 array에서 걸리는 시간도 비슷할까. 정사각형 출력과 길고 얇은 출력을 그려 보면, 연산량 하나가 지운 정보가 나타난다. GEMM aspect ratio는 M, N, K가 공간과 시간에 배치되는 방식에 관한 문제다. 세 차원의 곱을 유지해도 array를 사용하는 모양은 달라진다.

M과 N은 출력 자리, K는 그 자리에서 더할 횟수다

A의 크기가 M×K, B의 크기가 K×N이면 출력 C는 M×N이고 필요한 MAC은 MNK개다. 가상의 4×4 array에서 한 PE가 한 출력 원소를 담당하도록 정하자. M을 row, N을 column에 배치하고 K는 각 PE에서 순서대로 처리한다. 출력이 4보다 크면 여러 tile로 나눈다.

차원을 이렇게 배치한 것은 하나의 선택이다. K를 여러 PE로 나누고 partial sum을 합치는 다른 구조에서는 병렬화의 조건이 달라진다. MAC 수를 비교하기 전에 어느 차원을 공간에 펼쳤는지 밝혀야 한다. MAESTRO가 시간 방향과 공간 방향의 map을 구분하는 이유도 같은 연산이 서로 다른 배치·순서로 실행되기 때문이다. MAESTRO 논문, 4쪽 3.1절

같은 256MAC을 만드는 두 출력 모양

비교를 위해 PE 사이 입력 이동은 한 cycle, PE는 cycle당 MAC 하나, 입력 공급은 충분하다고 가정한다. 한 tile을 완전히 계산한 뒤 다음 tile을 시작하고 두 tile은 겹치지 않는다. tile의 유효 크기가 r×c이면 첫 MAC부터 마지막 MAC까지 K+r+c−2 cycle로 센다. padding 위치까지 고정 길이로 실행하는 제어는 이번 모델에 넣지 않는다. 결과 저장이나 preload도 별도다.

작업 A는 M=N=4, K=16이다. 출력이 4×4라 tile 하나이며 연산은 4×4×16=256MAC이다. 계산 구간은 16+4+4−2=22cycle이다. 16PE×22cycle의 352개 연산 자리 중 256개가 유효해 약 72.7%를 채운다.

작업 B는 M=1, N=16, K=16이다. 연산은 똑같이 256MAC이지만 출력 방향에 네 개의 1×4 tile이 필요하다. 각 tile의 마지막 유효 PE가 끝나는 데 16+1+4−2=19cycle이고, 순차로 네 tile을 실행하면 76cycle이다. 물리적 16PE의 자리는 1216개가 되며 유효 비율은 약 21.1%다.

차이는 B가 더 많은 곱셈을 요구해서 생기지 않았다. 이 배치에서는 각 tile에서 세 row가 쉴 뿐 아니라, tile을 네 번 시작한다. row를 다른 독립 작업에 배정할 수 있다고 가정하지 않았으므로 그 빈자리가 그대로 남는다. 이번 시간은 정의한 가상 실행 구간의 비교이며 특정 Gemmini나 상용 accelerator의 측정값이 아니다.

K가 짧아지는 변화는 출력 tail과 다르다

세 번째 작업 C를 M=N=16, K=1로 정하면 다시 256MAC이다. 4×4 출력 tile은 16개이고 매 tile의 모든 출력 자리가 유효하다. 공간만 세면 빈 출력 자리가 없다. 그러나 누산은 각 PE에서 한 번씩만 하고 끝난다.

같은 시작·완료 규칙에서 tile 하나는 1+4+4−2=7cycle이고, 순차 tile 16개는 112cycle이다. 앞뒤 이동이 긴 reduction으로 상쇄되지 않는다. 유효 MAC을 물리적 자리 16×112로 나누면 약 14.3%다. 출력 tail이 없다는 이유만으로 시간 이용률이 높다고 말할 수 없는 예다.

서로 독립적인 tile을 계속 흘려 앞뒤 파동을 겹치면 이 단순 합보다 줄어들 수 있다. 반대로 매 tile이 입력이나 저장을 기다리면 늘어난다. K를 나눈 partial sum의 보관과 출력 경로가 다음 tile의 시작을 허용하는지도 함께 확인한다. 시작 횟수와 pipeline의 겹침을 지운 채 이 숫자를 다른 구조에 옮기지 않는다.

같은 dataflow도 저장 공간의 제약을 받는다

긴 K는 한 출력 위치에서 더 오래 일하게 하지만, 필요한 입력량과 누산 상태도 늘린다. 모든 K를 local memory에 놓을 수 없다면 K 방향을 다시 나누어 공급해야 한다. 이를 위해 출력 partial sum을 유지하거나 저장하는 비용이 발생할 수 있다. K를 늘리면 항상 유리하다는 단순한 규칙은 성립하지 않는다.

Gemmini 논문은 layer 차원과 생성한 hardware의 크기를 바탕으로 runtime에 tile 크기를 정하고 scratchpad로 옮기는 양을 선택한다고 설명한다. 실제 tile 선택은 array의 모양뿐 아니라 저장 공간과 전송 경로에 붙은 제약도 풀어야 한다. Gemmini 논문, 3쪽 Data Staging and Mapping

M과 N을 교환하는 것도 항상 대칭은 아니다. 정사각형이고 row·column의 공급 조건이 같은 이번 모델에서는 대칭이지만, 실제 직사각형 array나 비대칭 SRAM port와 입력 재사용 경로에서는 다른 비용이 나올 수 있다. 출력 주소의 연속성도 달라진다.

곱 대신 세 차원을 보관한다

GEMM의 시간이 예상과 다르면 MNK만 적은 표를 M, N, K와 tile 배치가 보이는 표로 바꾼다. 출력 방향의 빈자리, 짧은 reduction의 시작 비용과 데이터 공급 시간을 각각 센다. aspect ratio를 읽는 일은 같은 연산량을 더 작은 숫자로 꾸미는 방법이 아니라, 같은 연산량이 서로 다른 일정이 되는 이유를 찾는 방법이다.