Output/Weight Stationary

개요

Output Stationary과 Weight Stationary는 시스토릭 어레이(Systolic Array)에서 행렬 곱셈(GEMM) 연산 시 데이터를 어떤 하드웨어 요소에 고정시킬 것인지를 결정하는 두 가지 핵심 데이터 흐름(Data Flow) 전략이다. 시스토릭 어레이는 병렬 컴퓨팅 아키텍처의 일종으로, 데이터가 인접한 처리 요소(PE) 사이에서 규칙적으로 주기적으로 흐르며 연산이 수행된다. 심장의 펌핑과 유사한 데이터 흐름 방식에서 이름이 유래되었다.

Output Stationary은 결과 행렬의 요소가 각 PE에 고정되어 누적되도록 설계하며, Weight Stationary은 가중치(가중치 행렬의 요소)가 PE에 미리 로드된 후 고정된 채 입력 데이터만 흐르도록 설계한다. 두 전략 모두 연산 재사용(Operand Reuse)을 극대화하여 외부 메모리 접근을 최소화하는 것이 목표이며, 실제 AI 가속기(Google TPU, NVIDIA Tensor Core 등)에서 널리 사용된다.

핵심 개념

시스토릭 어레이 기본 구조

시스템 어레이는 격자형으로 배치된 처리 요소(PE) 네트워크로 구성된다. 각 PE는 곱셈-누산(MAC: Multiply-Accumulate) 유닛을 포함하며, 인접 PE와 데이터를 주고받는다.

• PE(Processing Element): 곱셈과 덧셈을 수행하는 기본 연산 단위
• 데이터 흐름 방향: 입력 행렬의 행이 위→아래로, 열이 좌→우로 흐름
• 시계 동기화: 모든 PE가 동일한 클럭에 의해 동작 (systolic = 규칙적)
• 메모리 접근 최소화: 데이터가 PE 간에서 전달되므로 글로벌 메모리 접근이 크게 감소

Output Stationary 원리

Output Stationary에서 결과 행렬 C의 각 요소 C[i][j]는 특정 PE에 고정되어 계산된다. 입력 행렬 A의 요소와 B의 요소가 PE를 통과하면서 누적 연산이 수행된다.

동작 흐름:

PE[i][j] 초기 상태: accumulator = 0

클럭 0: PE[0][0]에 A[0][0] × B[0][0] 수신 → accumulator += A[0][0] × B[0][0]
클럭 1: PE[0][0]에 A[0][1] × B[1][0] 수신 → accumulator += A[0][1] × B[1][0]
...
클럭 K-1: PE[0][0]에 A[0][K-1] × B[K-1][0] 수신 → C[0][0] = accumulator 완성

각 PE는 K번의 클럭 동안 누적 연산을 수행한 후, 완성된 결과를 출력한다.

데이터 흐름:

• A의 행: 위에서 아래로 PE를 통과 (각 행이 하나의 클럭 지연 후 하위 행으로 전달)
• B의 열: 좌에서 우로 PE를 통과 (각 열이 하나의 클럭 지연 후 우측 열로 전달)
• C[i][j]: PE[i][j]에 고정, K 클럭 후 출력

Weight Stationary 원리

Weight Stationary에서는 가중치(보통 행렬 B 또는 필터 가중치)가 각 PE에 미리 로드되고 고정된다. 입력 데이터(행렬 A)와 누적 결과(부분 합)만이 PE 사이를 흐른다.

동작 흐름:

초기 단계: 모든 PE에 가중치 B[k][j]를 사전 로드

클럭 0: PE[0][0]에 A[0][0] 수신 → partial_sum += B[0][0] × A[0][0]
클럭 1: PE[0][0]에 A[0][1] 수신 → partial_sum += B[0][1] × A[0][1]
...

Weight Stationary의 핵심 장점은 가중치를 한 번 로드하면 여러 입력에 대해 재사용할 수 있다는 점이다. 특히 추론(Inference) 단계에서 동일한 가중치를 사용하는 반복 연산에 유리하다.

데이터 흐름:

• 가중치 B: PE에 고정 (초기 로드 후 변경 없음)
• 입력 A: 위에서 아래로 흐름
• 부분 합(partial sum): 좌에서 우로 흐름

비교/분석

두 전략 비교표

메모리 대역폭 분석

N×N 행렬 곱셈에서 두 전략의 외부 메모리 접근량을 비교하면:

Output Stationary:
- A 행렬: N² 요소 스트리밍 입력 (재사용으로 PE 간 전달)
- B 행렬: N² 요소 스트리밍 입력
- C 행렬: N² 요소 출력
- 총 외부 접근: O(3N²) (A, B, C 각각 N²)

Weight Stationary:
- A 행렬: N² 요소 스트리밍 입력
- B 행렬: N² 요소 초기 로드 (1회)
- C 행렬: N² 요소 출력
- 총 외부 접근: O(2N² + N² 초기 로드)

Weight Stationary은 B의 초기 로드 비용을 제외하면 외부 메모리 접근이 더 적으나, B의 크기가 작고 A의 재사용률이 높은 경우에 유리하다.

하드웨어 자원 활용 비교

동작 원리

4×4 Output Stationary 예시

2×2 행렬 곱셈 C = A × B에서 4×4 시스토릭 어레이를 사용하는 경우를 생각해보자.

A = [a00 a01]    B = [b00 b01]    C = [c00 c01]
    [a10 a11]        [b10 b11]        [c10 c11]

클럭별 동작 (2×2 배열):

이 과정에서 각 요소 a_{ij}, b_{kl}은 정확히 한 번만 글로벌 메모리에서 로드된다.

4×4 Weight Stationary 예시

동일한 2×2 행렬 곱셈에서 Weight Stationary 방식을 사용하는 경우:

초기 로드 단계:

PE[0][0] ← b00,  PE[0][1] ← b01
PE[1][0] ← b10,  PE[1][1] ← b11

연산 단계 (입력 A가 위에서 아래로 흐름):

Weight Stationary에서는 b_{kl}이 PE에 고정되어 있으므로 B의 글로벌 메모리 접근이 완전히 제거된다.

TPU v1에서의 적용

Google TPU v1는 256×256 시스토릭 어레이를 사용하며, Weight Stationary 방식을 채택했다:

• 가중치 로드: 메모리에서 256×256 가중치 메트릭스를 한 번에 어레이에 로드
• 입력 스트리밍: 활성값(activation)이 위에서 아래로 스트리밍
• 출력 수집: 완성된 결과가 아래쪽에서 출력
• 성능: 8-bit 정수 기준 92 TOPS (Tera Operations Per Second)
• 전력 효율: Von Neumann 아키텍처 대비 30~100배 높은 에너지 효율

TPU는 가중치를 어레이에 고정시킴으로써 메모리 대역폭 병목을 극적으로 줄였으며, 이것이 TPU의 높은 에너지 효율의 핵심 원인이다.

장단점

Output Stationary

장점:

• 행렬 곱셈의 모든 요소를 균등하게 분배하여 하드웨어 활용도가 높음
• 대규모 행렬 곱셈에 적합하며, 병렬성이 극대화됨
• 출력이 각 PE에서 직접 생성되므로 출력 라우팅이 단순함
• 학습(Training)과 추론(Inference) 모두에 활용 가능

단점:

• A와 B를 모두 스트리밍으로 공급해야 하므로 외부 메모리 대역폭이 필요
• 입력 데이터의 중복 전달이 발생할 수 있음
• 출력 버퍼링이 필요할 수 있음 (동시 출력 관리)

Weight Stationary

장점:

• 가중치를 한 번 로드하면 재사용되므로 외부 메모리 대역폭이 크게 절감
• 추론 단계에서 동일한 가중치를 반복 사용하는 워크로드에 최적
• 에너지 효율이 높음 (메모리 접근 에너지 소비가 가장 큰 비중)
• 구현이 상대적으로 단순

단점:

• 가중치 로드 네트워크가 필요 (초기 로드 비용)
• 가중치가 크면 PE당 레지스터/메모리 사용량이 증가
• 가중치 업데이트가 빈번한 학습 워크로드에서는 불리
• 가중치 행렬의 크기가 어레이 크기와 일치하지 않으면 로딩 효율이 저하

관련 기술

• Systems of Linear Equations: H. T. Kung, C. E. Leiserson, "Algorithms for VLSI processor arrays", Introduction to VLSI Systems, 1979
• Google TPU: Jouppi, N. P. et al., "In-Datacenter Performance Analysis of a Tensor Processing Unit", ISCA 2017
• NVIDIA Tensor Core: Mishra, A. et al., "Efficient Neural Network Inference with Mixed Precision and Tensor Cores", NVIDIA Technical Report, 2018
• Systolic Array for GEMM: Kung, H. T., "Why Systolic Architectures?", IEEE Computer, 1982
• Weight Stationary Optimization: Chen, Y. et al., "Eyeriss: An Energy-Efficient Accelerator for Deep Neural Networks", ISCA 2016

핵심 정리

1. Output Stationary은 결과 행렬의 각 요소를 특정 PE에 고정시켜 누적 연산을 수행하는 전략으로, 대규모 GEMM에 적합하다.
2. Weight Stationary은 가중치를 PE에 사전 로드하여 고정시키고 입력 데이터만 스트리밍하는 전략으로, 추론 워크로드에서 에너지 효율이 높다.
3. 두 전략 모두 외부 메모리 접근을 최소화하여 연산 재사용을 극대화하지만, 메모리 대역폭 절감 패턴이 다르다.
4. 실제 AI 가속기(Google TPU, NVIDIA Tensor Core)에서는 워크로드 특성에 따라 적합한 전략을 선택하여 구현한다.
5. 시스토릭 어레이는 Von Neumann 아키텍처의 병목을 극복하는 핵심 기술로, AI 추론과 학습 모두에서 에너지 효율을 크게 향상시킨다.