기술 요약: DCC – 메모리 내 처리 아키텍처를 위한 머신러닝 커널의 데이터 중심 컴파일

1. 문제 정의

메모리 내 처리 (Processing-In-Memory, PIM) 아키텍처는 DRAM 뱅크 근처에 처리 유닛을 배치함으로써 머신러닝 (ML) 의 메모리 병목 현상에 대한 유망한 해결책을 제공합니다. 그러나 고성능 호스트 프로세서 (예: GPU) 와 PIM 코어 간의 데이터 배치 요구 사항에 근본적인 불일치가 존재합니다.

• 호스트 요구 사항: 뱅크 수준 병렬성과 높은 대역폭을 활용하기 위해, 호스트 프로세서는 연속된 데이터 요소가 여러 DRAM 뱅크에 분산되도록 요구합니다.
• PIM 요구 사항: PIM 코어는 로컬 뱅크의 데이터에만 접근할 수 있습니다. 로컬 대역폭을 극대화하고 효율적인 다중 요소 접근을 가능하게 하려면, 연속된 요소가 동일한 뱅크 내에 배치되어야 합니다.

이로 인해 호스트 메모리 버스를 통해 수행되는 데이터 재배열(입력 및 출력) 이 필요하며, 이는 종단 간 실행 시간을 지배할 수 있는 상당한 데이터 이동 비용을 초래합니다. 현재 PIM 시스템을 위한 컴파일러 접근 방식은 두 가지 주요 한계를 겪고 있습니다:

1. 체계적인 최적화 부재: 기존 도구들은 종종 특정 백엔드 (예: UPMEM) 나 특정 커널을 대상으로 하며, 다양한 ML 워크로드와 여러 PIM 장치 (예: 삼성 HBM-PIM, SK 하이닉스 GDDR6-AiM) 를 지원하지 못합니다.
2. 연산 중심 튜닝: 기존 컴파일러 (예: ATiM) 는 데이터 재배열 비용과 독립적으로 연산 스케줄을 최적화합니다. 저자들은 연산 변환과 데이터 재배열이 상호 의존적임을 입증했습니다. 이를 분리하여 최적화하면 하위 최적의 성능을 초래하는데, 이는 연산 효율적인 스케줄이 비싼 데이터 재배열을 필요로 할 수 있는 반면, 덜 효율적인 연산 스케줄은 더 저렴한 데이터 이동을 가능하게 할 수 있기 때문입니다.

2. 방법론: DCC

저자들은 PIM 시스템을 위한 최초의 데이터 중심 ML 컴파일러인 DCC를 제안합니다. DCC 는 통합 튜닝 프로세스 내에서 데이터 재배열 전략과 루프 수준 연산 스케줄을 공동으로 최적화합니다. 이 시스템은 네 가지 핵심 구성 요소로 이루어져 있습니다:

2.1 다층 추상화

DCC 는 PIM 메모리 계층을 연산 계층으로 매핑하는 범용 추상화를 도입하여 컴파일러를 백엔드 특정 의미론에서 분리합니다. 이 추상화는 세 가지 수준에서 작동합니다:

• 시스템 수준: 호스트 xPU(예: GPU), 호스트 메모리 공간, 그리고 PIM 장치.
• PIM 그룹 수준: 동일한 메모리 채널 또는 랭크에 속하는 PIM 코어 그룹 (및 로컬 뱅크). 이는 대규모 최적화와 그룹 수준 명령어 브로드캐스트를 가능하게 합니다.
• PIM 코어 수준: 개별 처리 유닛과 로컬 뱅크로, 메모리 접근 및 로컬 데이터 배치를 위한 세밀한 최적화를 가능하게 합니다.

2.2 데이터 중심 스케줄 생성기

루프 변환으로 시작하는 전통적인 컴파일러와 달리, DCC 는 데이터 우선 전략을 채택합니다:

1. 차원 집합 구성: 텐서 차원을 루프 변수와 상관관계시켜 "복합 차원 집합"을 형성합니다.
2. 데이터 타일 구성: 시스템은 PIM 그룹과 코어 전반에 걸쳐 후보 데이터 분할 (타일) 을 포괄적으로 생성합니다.
3. 초안 가지치기: 세 가지 규칙으로 유망하지 않은 초안을 제거합니다:
   • 중복성: 동일한 데이터 분할과 매핑을 가진 초안을 제거합니다.
   • SIMD 정렬: 분할 크기가 백엔드의 SIMD 너비 (예: 16 배) 와 정렬되지 않은 초안을 가지치기합니다.
   • 성능 동등성: 여러 초안이 동일한 최악의 경우 코어당 성능을 산출할 경우 하나만 유지합니다.
4. 매핑: 남은 초안은 연산 타일에 매핑되어 각 코어에 대한 루프 범위와 텐서 인덱스를 결정합니다.

2.3 PIM 특화 코드 최적화기

이 구성 요소는 PIM 제약 조건에 맞춘 최적화를 적용합니다:

• 데이터 재배열 최적화: 제어 가능한 온칩 메모리 (예: GPU 공유 메모리) 를 활용하여 데이터를 재구성합니다. 이는 호스트 메모리 채널에서의 직렬 읽기를 수행한 후 PIM 채널로의 병렬 쓰기를 수행하여 읽기 및 쓰기 작업 모두에 대한 채널 병렬성을 극대화합니다.
• 연산 코드 최적화: 계층적 명령어 생성을 활용합니다. 명령어 버스 오버헤드를 줄이기 위해 그룹 수준의 SIMD 스타일 명령어 (그룹 내 모든 코어에 단일 명령어 브로드캐스트) 를 우선시하고, 필요한 경우에만 뱅크 수준 명령어로 대체합니다.

2.4 결합 성능 예측기

수천 개의 후보 초안을 프로파일링하는 prohibitive 한 비용을 피하기 위해 DCC 는 학습 기반 예측기 (XGBoost) 를 사용합니다.

• 기능: 데이터 재배열 비용과 연산 실행 시간을 공동으로 추정하여 최고의 종단 간 구성을 선택합니다.
• 학습: 모델은 초안의 다양한 하위 집합에 대해 오프라인으로 학습됩니다.
• 런타임: 정적 텐서 크기의 경우, 최상의 스케줄이 사전 선택됩니다. 동적 크기 (예: LLM 의 경우) 의 경우, 예측기가 실시간으로 성능을 추정합니다.

2.5 프로그래밍 인터페이스

DCC 는 PyTorch 와 호환되는 Python 인터페이스를 제공합니다. 사용자는 커널에 @DCC_kernel로 주석을 다는 것입니다. 컴파일러는 이러한 고수준 정의를 파싱하여 데이터 및 연산 타일이 포함된 중간 표현 (IR) 을 생성하고, 저수준 백엔드 명령어를 출력합니다.

3. 주요 결과

저자들은 삼성 HBM-PIM과 AttAcc(특수 GEMV/Softmax 유닛을 갖춘 연구용 프로토타입) 라는 두 가지 최첨단 PIM 백엔드에서 DCC 를 평가하였으며, GPU 전용 실행 및 기존 PIM 구현과 비교했습니다.

• 개별 ML 커널:
  • HBM-PIM: DCC 는 GPU 실행 대비 최대 7.68 배 (평균 2.21 배) 의 속도 향상을 달성했습니다.
  • AttAcc: DCC 는 GPU 실행 대비 최대 13.17 배 (평균 3.92 배) 의 속도 향상을 달성했습니다.
  • 기존 PIM 대비 비교: DCC 는 GEMV 및 RED 와 같은 특정 커널에 대해 기존 AttAcc 구현을 최대 2.74 배 (평균 1.75 배) 까지, HBM-PIM 을 최대 2.09 배 (평균 1.53 배) 까지 개선했습니다.
• 종단 간 LLM 추론:
  • AttAcc 에서 DCC 는 GPU 전용 실행 대비 GPT-3및 LLaMA-2추론을 평균 4.52 배 (LLaMA-2 의 경우 최대 7.71 배) 가속화했습니다.
  • DCC 는 이전에 GPU 에서 병목 현상이 발생했던 QKV 생성 및 프로젝션 레이어의 효율적인 실행을 가능하게 함으로써 고정된 타일링을 사용하는 기존 AttAcc 구현을 크게 능가했습니다.
• 에너지 효율성: DCC 는 데이터 이동 감소로 인해 기존 AttAcc 시스템 대비 평균 1.24 배 에너지 소비를 줄였습니다.
• 컴파일 오버헤드: 오프라인 컴파일 시간은 낮으며 (초 단위), 가지치기 방식은 검색 공간을 약 9 배 줄입니다. 결합 예측기는 높은 정확도 (LLM 의 경우 97.37%) 를 달성하며 포괄적 탐색 대비 컴파일 시간을 약 3.8 배 단축합니다.

4. 중요성 및 주장

이 논문은 DCC 가 PIM 시스템을 위한 최초의 데이터 중심 ML 컴파일러라고 주장하며, 다음과 같은 특징을 가집니다:

1. 공동 최적화: 데이터 재배열과 연산 스케줄을 공동으로 최적화하여, 이 두 요소가 상호 의존적이며 하위 최적의 성능을 피하기 위해 함께 튜닝되어야 함을 입증합니다.
2. 다중 백엔드 지원: 다층 추상화를 통해 다양한 PIM 아키텍처 (예: HBM-PIM, AttAcc) 간 이식성을 가능하게 하며, 수동 백엔드 특정 튜닝을 요구하지 않습니다.
3. 고성능 구현: 메모리 집약적 ML 커널 및 대규모 LLM 추론을 위한 고성능을 가능하게 하여, GPU 전용 실행 및 기존 PIM 컴파일러 접근 방식을 크게 능가합니다.
4. 프로그래밍 용이성 향상: PIM 장치의 저수준 데이터 배치 관리 및 명령어 스케줄링의 복잡성을 추상화하는 고수준 Python 인터페이스를 제공하여 프로그래밍 용이성을 개선합니다.

저자들은 DCC 가 고성능과 프로그래밍 용이성 사이의 균형을 맞추며 신흥 PIM 생태계에서 실용적이고 효과적인 이종 ML 가속 경로를 제공한다고 결론지었습니다. 해당 코드는 추가 연구를 촉진하기 위해 오픈 소스로 공개되었습니다.