[Paper] 머신러닝 접근을 통한 행렬 곱셈의 런타임 최적화

Source: arXiv - 2601.09114v1

개요

이 논문은 최신 멀티코어 CPU에서 일반 행렬 곱셈(GEMM) 루틴을 실행할 때 최적의 스레드 수를 자동으로 선택하는 머신러닝 기반 기법을 제안한다. 경량 모델을 실시간으로 학습함으로써, 저자들은 인기 있는 BLAS 라이브러리에서 사용되는 기본 스레드 선택 휴리스틱에 비해 두 개의 최신 HPC 노드에서 25‑40 %의 속도 향상을 달성하였다.

주요 기여

• ADSALA 라이브러리 프로토타입 – “Architecture and Data‑Structure Aware Linear Algebra” 프레임워크로, ML 모델을 BLAS 호출에 연결합니다.
• 실시간 모델 학습 – 초기 GEMM 실행 중에 런타임 특성(행렬 크기, 캐시 사용량, 코어 토폴로지)을 수집하고 회귀 모델을 지속적으로 정제합니다.
• 스레드 수 예측 – 모델이 각 GEMM 호출에 대해 최적의 OpenMP 스레드 수를 출력하여 정적 또는 단순 휴리스틱을 대체합니다.
• 크로스 아키텍처 검증 – Intel Cascade Lake (2 × 18 코어)와 AMD Zen 3 (2 × 16 코어)에서 실험을 수행하여 ~100 MiB 이하 행렬 워크로드에 대해 일관된 성능 향상을 보여줍니다.
• 오픈소스 개념 증명 – 저자들은 코드와 학습 데이터를 공개하여 재현성과 추가 확장을 가능하게 합니다.

방법론

1. Feature extraction – 각 GEMM 호출에 대해 라이브러리는 다음을 기록합니다:

   • 행렬 차원 (M, N, K)
   • 추정 작업 집합 크기 (바이트)
   • CPU 토폴로지 (소켓, 코어 수, 하이퍼스레딩 상태)
   • 현재 시스템 부하 (선택 사항)

2. Model choice – 간단한 결정 트리 회귀 모델(또는 그래디언트 부스팅 트리)을 학습시켜, 이러한 특성을 짧은 “탐색 단계” 동안 테스트된 각 스레드 수에 대해 관찰된 실행 시간에 매핑합니다.

3. Exploration phase – 새로운 문제 크기에 대해 처음 몇 번의 GEMM 실행을 다양한 스레드 수(예: 1, ½ 코어, 전체 코어)로 수행합니다. 그 실행 시간은 학습 데이터셋을 채웁니다.

4. Prediction & deployment – 모델이 사전에 정의된 신뢰 임계값에 도달하면, 이후 GEMM 호출에 대해 최적의 스레드 수를 예측하고, 해당 설정으로 실행합니다.

5. Continuous adaptation – 새로운 행렬 크기나 시스템 상태가 학습 분포와 크게 달라지면, 라이브러리는 짧은 재탐색을 수행하여 모델을 업데이트합니다.

전체 파이프라인은 몇 밀리초 수준의 오버헤드만 추가하므로, 목표로 하는 수초 단위 GEMM 커널에 비해 무시할 수 있습니다.

결과 및 발견

• 최적 스레드 수는 종종 순진한 “모든 코어 사용” 규칙과 달랐습니다; 많은 중간 규모 행렬에 대해 코어의 절반만 사용하면 공유 캐시와 메모리 대역폭에 대한 경쟁을 줄일 수 있었습니다.
• 예측 정확도(선택된 스레드 수 vs. 무차별 최적값)는 초기 탐색 단계 이후 90 % 이상이었습니다.
• 모델 학습 및 추론 오버헤드는 테스트된 워크로드 전체 실행 시간의 2 % 미만이었습니다.

Practical Implications

• Performance‑critical libraries – 과학 파이썬(Numpy, SciPy) 개발자, 머신러닝 프레임워크(TensorFlow, PyTorch) 혹은 맞춤형 HPC 커널 개발자는 ADSALA 방식을 내장하여 수동 벤치마킹 없이 자동으로 아키텍처 인식 스레드 튜닝을 얻을 수 있다.
• Cloud & container environments – 가상화되거나 컨테이너화된 환경에서 프로세스가 볼 수 있는 물리 코어 수가 런타임에 변할 수 있는 경우, 자체 튜닝 BLAS가 실시간으로 적응하여 일관된 처리량을 제공한다.
• Energy efficiency – 메모리 대역폭을 놓고 경쟁하게 될 경우 스레드 수를 줄여 실행하면 전력 소모가 감소하며, 이는 친환경 컴퓨팅 이니셔티브에 매력적인 부수 효과이다.
• Ease of deployment – 모델이 가벼운(결정 트리) 온라인 학습 방식이므로 대규모 오프라인 자동 튜닝 데이터베이스가 필요 없으며, 라이브러리를 기존 BLAS 호출의 즉시 교체용으로 제공할 수 있다.

제한 사항 및 향후 작업

• 범위는 GEMM ≤ 100 MiB 로 제한 – 메모리 대역폭을 포화시키는 더 큰 행렬은 다른 스케일링 동작을 보일 수 있으며, 접근 방식을 전체 GEMM 스펙트럼으로 확장하는 작업은 진행 중입니다.
• 단일 노드에 초점 – 다중 노드 분산 GEMM(예: ScaLAPACK)은 여기서 다루지 않은 네트워크 지연 및 프로세스 배치 요소를 도입합니다.
• 모델 단순성 – 결정 트리는 빠르지만, 보다 표현력이 풍부한 모델(예: 신경망 또는 가우시안 프로세스)은 캐시 계층과 스레드 친화성 사이의 미묘한 상호 작용을 포착할 수 있습니다.
• 탐색 오버헤드 – 초기 프로파일링 실행은 첫 몇 번의 호출에 지연을 추가합니다; 향후 작업에서는 애플리케이션 간에 과거 데이터를 재사용하거나 중앙 저장소를 통해 모델을 공유하는 것을 목표로 합니다.

전반적으로, 이 논문은 소량의 머신러닝 인텔리전스가 손수 튜닝된 BLAS 라이브러리와 끊임없이 변화하는 멀티코어 CPU 아키텍처 사이의 격차를 메울 수 있음을 보여주며, 개발자에게 더 빠른 선형대수 커널을 구현할 실용적인 경로를 제공합니다.

저자

• Yufan Xia
• Marco De La Pierre
• Amanda S. Barnard
• Giuseppe Maria Junior Barca

논문 정보

• arXiv ID: 2601.09114v1
• Categories: cs.DC, cs.LG
• Published: 2026년 1월 14일
• PDF: PDF 다운로드