입력의 도착 순서가 systolic array를 채우는 방식

작은 systolic array에 곱셈기가 네 개 있으면 매 cycle 네 번의 MAC이 일어날 것 같다. 충분히 긴 계산의 가운데에서는 가능하지만, 첫 입력이 들어온 순간부터 그렇게 움직이지는 않는다. 데이터가 이웃 PE로 이동하는 데 시간이 필요하기 때문이다. systolic wavefront를 보면 peak 처리량과 한 행렬의 완료 시간이 왜 다른 숫자인지 이해할 수 있다.

같은 k의 두 입력을 만나게 한다

행렬 곱 C=A×B에서 C[i,j]는 A[i,k]와 B[k,j]의 곱을 k에 대해 더한 값이다. 한 PE에 C의 한 원소를 맡기는 가상 구조를 정하자. A는 왼쪽에서 오른쪽으로, B는 위에서 아래로 이동한다. 각 PE 사이 이동은 한 cycle이고, PE는 두 입력이 유효하게 만난 cycle에 MAC 하나를 끝낸다. 초기 누산값은 0이며 입력이나 내부 이동이 멈추는 경우는 제외한다.

좌표와 k를 0부터 세면 A[i,k]를 왼쪽 경계에 넣는 시점은 k+i, B[k,j]를 위쪽 경계에 넣는 시점은 k+j로 정한다. 그러면 A가 j번 이동하고 B가 i번 이동한 뒤 둘 다 t=k+i+j에 PE[i,j]에 도착한다. row와 column의 입구에서 붙이는 지연이 같은 k를 맞추는 역할을 한다.

Gemmini의 원문도 입력별 shift register가 두 행렬의 올바른 원소를 올바른 시점에 만나게 한다고 설명한다. 다만 실제 Gemmini는 tile 안의 조합 연결과 tile 사이 pipeline을 구분한다. 여기의 모든 PE 사이 한 cycle이라는 조건은 계산을 보여 주는 별도 모델이다. Gemmini, Systolic Array and Transposer

네 PE가 일하는 시점은 대각선으로 이동한다

2×2 array로 2×3 행렬 A와 3×2 행렬 B를 곱해 보자.

A = [[1, 2, 3],   B = [[ 7,  8],
     [4, 5, 6]]        [ 9, 10],
                       [11, 12]]

PE[0,0]은 t=0,1,2에 1×7, 2×9, 3×11을 더해 58을 만든다. PE[0,1]은 같은 A row와 B의 두 번째 column을 t=1,2,3에 만나 64를 만든다. 아래 row의 결과는 139와 154이며 마지막 PE[1,1]은 t=4에 세 번째 곱을 끝낸다.

각 cycle의 유효 MAC 수를 세면 1, 3, 4, 3, 1이다. 전체 12MAC은 맞지만, 네 PE가 모두 일하는 cycle은 가운데 하나뿐이다. t=0부터 t=4까지 다섯 cycle에 물리적으로 가능한 20MAC 자리 중 12개를 사용하므로 이 구간의 유효 이용률은 60%다. 이는 메모리나 명령 overhead까지 넣은 제품 성능 수치가 아니다.

출력 tile이 R×C이고 reduction 길이가 K인 이 모델에서는 마지막 곱이 K−1+R−1+C−1에 끝난다. 첫 유효 MAC부터 마지막 유효 MAC까지는 K+R+C−2 cycle이다. 마지막 PE가 끝나기 전에 다른 PE의 결과가 준비될 수 있지만, 결과 이동과 외부 저장까지 완료됐다는 뜻은 아니다. 여기서는 그 시간을 세지 않았다.

지연을 빼면 빨라지는 대신 짝이 어긋날 수 있다

입구의 지연이 불필요해 보여 A와 B의 모든 row·column을 같은 시점에 넣는다고 하자. PE[0,1]에서 A[0,0]은 오른쪽으로 한 번 이동해 t=1에 도착한다. 그러나 B[0,1]은 위쪽 경계의 같은 PE에 t=0에 도착하고, t=1에는 B[1,1]이 온다. 그대로 곱하면 1×8이어야 할 첫 항을 1×10으로 바꿀 수 있다.

원소를 빨리 공급한 것과 올바른 원소를 동시에 공급한 것은 다르다. 이런 모델에서 valid만 보고 두 값을 곱하면 데이터 순서의 오류를 숨길 수 있다. 입력에 k 표식을 붙여 같은 k끼리 만나게 되는지 확인하면 수치 결과에 앞서 정렬 계약을 검사할 수 있다.

실제 구조가 backpressure를 허용한다면 데이터뿐 아니라 valid와 제어 상태도 함께 정렬해야 한다. 한쪽만 멈추고 다른 쪽은 계속 이동하는 상황은 지금의 일정 간격 가정을 깨뜨린다. 필요한 buffering과 정지 규칙은 구현에서 별도로 정의한다.

peak를 읽을 때 파동의 앞뒤도 포함한다

TPU의 원 논문은 systolic 계산이 대각선 wavefront로 진행하고, software가 기능적으로는 그 내부 이동을 보지 않아도 성능에서는 latency를 고려한다고 설명한다. 그 논문의 weight가 미리 놓이는 구조와 이번 output 누산 모델을 같은 timing으로 취급할 필요는 없다. 공통 질문은 입력을 array 전체의 작업으로 바꾸는 데 걸리는 시간이다. TPU 논문, 4쪽

K가 길어지면 이 모델의 앞뒤 빈 구간이 전체에서 차지하는 비중은 작아진다. 하지만 짧은 tile을 자주 시작하거나 tile 사이에 데이터를 기다리면 다시 커진다. 여러 tile의 파동을 겹칠 수 있는지도 누산 상태와 출력 경로의 조건에 달려 있다.

첫 입력 뒤 array가 천천히 채워지는 이유는 MAC이 부족해서가 아니라 필요한 두 입력이 각 위치에 도착해야 해서다. 성능을 계산할 때 첫 MAC, 마지막 MAC과 외부에서 결과를 읽을 수 있는 시점을 나누어 세면, 같은 peak 수치 아래에서 달라지는 완료 시간을 설명할 수 있다.