Asahi Linux에서 ANE(Apple Neural Engine) 구조 이해하기

Intro

저는 M1 맥을 사용하고 있습니다.
MacOS는 거의 안 쓰고, 백업용으로 외장 SSD에만 설치해둔 상태로 Asahi Fedora Remix를 내부 NVMe에 설치하고 사용합니다.
그런데 요즘 애플에서 Apple Inteligence같은 로컬 LLM을 MacOS에 포함한다던가(무슨 이유에서인지 외장 장치에 설치한 MacOS에서는 못 설치하게 되어있어서 실제로 써보진 못했지만), AI에 집중하는 것에 대해 좀 찾아보니 처음 맥을 살 때는 그냥 이런게 있는가보다 하고 넘긴 NPU가 눈에 들어왔습니다.
FP16인데도 11tops가 나옵니다. 동시대 모바일 NPU 중에서는 꽤 높은 편입니다.
그럼 왜 아무도 직접 쓰지 않을까요? 리눅스에서 사용할 수는 없을까요?
이 글은 이에 대한 답을 얻기 위해 알아본 것들에 대해 정리한 내용입니다.

Disclaimer

저는 지금까지 CS나 반도체공학에 대해 전문적으로 배운 바가 없습니다. 제 블로그에서 다룬 내용은 모두 취미로 독학한 내용입니다. (Just for Fun!)
NPU 아키텍처에 대해서도 거의 아는 바가 없고, ML은 찍먹만 해봤습니다. 이 글은 그것을 공부하는 과정에서 쓰였기에 부족한 부분이 있을 수 있습니다. 잘 아시는 분이 계시다면 피드백을 해주신다면 감사하겠습니다.

Body

찾아보니 이미 Asahi Linux에서 리버싱을 통해 M1/M2의 Neural Engine에서 동작하는 드라이버를 구현한 분이 있었습니다. #
m1n1에서 어떻게 어떻게 하면 ANE의 동작을 관찰 가능하고 어떻게 해서 구현이 된다는데 어려워서 잘 모르겠고, 이미 누군가 해결한 문제를 다시 겪을 필요는 없겠죠.
이제 이 드라이버를 보면 ANE가 어떻게 동작하는지 대략적으로 알 수 있겠네요.

Look Around

대략적으로 코드를 살펴보고 알 수 있는 점은, ANE는 놀랍도록 추상화되고 단순한 인터페이스를 가졌다는 것입니다.
GPGPU와는 꽤 다른 방식으로 동작하며(instruction이 아니라 무언가 binary blob을 직접 하드웨어에 전달하는 것 같습니다), 주로 CNN에 최적화된 구조인 것 같고, 아무래도 LLM(특히 트랜스포머 디코더인 GPT)을 구동하기에는 효율이 낮아 보입니다.
이는 아마도 원래 FaceID나 영상, 사진 처리를 위한 연산을 주로 수행하던 Neural Engine의 계통적 특성이 아닐까 싶네요.
제가 NPU 아키텍처에 대하여 아는 것이 전혀 없기에 사실인지는 모르겠습니다만, 좀 찾아보니 LLM 이전의 NPU는 보통 이런 식이었다고 합니다.
아래는 드라이버의 IOCTL을 정의하는 헤더로, 보면 init, free, submit의 3가지 기능밖에 없습니다.
이는 드라이버가 미완성된 것이 아니라 저것 외에 기능이 따로 필요하지 않을 정도로 하드웨어나 펌웨어 레벨에서 추상화되어 있기 때문입니다.
물론 내부적으로는 훨씬 복잡하지만, 그래도 상당 부분을 하드웨어가 알아서 처리한다는 점에서 특이합니다. (NPU에 대해 아는게 없어서 이런 구조가 일반적인지는 모르겠네요.)

// SPDX-License-Identifier: GPL-2.0-only OR MIT
/* Copyright 2022 Eileen Yoon <eyn@gmx.com> */

#ifndef __ANE_ACCEL_H__
#define __ANE_ACCEL_H__

#if defined(__cplusplus)
extern "C" {
#endif

#define ANE_TILE_COUNT	0x20
#define ANE_FIFO_NID	0x40
#define ANE_CMD_GRAN	0x10

#define DRM_ANE_BO_INIT 0x1
#define DRM_ANE_BO_FREE 0x2
#define DRM_ANE_SUBMIT	0x3

struct drm_ane_bo_init {
	__u32 handle;
	__u32 pad;
	__u64 size;
	__u64 offset;
};

struct drm_ane_bo_free {
	__u32 handle;
	__u32 pad;
};

struct drm_ane_submit {
	__u64 tsk_size;
	__u32 td_count;
	__u32 td_size;
	__u32 handles[ANE_TILE_COUNT];
	__u32 btsp_handle;
	__u32 pad;
};

#define DRM_IOCTL_ANE_BO_INIT \
	DRM_IOWR(DRM_COMMAND_BASE + DRM_ANE_BO_INIT, struct drm_ane_bo_init)
#define DRM_IOCTL_ANE_BO_FREE \
	DRM_IOWR(DRM_COMMAND_BASE + DRM_ANE_BO_FREE, struct drm_ane_bo_free)
#define DRM_IOCTL_ANE_SUBMIT \
	DRM_IOWR(DRM_COMMAND_BASE + DRM_ANE_SUBMIT, struct drm_ane_submit)

#if defined(__cplusplus)
}
#endif

#endif /* __ANE_ACCEL_H__ */

Whisper.cpp

이 드라이버의 작성자가 Whisper.cpp에 제출한 pull request의 Convarsation에는 구조에 대해 많은 내용이 설명되어 있습니다. #

... The ANE, having no ISA, cannot execute instructions (e.g. shaders or openCL). All actions must be known and hard-coded at compile time as "configuration states", the poor man's instructions. The entire computation sequence must be encoded in the microcode (model.hwx), including DMA memory movement. Think of it as a stream of circuit actions triggered by one big red nuclear button (which acutally exists). I guess the ELI5 is: with static comes speed. Notably, flexible/dynamic shapes are impossible. At best, shapes can be padded and computed in full; the Stable Diffusion repo does that. High-precision math may be limited (e.g. trig, discussed below), but I wouldn't worry about representing instructions as multiply-adds. The compiler is capable of approx 97% of these.

... ANE는 ISA가 없으며, 명령어를 실행할 수 없습니다(예를 들어 쉐이더나 OpenCL같은). 모든 동작은 "설정된 상태들"로 컴파일 타임에 명시되고 하드코딩되어야 합니다. 빈자의 명령어죠. 전체 계산 시퀀스는 마이크로코드로 인코딩되어야 하며 (model.hwx), DMA(역주 - 직접 메모리 접근) 메모리 이동을 포함해야 합니다. 마치 하나의 거대한 핵버튼(실제로 존재함)에 의해 촉발되는 거대한 회로의 연쇄라고 할 수 있습니다. 5살짜리도 알아듣게 설명하자면, 정적인 상태가 속도를 만든다는 것입니다. 주목할 만한 것으로, 유연하고 동적인 구조는 불가능합니다. 가장 최선의 경우, 차원들은 패딩되고 꽉 찬 상태에서 연산될 수 있습니다. Stable Diffusion 레포지토리는 그렇게 합니다. 고정밀도 연산은 제한적일 수 있으나 (예시로 하기의 내용처럼 삼각함수 등), 곱셈-누산(역주 - NPU의 기본적 동작)을 명령어로 표현하는 것에 대해서는 걱정하지 않을 것입니다. 컴파일러는 그 동작의 97%에 대응 가능합니다.

ANE is IEE-754 F16. There's no F32 mode. I (hastily) convert F32 to F16 using ggml there. INT8 mode exists, but I've never been able to trigger it. I haven't looked at INT8 further, but nothing changes at the driver level to support it (e.g. ANE-specific NCHW tiling is precision-independent). Apple doesn't like admitting its low precision hardware limitations, so they'll only accept F32 params and sneakily downcast underneath.

ANE는 IEEE-754표준의 반정밀도 수를 사용합니다. 단정밀도 모드는 없습니다. 저는 (대충) ggml을 사용하여 여기서 단정밀도 수를 반정밀도 수로 변환했습니다. 8비트 정수 모드는 존재하지만, 트리거할 수 없었습니다. 8비트 정수를 더 하지는 않았지만, 드라이버 수준에서 그것을 지원하기 위해 바뀌는 것은 없을 것입니다 (일례로 ANE에 국한된 NCHW 타일링은 정밀도-독립적입니다.). 애플은 저정밀도에서의 하드웨어 제약에 대해 인정하고 싶어하지 않기 때문에, 단정밀도 파라미터만을 받아들이고 몰래 다운캐스팅합니다. (역주 - 여기서 정밀도는 자료구조에서의 맥락입니다.)

Direct Access in MacOS

tinygrad라는 ML 백엔드 + 프론트엔드 프레임워크를 개발하고 있는 GeoHot(PS3 jailbreak 하다가 소니한테 고소당한 걸로 유명한 그 분 맞습니다)이 MacOS에서도 ANE에 직접 접근하기 위한 시도를 하고 일정 부분 성과를 얻은 것을 발견했습니다. #
업스트림에서 해당 백엔드는 제거되었지만, 흥미로운 사실들을 다루고 있고요.

The Apple Neural Engine is a fancy DMA Engine that is based around convolutions. We don't have all the details worked out yet, but we can do some things with it. At its core, it runs through 0x300 ops in an hwx file. See aneregs for the registers used in each op.
애플 뉴럴 엔진은 컨볼루션에 주로 목적을 둔 멋진 DMA 엔진입니다. 우리는 아직 모든 세부 사항을 밝혀내지 못했지만, 그것으로 무언가 하는 것은 가능했습니다. 코어에서 그것은 최대 0x300개 작업을 하나의 hwx 파일로부터 실행합니다. 개별 작업에서 사용된 레지스터를 aneregs 폴더에서 확인하세요.

It operates out of RAM or its 4MB L2 cache. The L2 "cache" appears to be manually managed, and only applies to the input and output, not the weights. The weights are usually included in the program, and it's unclear where they are copied to.
이것은 RAM 또는 그 스스로의 4MB의 L2 캐시에서 작동합니다. L2 "캐시"는 수동으로 관리되는 것으로 보이며, 입력과 출력에만 적용되고 가중치에는 적용되지 않습니다. 가중치들은 일반적으로 프로그램 내부에 포함되어 있으며, 그것들이 어디에 복사되는지는 확실하지 않습니다.

The 16 cores likely refer to the 16 wide Kernel DMA engine. They claim 11 TOPS total, which would be 687.5 GOPS/core. Perhaps it's a 32x32 MAC running at 335 MHz. That clock speed matches the cycle count time ratio from the debug perf stats.
16 코어라는 것은 아마도 16개의 넓은 커널(역주 - 컨볼루션에서 n x m 크기의 행렬을 의미합니다) DMA 엔진인 것으로 보입니다. 그들은 코어당 687.5 GOPS가 되는 총 11 TOPS를 주장합니다. 어쩌면 이것은 335 MHz로 작동하는 32x32 크기의 MAC(곱셈-누산기)일 것입니다. 이 클럭은 디버그 성능 통계에서의 사이클 수 비율과 일치하는 속도입니다.

It works with 5D Tensors, you specify the stride for the latter 4. All strides must be a multiple of 0x40 bytes
뉴럴 엔진은 5차원 텐서와 작동하며, 후단의 4개 차원에 대해서는 폭을 명시할 수 있습니다. 모든 폭은 0x40바이트의 배수여야 합니다.

Column (width) -- aneRegs.Common.InDim.Win / aneRegs.Common.OutDim.Wout
Row (height) -- aneRegs.Common.InDim.Hin / aneRegs.Common.OutDim.Hout
Plane (channels) -- aneRegs.Common.Cin.Cin / aneRegs.Common.Cout.Cout
Depth
Group (batch) -- aneRegs.Common.GroupConvCfg.NumGroups

It works with 3 data types
3개의 자료형으로 작동합니다

UInt8
Int8
Float16

The ops have several parts
각 작업들은 몇 부분으로 나누어집니다

Header -- The base addresses for the DMA engines
KernelDMASrc -- 16x wide DMA engine for the weights/bias/nonlinearity
Common -- Specifies the parameters for the convolution
TileDMASrc -- Input DMA engine
L2 -- Use the L2 cache for Source/Result instead of RAM
NE -- Configure Kernel/MAC/Post
TileDMADst -- Output DMA engine

It can work with 8 base addresses for the DMA streams per OP
작업당 DMA 스트림을 위한 8개의 베이스 주소와 작동할 수 있습니다.

2x Read, both used for things like sum
1x Write
1x T?
4x Kernel, though only the first one seems used

Conclusion

1. 아무도 안 쓰는 이유: 변환이 너무 복잡하고 백엔드 구현이 불편함, 범용성이 떨어짐
2. 접근성만 해결되면 쓸 것 같은데 그 접근성 해결이 매우 어려움 (본문에는 안 적었는데 위에서 언급한 PR에 보면 CPU fallback을 하려면 m1n1을 써서 .hwx를 추출해야 한다고 합니다. 결국 리눅스에서 모델 변환부터 컴파일까지 전부 가능한 툴체인이 개발되어야 할 건데, 고작 11 tops짜리 NPU를 쓰기 위해서 아키텍처를 리버싱해서 컴파일러를 작성할 사람은 잘 없으니까요...)
3. 이 이상은 제 능력이 안 되니까 더 배워서 알아보도록 하겠습니다. 드라이버 코드는 그럭저럭 읽겠는데 ML 용어랑 관련 구조들을 하나도 모르니까 진도가 안 나갑니다.

Credit & References

Eileen Yoon
GeoHot
1. https://github.com/eiln/ane
2. https://github.com/ggml-org/whisper.cpp/pull/1021#issuecomment-1597676195
3. https://github.com/tinygrad/tinygrad/tree/d0e752003da3fc023fa85094d7f5b65b47dd5091/extra/accel/ane