Cascade Pipeline for Leading-Order Matrix Element Evaluation on AMD Versal AI Engine Arrays

본 논문은 $\gamma\gamma \to t\bar{t}g$ 과정에 대한 선도 차수 행렬 요소를 효율적으로 평가하기 위해 AMD Versal AI 엔진 어레이에 구현된 5 단계 캐스케이드 파이프라인 아키텍처를 제시하며, 단일 CPU 코어 대비 $34\times$의 속도 향상과 $7.7\times$의 에너지 효율성 개선을 달성하면서 백만 분의 일 수준의 수치 정확도를 유지하여 초당 $1.0\times10^6$회의 평가에 달하는 예상 처리량을 실현합니다.

핵심

거대 입자 가속기 내부에서 두 개의 작은 입자 (예: 양성자) 사이에서 발생하는 거대하고 혼란스러운 충돌의 결과를 예측하려 한다고 상상해 보세요. 이를 위해 물리학자들은 '행렬 요소 (matrix element)'라는 복잡한 수학적 공식을 사용합니다. 이 공식을 계산하는 것은 거대한 다단계 퍼즐을 푸는 것과 같습니다. 문제는 신뢰할 수 있는 답을 얻기 위해, 매번 약간씩 다른 초기 조건으로 이 같은 퍼즐을 수백만 번 풀어야 한다는 점입니다.

현재 표준 컴퓨터 프로세서 (CPU) 로 이를 수행하는 것은 한 사람이 퍼즐 하나씩을 해결하려는 것과 같습니다. 정확하지만, 특히 입자 가속기가 더 강력해질수록 매우 느리고 많은 에너지를 소모합니다.

본 논문은 AMD Versal AI 엔진이라는 특수한 유형의 컴퓨터 칩을 사용하여 이러한 퍼즐을 해결하는 새로운 방법을 제시합니다. 한 사람이 전체 퍼즐을 해결하는 대신, 저자들은 칩 내부에 공장 조립 라인을 구축했습니다.

다음은 그들의 해결책이 작동하는 방식을 간단한 개념으로 분해한 것입니다:

1. "조립 라인" 문제

이 특정 입자 충돌 (두 개의 글루온이 탑 쿼크, 반탑 쿼크, 그리고 또 다른 글루온으로 변하는 과정) 에 대한 수학적 공식은 칩 내 단일 소형 프로세서의 메모리에 담기에는 너무 큽니다. 16 페이지만 들어가는 주머니에 38 페이지 분량의 설명서를 넣으려 하는 것과 같습니다.

해결책: 저자들은 설명서를 5 개의 장으로 나누었습니다. 그리고 5 단계 조립 라인을 구축했습니다.

• 1 단계: 원재료 (충돌 데이터) 를 읽고 초기 몇 단계를 준비합니다.
• 2 단계 및 3 단계: 작업을 라인 아래로 전달하며 계산에 더 많은 단계를 추가합니다.
• 4 단계 및 5 단계: 최종 계산을 완료하고 답을 출력합니다.

2. "컨베이어 벨트" (캐스케이드 파이프라인)

이 5 단계는 캐스케이드 인터페이스라는 초고속 전용 컨베이어 벨트로 연결되어 있습니다.

• 상자를 다음 사람에게 넘길 때 말하거나 허락을 기다리지 않고, 상자를 즉시 미끄럼틀을 통해 미끄러뜨리는 공장을 상상해 보세요.
• 이 칩에서 "상자"는 **토큰 (tokens)**이라고 불리는 데이터 덩어리입니다.
• 저자는 작업자들이 서로 기다리며 멈추지 않도록 엄격한 규칙 (결정적 계약) 을 설계했습니다. 모든 작업자는 상자를 언제 전달하고 언제 받아야 하는지 정확히 알기 때문에 라인이 결코 막히지 않습니다.

3. "슈퍼 공장" (동시 80 개 라인)

사용된 칩 (VCK190) 은 **400 명의 작은 작업자 (타일)**를 보유한 거대한 창고와 같습니다.

• 조립 라인을 하나만 만드는 대신, 그들은 80 개의 동일한 조립 라인을 나란히 구축했습니다.
• 각 라인에는 5 명의 작업자가 있습니다. $80 \text{ 라인} \times 5 \text{ 작업자} = 400 \text{ 작업자}$.
• 그들은 모두 동시에 작동하여 80 개의 서로 다른 퍼즐을 동시에 해결합니다.

4. 결과: 속도와 효율성

저자들은 이 "공장"을 표준 컴퓨터 프로세서 (CPU) 와 고성능 그래픽 카드 (GPU) 라는 두 가지 다른 방법과 비교하여 테스트했습니다.

• 속도: 그들의 80 개 라인 공장은 단일 표준 컴퓨터 코어보다 34 배 더 빠릅니다.
  • 참고: 최상급 그래픽 카드 (GPU) 는 전체적으로 여전히 더 빠릅니다 (그들의 칩보다 약 22 배 빠름). 하지만 GPU 는 훨씬 더 크고 비싼 기계입니다.
• 에너지: 이 부분에서 그들의 방법이 빛을 발합니다. 조립 라인이 매우 효율적이고 전문화되어 있기 때문에 매우 적은 전력을 사용합니다.
  • 퍼즐 하나를 풀기 위해 그들의 칩은 표준 컴퓨터 프로세서보다 7.7 배 적은 에너지를 사용합니다.
  • 거대한 GPU 보다는 에너지 효율이 낮지만, GPU 는 이를 수행하기 위해 막대한 양의 전력을 소비합니다. 칩의 방법은 속도가 필요하지만 거대한 전력을 소모하는 기계를 연결할 수 없는 상황에서 이상적인 "적정점"을 제공합니다.

5. 정확도 확인

저자들은 그들의 "조립 라인"이 실수를 하지 않았는지 확인했습니다. 그들의 칩에서 나온 답을 "골드 스탠다드"인 이중 정밀도 (double-precision) 계산 결과와 비교했습니다.

• 결과는 거의 완벽하게 일치했습니다. 차이점은 약 100 만 분의 1 정도로 매우 미미하여 그들이 수행하는 물리 계산에서는 무시할 수 있는 수준으로 간주됩니다.

요약

간단히 말해, 저자들은 단일 컴퓨터 칩에 담기에는 너무 복잡한 물리 계산을 5 개의 관리 가능한 조각으로 잘라내고, 이를 모두 한 번에 해결하기 위해 80 개의 병렬 조립 라인을 구축했습니다. 이 접근 방식은 높은 속도와 낮은 에너지 소비라는 이상적인 "적정점"을 창출하여, 대형 강입자 충돌기 (Large Hadron Collider) 에서 우주를 이해하는 데 필요한 시뮬레이션을 실행하기 위한 강력한 대안을 제공합니다.

기술 요약

"AMD Versal AI Engine 어레이에서의 Leading-Order 행렬 요소 평가용 캐스케이드 파이프라인"에 대한 상세한 기술 요약입니다.

1. 문제 제기

현대 고에너지 물리학 (HEP) 이벤트 생성기인 MadGraph5_aMC@NLO(MG5aMC) 는 입자 충돌에 대한 행렬 요소 ($|M|^2$) 평가에서 상당한 계산 병목 현상에 직면해 있습니다. 대형 강입자 충돌기 (LHC) 가 고광도 (High-Luminosity) 단계로 진입함에 따라 이러한 계산에 대한 수요는 비선형적으로 증가하는 반면, CPU 확장성은 제한적입니다.

• 병목 현상: 행렬 요소 평가는 특히 추가적인 실제 방출 (real emissions) 이 포함된 다중 제트 (multi-jet) 과정의 경우 전체 이벤트 생성 시간의 30~40% 를 차지합니다.
• 과제: 기존 GPU 기반 솔루션 (예: CUDACPP) 은 높은 처리량을 제공하지만 상당한 전력을 소비합니다. 필드 프로그래머블 게이트 어레이 (FPGA) 는 에너지 효율성이 뛰어나지만, 최신 AI Engine 어레이의 타일당 16kB 프로그램 메모리 (PM) 제약으로 인해 어려움을 겪습니다. $gg \to t\bar{t}g$와 같은 복잡한 과정의 단일화 (monolithic) 구현은 이 메모리 한도를 초과하여 단일 타일에 직접 매핑하는 것을 불가능하게 합니다.

2. 방법론

저자들은 메모리 제약을 극복하고 병렬성을 극대화하기 위해 AMD Versal AI Engine (AIE) 어레이 (특히 VCK190 플랫폼) 상에 캐스케이드 파이프라인 아키텍처를 제안합니다.

A. 대상 플랫폼 및 아키텍처

• 하드웨어: 1.25 GHz 클럭으로 동작하는 $50 \times 8$ 그리드 형태로 배치된 400 개의 AI Engine 타일을 갖춘 AMD Versal XCVC1902 ACAP.
• 파이프라인 분해: 16 개의 파인만 다이어그램과 10 개의 고유한 HELAS 함수를 포함하는 $gg \to t\bar{t}g$ 과정을 5 단계 파이프라인으로 분해합니다.
  • 1 단계: 파동함수 생성 (외부 스피너/벡터) 및 토큰 초기화.
  • 2 단계 및 3 단계: 페르미온 - 벡터 정점 평가 (메모리 균형을 맞추기 위해 12 개의 다이어그램 분할).
  • 4 단계: 3 글루온 정점 평가 (지연 오프 - 쉘 보손 생성기 포함).
  • 5 단계: 4 글루온 접촉 항 및 컬러 행렬 축소.
• 메모리 관리: 16kB 한도 내에 맞추기 위해 저자들은 프로그램 메모리 분할과 지연 평가를 활용했습니다. 예를 들어, 오프 - 쉘 보손 생성기 ($FFV1P0\_3$) 를 1 단계에서 4 단계로 이동시켜 1 단계의 메모리 사용량을 17.8kB 에서 15.5kB 로 줄였습니다.

B. 타일 간 통신 (캐스케이드 프로토콜)

• 메커니즘: 단계 간 통신은 384 비트 단방향 캐스케이드 인터페이스(대역폭 60 GB/s) 를 통해 수행됩니다.
• 토큰 프로토콜: 파동함수와 부분 진폭을 포함하는 구조화된 "토큰"을 교환하는 결정론적, 데드락 없는 프로토콜을 사용합니다.
  • 확장된 토큰 (1~4 단계): 5 개의 외부 파동함수, 3 개의 사전 계산된 전파자, 6 개의 컬러 플로우 진폭을 운반합니다 (18 비트/헬리시티).
  • 축소된 토큰 (4~5 단계): 3 글루온 정점의 로컬 평가 후 5 개의 파동함수와 진폭만 운반합니다 (12 비트/헬리시티).
• 결정론성: 시스템은 제로 오버헤드 및 흐름 제어 없는 작동을 보장하기 위해 동일한 루프 구조, 무조건적 쓰기, 그리고 정적으로 일치하는 토큰 수를 가진 "캐스케이드 계약"을 강제합니다.

C. 소프트웨어 적응

• HELAS 라이브러리 포트링: 기존 MG5aMC HELAS 라이브러리 (원래 스칼라 배정밀도 C++) 를 AI Engine 벡터 내장 함수(단일 정밀도 float32) 로 포트했습니다.
• 최적화:
  • 벡터화: 파동함수를 8 개 너비 SIMD 벡터에 매핑했습니다.
  • 복소수 나눗셈: 스미스 방법 (나눗셈 2 회) 을 단일 하드웨어 역수 명령어로 대체했습니다.
  • 헬리시티 캐싱: 32 개 헬리시티 구성에 대해 10 개의 파동함수를 사전 계산하여 비트 인덱싱된 룩업으로 선택함으로써 평가 횟수를 16 배 줄였습니다.
  • 컬러 축소: 컬러 정규화 나눗셈을 컴파일 타임 상수로 통합했습니다.

D. 시스템 배포

• 규모: 400 개의 사용 가능한 타일에 80 개의 독립적인 파이프라인이 매핑되었습니다 (파이프라인당 5 개 타일).
• I/O: 프로그래머블 로직 (PL) 내의 패킷 스위칭 아키텍처가 위상 공간 점을 파이프라인에 분배하고 결과를 수집합니다.

3. 주요 기여

1. 메모리 기반 파이프라인 아키텍처: 16kB PM 제약을 극복하기 위해 복잡한 다중 다이어그램 행렬 요소 계산을 여러 AI Engine 타일에 성공적으로 분할하는 새로운 5 단계 캐스케이드 파이프라인을 도입했습니다.
2. 결정론적 캐스케이드 계약: 파동함수 토큰과 정적 루프 구조를 사용하는 데드락 없는 통신 프로토콜을 개발하여 복잡한 흐름 제어 하드웨어의 필요성을 제거했습니다.
3. 완전한 HELAS 포트: 이진 인덱싱 헬리시티 캐싱 및 축소된 복소수 나눗셈과 같은 복잡한 최적화를 포함하여 완전한 HELAS 진폭 라이브러리를 AI Engine 벡터 내장 함수로 성공적으로 포트했습니다.
4. 확장 가능한 배포: VCK190 의 AI Engine 컴퓨팅 자원을 100% 활용하는 80 개 파이프라인의 이론적 배포를 시연했습니다.

4. 결과

• 처리량: 80 개 파이프라인 어레이의 예상 처리량은 초당 **$1.0 \times 10^6$개의 행렬 요소 평가 (ME/s)**입니다.
  • 이는 단일 스레드 CPU 코어 (Intel i5-10600) 대비 34 배의 속도 향상을 의미합니다.
  • NVIDIA A100 GPU ($2.18 \times 10^7$ ME/s) 보다는 낮지만, AI Engine 솔루션은 훨씬 더 에너지 효율적입니다.
• 에너지 효율성:
  • AI Engine: 54.8 µJ/ME (54.8 W AIE 도메인 전력 기준).
  • CPU: 422 µJ/ME.
  • GPU: 7.3 µJ/ME (하지만 159 W 전력 소비).
  • 개선: AI Engine 은 CPU 기준 대비 에너지 효율성이 7.7 배 향상되었습니다.
• 정밀도: MG5aMC 배정밀도 참조값과 비교하여 검증되었습니다.
  • 평균 상대 오차: 1.43 ppm(백만 분의 일).
  • 최대 상대 오차: 168 ppm.
  • 이 수준의 정밀도는 물리적 불확실성 (스케일 변화, PDF 등) 이 수치적 오차보다 우세한 Leading-Order (LO) 계산에 충분하다고 판단됩니다.
• 자원 활용도:
  • 프로그램 메모리: 1 단계가 94.7% 활용도 (15,514 바이트) 로 병목 현상을 일으킵니다.
  • 프로그래머블 로직: modest한 사용량 (LUT 4.72%, 레지스터 2.87%) 으로 추가 로직을 위한 여유 공간이 있습니다.

5. 의의 및 향후 작업

• 의의: 이 연구는 AI Engine 어레이가 GPU 전력 봉투가 금지적인 환경 (예: LHC 의 온라인 트리거 시스템 또는 엣지 컴퓨팅) 에서 특히 전력 제약이 있는 환경에 대한 고통처리량, 고효율 HEP 이벤트 생성에 유효함을 입증했습니다. 이는 제약된 메모리 타일에 걸쳐 복잡한 물리 커널을 분할하는 체계적인 방법론을 확립합니다.
• 한계: 현재 구현은 Leading-Order (LO) 과정으로 제한됩니다. 지연 시간 수치는 전체 어레이의 직접 하드웨어 타이밍이 아닌 사이클 근사 시뮬레이터에 기반합니다.
• 향후 방향:
  • 헬리시티 필터링: 내부 루프 반복 횟수를 줄이기 위해 유효한 헬리시티 마스크를 사전 계산하여 처리량을 잠재적으로 두 배로 늘리는 것.
  • 더 높은 다중성: $t\bar{t}ggg$와 같은 더 복잡한 과정을 위한 파이프라인 깊이 확장.
  • NLO 통합: 루프 적분을 포함하는 Next-to-Leading-Order 계산을 위한 아키텍처 적응.
  • 하드웨어 진화: 더 큰 어레이나 더 높은 클럭 속도를 가진 차세대 Versal 장치 활용.

결론적으로, 이 논문은 복잡한 행렬 요소 계산에 내재된 메모리 분할 문제를 해결하기 위해 AMD Versal AI Engine 의 고유한 캐스케이드 기능을 활용하여, 특정 HEP 워크로드에 대한 GPU 가속의 견고하고 에너지 효율적인 대안을 제시합니다.