[Paper] Tenstorrent Wormhole에서 스텐실 연산

Source: arXiv - 2605.07599v1

개요

이 논문은 RISC‑V 기반 AI 데이터‑플로우 가속기인 Tenstorrent Wormhole이 고전적인 HPC 커널인 2‑D 5‑점 스텐실을 얼마나 잘 실행할 수 있는지를 평가합니다. 스텐실을 요소별 행렬 연산 시퀀스(Axpy) 또는 행렬 곱셈(MatMul)으로 재구성함으로써, 저자들은 성능, 에너지 사용량 및 병목 현상을 기존 CPU 구현과 비교합니다.

주요 기여

• AI 가속기에서 스텐실 워크로드를 위한 두 가지 새로운 매핑 전략:
  1. Axpy – 스텐실을 요소별 서브‑매트릭스 연산 파이프라인으로 분해합니다.
  2. MatMul – 스텐실을 밀집 행렬‑곱으로 재구성하여 하드웨어의 GEMM 엔진을 활용합니다.
• 포괄적인 프로파일링을 통해 Wormhole에서의 계산 시간과 호스트‑측 오버헤드(PCIe 전송, 디바이스 초기화, 전처리)를 분리했습니다.
• 에너지 효율성 분석 결과, 큰 문제 규모에서는 Axpy 구현이 실제 벽시계 시간은 느리지만 CPU 기준선보다 전력을 적게 소비함을 보여줍니다.
• 아키텍처 및 소프트웨어 병목 현상 식별(메모리 대역폭, PCIe 지연, 네이티브 스텐실 프리미티브 부족) 및 AI 가속기를 HPC 친화적으로 만들기 위한 하드웨어‑소프트웨어 공동 설계에 대한 구체적인 권고안을 제시했습니다.

방법론

1. Benchmark selection – 저자들은 표준 2‑D 5‑점 스텐실(열‑확산, CFD, 이미지‑처리 코드에서 흔히 사용)을 사용한다.
2. Kernel redesign –
   • Axpy: 스텐실 업데이트 out[i,j] = a*in[i,j] + b*in[i‑1,j] + c*in[i+1,j] + d*in[i,j‑1] + e*in[i,j+1] 를 다섯 개의 별도 요소‑별 행렬 덧셈/곱셈으로 분해하여 Wormhole의 벡터 ALU에 자연스럽게 매핑한다.
   • MatMul: 입력 그리드를 패딩하고 재배열함으로써 스텐실을 희소 행렬‑벡터 곱으로 표현하고, 가속기는 이를 적절한 타일링 후 밀집 GEMM으로 처리할 수 있다.
3. Implementation stack – 커널은 Tenstorrent의 SDK(Python‑based API)로 작성되고 온‑칩 데이터‑플로우 ISA로 컴파일된다. 호스트 코드는 x86‑64 CPU에서 실행되며 PCIe를 통해 데이터 이동을 처리한다.
4. Baseline – OpenMP와 캐시‑블로킹을 사용한 고도로 최적화된 멀티‑스레드 CPU 버전이 기준으로 사용된다.
5. Metrics collected – 엔드‑투‑엔드 실행 시간, 가속기 별도 연산 시간, PCIe 전송량, 전력 소모(온보드 센서 이용), 그리고 에너지‑투‑솔루션을 수집한다.

결과 및 발견

요약: 데이터가 가속기에 상주하면 Wormhole은 스텐실을 최신 CPU만큼 빠르게 계산할 수 있다. 주요 지연은 PCIe를 통한 데이터 이동 및 일회성 장치 설정에서 발생한다. Axpy 변형은 전체적으로는 느리지만 큰 문제 크기에서는 에너지 효율면에서 우수하다.

실용적인 시사점

• Accelerator‑first HPC 파이프라인 – 동일한 데이터가 여러 스텐실 패스에서 재사용되는 워크로드(예: 시간 단계 시뮬레이션)의 경우, 데이터를 칩 내에 유지하면 PCIe 페널티를 상쇄하고 AI 가속기를 경쟁력 있게 만들 수 있습니다.
• 에너지 제한 엣지 슈퍼컴퓨팅 – Axpy의 낮은 줄당(Joule-per-solve) 에너지는 전력 예산이 중요한 원격 또는 임베디드 HPC 노드에 AI 칩이 매력적일 수 있음을 시사합니다.
• 소프트웨어 스택 고려사항 – 개발자는 데이터 이동 비용을 고려해야 하며, 통합 메모리 또는 NVMe 기반 스테이징을 사용하면 호스트 측 격차를 줄일 수 있습니다.
• 알고리즘 재설계 – 전통적인 커널을 가속기의 강점에 맞는 형태(예: GEMM 친화적)로 재구성하면 MatMul 매핑에서 보여준 바와 같이 숨겨진 성능을 끌어낼 수 있습니다.
• 툴링 – 이 연구는 호스트 측과 디바이스 측 비용을 구분하는 프로파일링 도구의 필요성을 강조하며, 이는 현실적인 성능 예산 책정에 필수적입니다.

제한 사항 및 향후 작업

• PCIe 병목 – 현재 Wormhole 플랫폼은 비교적 느린 PCIe Gen3 링크에 의존하고 있으며, 최신 인터커넥트(CXL, PCIe Gen4/5)는 종단 간 시간을 크게 개선할 수 있습니다.
• 메모리 계층 구조 – 온칩 SRAM이 제한적이며, 더 큰 스텐실은 오프칩 DRAM으로 넘쳐나게 되어 논문에서 완전히 탐구되지 않은 지연을 초래합니다.
• 단정밀도 중심 – 실험은 FP32를 사용했으며, 혼합 정밀도 또는 정수 스텐실 변형(일부 CFD 코드에서 흔함)은 아직 테스트되지 않았습니다.
• 확장성 – 연구는 단일 가속기에만 국한되어 있으며, 다중 노드 확장, 집합 통신 및 작업 부하 분할은 아직 미해결 질문입니다.
• 소프트웨어 성숙도 – Tenstorrent SDK는 아직 발전 중이며, 더 풍부한 프리미티브(네이티브 스텐실 연산, 향상된 DMA 스케줄링)로 수동 커널 재구성 필요성을 줄일 수 있습니다.

향후 방향은 저자들이 제시한 바와 같이 CPU‑가속기 통합 강화, halo 교환을 위한 하드웨어 지원, 그리고 스텐실 DSL을 AI 칩에 최적화된 데이터 흐름 그래프로 자동 변환하는 컴파일러 확장을 포함합니다.

저자

• Lorenzo Piarulli
• Daniele De Sensi

논문 정보

• arXiv ID: 2605.07599v1
• 분류: cs.DC, cs.ET
• 출판일: 2026년 5월 8일
• PDF: Download PDF