ANEForge: Python for direct computation on the Apple Neural Engine

기술 요약: ANEForge

문제 정의

Apple Neural Engine(ANE)은 모든 최신 Apple 기기에 탑재된 고정 기능 신경 가속기입니다. 그러나 이 하드웨어에 접근할 수 있는 유일한 공인된 방법은 Apple의 CoreML 프레임워크를 통하는 것입니다. CoreML은 ANE를 확정적인 타겟이 아닌 스케줄링 옵션 중 하나로 취급하며, 런타임 시 휴리스틱(heuristics)에 따라 모델을 컴파일하고 연산 유닛(CPU, GPU 또는 ANE)을 선택합니다. 결과적으로 ANE 실행을 강제하는 명시적인 설정이 존재하지 않으며, CPU가 허용된 폴백(fallback)으로 남게 됩니다. 또한, 런타임은 특정 호출이 어떤 유닛에서 실행되었는지 보고하지 않으므로, 가장 전력 효율적인 하드웨어에서 효율적인 워크로드가 실행되는지 보장하는 것이 불가능합니다.

ANE의 진정한 프로그래밍 표면은 Espresso, aned, 그리고 커널 드라이버를 포함한 Apple의 비공개 프레임워크 스택 내에 존재합니다. 이 스택은 외부로 노출되지 않으므로, 엔진의 역량을 직접 규명하거나 ANE에 대한 결정론적 배치를 요구하는 커스텀 워크플로우를 개발하는 것을 방해합니다.

방법론

ANEForge는 비공개 프레임워크 스택과 인터페이스하여 CoreML을 우회하고 ANE를 직접 프로그래밍하도록 설계된 Python 패키지입니다. 이 방법론은 다음과 같은 아키텍처 구성 요소를 포함합니다:

• 레이지 텐서 그래프 컴파일 (Lazy Tensor Graph Compilation): 사용자는 정의된 58개의 퓨즈드 연산자(fused operators)와 19개의 네이티브 브릿지 연산자(native bridge operators)를 사용하여 레이지 연산자 그래프를 구축합니다. 컴파일러는 이 전체 그래프를 모델 중간 언어(MIL)를 사용하여 단일 ANE 프로그램으로 낮춥니다(lower).
• 비공개 심볼을 통한 직접 디스patch (Direct Dispatch via Private Symbols): 이 패키지는 사용자의 머신에서 컴파일된 작은 Objective-C++ 쉠(shim)을 활용하여, Espresso 런타임에서 내보내는 e5rt C 인터페이스와 같은 Apple 비공개 프레임워크의 공개 심볼에 접근합니다. 이를 통해 Apple의 권한(entitlements)을 요구하거나 시스템 무결성 보호(SIP)를 수정하지 않고도 aned 시스템 데몬 및 ANE 커널 드라이버로 디스패치할 수 있습니다.
• 연산자 표면 (Operator Surface):
  • 퓨즈드 경로 (Fused Route): 58개의 연산자(컨볼루션, 행렬 곱셈, 활성화 함수, 정규화, 어텐션 포함)를 단일 프로그램으로 낮춥니다.
  • 네이티브 브릿지 경로 (Native Bridge Route): 그래프 컷(graph cuts)을 통해 19개의 네이티브 연산자(퓨즈드 어텐션, argmax, topk, 기하학 레이어 등)를 서브 프로그램으로 실행하여, 공개 툴체인에는 노출되지 않은 하드웨어 레이어에 접근합니다.
• 학습 및 상태 관리 (Training and State Management): 이 패키지는 엔진 상에서 순전파(forward pass), 역전파(backward pass), 옵티마이저 업데이트를 모두 수행할 수 있는 역방향 자동 미분(reverse-mode autograd) 모듈을 구현합니다. 엔진 상에 키-값 캐시(자기회귀 디코딩을 위한) 및 옵티마이저 상태를 유지하기 위해 출력-대-입력 버퍼 에일리어싱(output-to-input buffer aliasing)을 활용하여 호스트-디바이스 간의 데이터 전송을 방지합니다.
• 가중치 압축 (Weight Compression): 시스템은 컴파일 타임에 반정밀도로 폴딩하는 대신, 엔진의 역양자화(dequantization) 경로로부터 int8, int4(룩업 테이블), 그리고 희소(sparse) 가중치를 직접 스트리밍하는 것을 지원합니다.

주요 기여

1. CoreML 없는 접근 (CoreML-Free Access): ANEForge는 스케줄링을 위해 CoreML에 의존하지 않는 최초의 ANE Python 인터페이스를 제공하여, 하드웨어 가속기에서의 결정론적 실행을 보장합니다.
2. 네이티브 퓨즈드 어텐션 (Native Fused Attention): 공개 CoreML 컴파일러가 표준 행렬 곱셈과 소프트맥스 연산으로 분해해버리는 엔진 전용 퓨즈드 어텐션 레이어(sdpa)를 노출합니다.
3. 온디바이스 학습 (On-Device Training): 이 패키지는 그래디언트와 옵티마이저 상태를 디바이스에 유지하면서 전체 학습 루프(순전파, 역전파, 옵티마이저)를 ANE에서 가능하게 합니다.
4. 가중치 스트리밍 (Weight Streaming): 실행 경로를 손상시키지 않으면서 int4, int8, 희소 가중치를 엔진으로 직접 스트리밍하는 메커니즘을 도입하여 메모리 점유율을 최대 4배까지 줄입니다.
5. 수치 해석 방법 (Numerical Methods): ANE의 유용성을 신경망을 넘어 고정 반복 수치 선형 대수(예: Krylov 솔버) 및 스펙트럼 방법(FFT)을 정적 데이터플로우 그래프로서 포함하도록 확장합니다.
6. 기계 검증 가능 역량 레지스트리 (Machine-Checked Capability Registry): 레지스트리의 모든 연산자가 하드웨어에서 올바르게 컴파일되고 실행되는지 확인하는 적합성 테스트 스위트를 통해 새로운 릴리스에 대한 회귀 게이트 역할을 수행합니다.

결과

본 논문은 Apple Silicon (M1 Max 및 M5 Pro) 환경에서 참조 프레임워크(NumPy, PyTorch, torchvision, transformers, diffusers)를 대상으로 패키지의 성능을 검증합니다:

• 성능: 작은 퓨즈드 프로그램은 약 90 µs 내에 호출을 완료하며, 이는 엔진의 디스패치 하한선인 ~70 µs에 근접합니다. 사전 훈련된 ResNet-18 순전파는 ANE에서 엔드 투 엔드로 0.33 ms 만에 완료되며, 이는 GPU(MPS)의 2.0 ms 및 CPU의 6.0 ms와 비교됩니다.
• 정확도: ResNet-18, ViT-B/16, MiniLM을 포함한 사전 훈련된 모델들은 float32 프레임워크 참조값에 대해 1.0000의 코사인 유사도를 달성합니다.
• 학습: ANEForge를 사용하여 CIFAR-10 데이터셋으로 훈련된 컨볼루션 네트워크는 **71.2%**의 테스트 정확도를 기록했으며, 이는 PyTorch 베이스라인(72.9%)과 1.7 퍼센트 포인트 차이이며, 그래디언트는 1.0 코사인 유사도로 일치합니다.
• 최적화: 어텐션 블록을 단일 프로그램으로 퓨징하는 옵티마이저 리라이트(rewrite)를 통해 다양한 시퀀스 길이에 걸쳐 3.7–5.3x의 속도 향상을 얻었습니다.
• 전력: ResNet-18 핫 루프(hot loop) 동안 ANE는 4.5 W를 소모하며, CPU/GPU 대안에 비해 효율적인 전력 사용을 입증했습니다.

의의 및 범위

본 논문은 ANEForge를 연구 및 로컬 개발을 위한 도구로 규정하며, 비공개 프레임워크에 의존하기 때문에 App Store를 통해 배포될 수 없음을 명시적으로 밝힙니다. 주요 의의는 다음과 같습니다:

• 특성 규명 (Characterization): ANE가 수용하는 연산자에 대한 기계 검증된 조사와 제어된 조건 하에서의 속도 및 에너지 효율성에 대한 엄격한 측정 기능을 제공합니다.
• 프라이버시 보존형 개인화 (Privacy-Preserving Personalization): 데이터, 그래디언트, 옵티마이저 상태가 디바이스를 떠나지 않는 온디바이스 학습 및 개인화를 위한 기반을 제공하며, SoC의 최저 전력 컴퓨팅 블록을 활용합니다.
• 하드웨어 투명성 (Hardware Transparency): 런타임 휴리스틱의 불확실성을 제거하는 결정론적 실행 경로를 제공합니다.

저자들은 컴퓨팅 경로가 반정밀도(fp16)로 제한되어 있어 일반적인 과학 계산에 제약이 있고, 누적 오차로 인해 클래시파이어 프리 가이던스(classifier-free guidance) 확산 샘플링 루프를 완전히 닫을 수 없다는 점을 인정하며 신중한 범위를 유지합니다. 이 패키지는 사용자가 적합성 게이트에 맞춰 연산자를 추가하거나 측정 도구를 구축할 수 있도록 확장 가능하게 설계되었으나, 검증된 macOS 릴리스의 특정 하드웨어 역량 및 비공개 API 버전에 종속되어 있습니다.