output-stationary에서 partial sum의 소유권을 지키는 방법

같은 행렬 곱을 반복하는 PE가 있다. 입력은 계속 바뀌는데 PE 안의 누산값은 이전 곱을 기억한다. output-stationary라는 이름에서 고정되는 것은 최종 숫자 자체가 아니다. 같은 출력 원소를 담당하는 위치와 그곳에 남겨 두는 partial sum이다. 숫자는 reduction이 진행될 때마다 바뀐다.

출력의 소유권을 한 위치에 유지한다

출력 원소 하나를 C[i,j]=D[i,j]+Σk A[i,k]B[k,j]로 쓰자. 한 PE가 이 원소를 맡으면 초기값 D를 넣고, k가 바뀔 때마다 곱을 더한다. A와 B는 PE에 새로 공급하고 누산값은 같은 위치에 보관한다. 곱마다 partial sum을 먼 memory로 내보냈다가 다시 가져오는 경로를 줄이는 선택이다.

정해진 위치에 유지한다는 말은 반드시 하나의 물리적 register만 쓴다는 뜻도 아니다. 출력 tile이 local buffer에 머무르면서 그 안에서 여러 원소를 교대로 처리할 수 있다. 어느 범위와 시간 단위에서 stationary인지 명시해야 dataflow 이름을 실제 재사용 경로로 읽을 수 있다. MAESTRO도 시간에 따른 출력 reduction과 이를 지원하는 read–modify–write buffer를 구분해 설명한다. MAESTRO 논문, 4쪽 그림 4와 7쪽 표 2

세 번의 곱이 끝나기 전에는 숫자가 결과가 아니다

가상의 한 출력에서 D=5, A의 원소를 2, −1, 3, B의 대응 원소를 4, 6, −2로 정하자. 계산은 아래처럼 진행된다.

초기 partial sum: 5
k=0: 5 + 2×4    = 13
k=1: 13 + (−1)×6 = 7
k=2: 7 + 3×(−2)  = 1

5, 13, 7, 1은 모두 같은 출력 위치의 상태다. 최종 답은 모든 k를 처리한 뒤의 1이다. 값을 local memory에 쓸 수 있다는 사실만으로 7이 최종 결과가 되지는 않는다. 완료 여부에는 reduction의 진행 상태도 포함된다.

입력을 k=0,1과 k=2의 두 묶음으로 나눠 공급해도 이 정수 계산의 답은 같아야 한다. 첫 묶음 뒤 7을 유지하고 두 번째 곱 −6을 더하면 1이다. 두 번째 묶음을 시작할 때 다시 D=5를 넣으면 −1이 되어 앞의 계산을 잃는다.

또 두 묶음을 각각 D=5에서 계산한 7과 −1을 합치면 6이다. bias 5를 두 번 포함했기 때문이다. partial sum을 분할해 저장하는 구현에서는 첫 묶음의 초기화, 이후 묶음의 누산, 마지막 방출을 별도 상태로 다루어야 한다.

Gemmini 문서는 output-stationary 모드에서 D를 preload하고 A와 B를 통과시키며, matmul.compute.accumulated가 array의 이전 C에 이어 누산한다고 설명한다. 이 인터페이스는 초기값을 넣는 동작과 이미 가진 출력에 더하는 동작을 구분한다. 위 숫자는 문서의 benchmark가 아니라 그 구분을 보여 주는 자체 예다. Gemmini, Preloading 및 Computing

입력이 저장된 모양과 PE가 원하는 순서는 다를 수 있다

PE[i,j]가 계속 같은 출력 원소를 담당하려면 A의 같은 row에서 k가 달라지는 값과 B의 같은 column에서 k가 달라지는 값이 만나야 한다. scratchpad 한 row에 A의 여러 원소가 나란히 있다고 해서 그 row를 여러 PE에 한 번에 뿌리는 것이 항상 올바른 공급은 아니다. 특정 PE에는 그 원소들이 시간 순서대로 들어와야 할 수 있다.

Gemmini의 해당 dataflow는 programmer가 행렬의 수학적 transpose를 요청하지 않아도 입력의 공급 순서를 바꾸기 위해 transposer를 사용한다. memory에서 읽은 병렬 묶음을 array가 필요한 시간 방향의 흐름으로 바꾸는 작업이다. 수학적 A를 바꾼 것과 전달 순서를 바꾼 것을 혼동하면 주소와 결과의 대응을 잘못 읽기 쉽다. Gemmini, Systolic Array and Transposer

머무는 동안의 폭과 다음 작업의 경계도 지킨다

partial sum을 오래 유지하면 작은 입력값을 여러 번 더한 결과를 담을 폭이 필요하다. 예를 들어 127×127을 32번 더하면 516,128이다. signed 16bit 누산기의 최대값 32,767을 크게 넘는다. 두 입력이 각각 8bit 범위에 들어간다는 사실이 누산 결과의 폭까지 보장하지 않는다. overflow 규칙이나 충분한 폭을 먼저 정해야 한다.

고정된 위치는 다음 출력에도 재사용한다. 이전 결과를 외부로 옮기는 동안 같은 상태를 새 초기값으로 덮어도 되는지 확인한다. 출력 경로가 멈췄는데 아직 수거하지 못한 값을 다음 계산이 지우면, MAC 계산이 정확해도 결과를 잃는다. double buffering이나 방출 완료 조건은 그 경계를 다루는 방법이다.

reduction 순서를 바꾸는 문제도 남는다. 지금 예의 정수 계산은 충분한 폭이라는 조건에서 같은 답을 준다. 부동소수점 rounding이나 중간 saturation을 포함한 계산은 묶는 순서에 따라 달라질 수 있다. stationary라는 배치 이름이 수치 동등성을 자동으로 보증하지 않는다.

output-stationary를 이해하는 핵심은 값이 안 움직인다는 그림보다, 누가 그 값을 계속 소유하고 언제 최종값으로 내놓는지를 따라가는 것이다. 초기화와 이어 누산하는 시점, 유효한 k의 수와 방출 조건을 함께 기록하면 local 재사용의 이득과 정확한 결과를 만드는 계약을 동시에 설명할 수 있다.