사후 분석을 통한 Apple Neural Engine 역공학

3년 전, 저는 역공학된 Apple Neural Engine (ANE) 드라이버 작업을 중단했습니다. ANE 블록이 그다지 유용하지 않다는 슬픈 깨달음을 얻었고, 더 쓸모 있는 일을 할 수 있게 되자 다른 더 유용한 블록들을 업스트림(upstreaming)하는 작업으로 넘어갔습니다. ANE의 아키텍처는 너무 의견이 강해서 이를 중심으로 범용 가속기 플랫폼을 구축하기 어려웠고, ANE 하드웨어 API 접근을 효과적으로 여는 리눅스 드라이버로는 이 블록이 수행할 수 있는 워크로드의 종류를 넓힐 수 없었습니다. 심지어 macOS조차도 Finder에서 업샘플링된 미리보기 이미지를 생성하는 데 자체 ANE만을 정기적으로 사용합니다.

M1 die

M5 (2025)의 주요 기능은 'LLM 성능'이었고, 그들은 ANE 코어들을 GPU 코어 안에 편리하게 통합했습니다. 저는 이것이 올 것이라고 알고 있었지만, 공식적으로는 독립형 NPU(신경 처리 장치)의 종말처럼 느껴집니다. 그래서 ANE의 명백한 몰락을 기념하며, 우리는 훨씬 더 쓸모없는 일을 할 것입니다: M1의 ANE를 되돌아가서 역공학하고, 우리가 시작했던 것을 마무리할 것입니다. 3년이나 지났습니다(젠장), 그리고 저는 처음 이 작업을 했을 때보다 더 많이 알아야 합니다.

만약 3년 전 목표가 ANE에서 연산을 실행하여 이를 유용하게 만드는 것이었다면; 이번에는 전체 내부 아키텍처—컴퓨트, 데이터패스, 스케줄러, 메모리 및 실행 모델—를 매핑하는 것이 더 중요합니다. 왜냐하면 이러한 내부 설계 결정들은 Apple이 A11 Bionic (2017)에서 실리콘에 우선적으로 전념할 의향이 있었던 ML 워크로드에 대한 가정을 드러내고, 이것이 CNN 시대의 NPU에서 트랜스포머(transformer) 워크로드를 실행하는 오늘날의 GPU로의 전환에 대해 무엇을 말해주는지를 보여주기 때문입니다.

16개의 컴퓨트 코어는 아마도 ANE에서 가장 흥미롭지 않은 부분일 것입니다. Apple은 원래 예측 가능한 재사용이 있는 밀집 텐서 축소(dense tensor reductions)를 포함하는 밀집 이미지 처리 CNN 워크로드를 목표로 했습니다. M1 ANE 컴퓨트 코어는 곱셈-누산(MAC, multiply-accumulate) 유닛의 대규모 병렬 배열이지만, 이것만으로는 이 블록이 어떤 워크로드에 대해 설계되었고 가속화하는지에 대해 거의 아무것도 말해주지 못합니다.

ANE die layout

합성곱 계층(convolutional layer)은 활성화 창(activation window)과 학습된 커널 가중치(learned kernel weights) 간의 내적(dot product)을 수행하고, 어텐션(attention)은 쿼리(query)와 키 벡터(key vector) 간의 내적을 수행합니다. 내적은 내적이며, MAC(Multiply-Accumulate)는 그저 그것만을 수행합니다. 2017년 CNN 모델에 특화된 ANE이 MAC 자체가 아니라, MAC 주변의 데이터 흐름(dataflow)입니다: 언제, 어디서 MAC의 입력과 출력이 진입하고, 머무르고, 이동하는가. 트랜스포머(transformers)가 무너졌다는 가정은 특히 자기회귀 디코드(autoregressive decode)를 할 때 예측 가능한 재사용 패턴이었으며, ANE은 이 패턴을 활용하여 휴대폰에서 실행할 수 있을 만큼 충분히 효율적인 데이터 흐름으로 설계되었습니다. M5 결정은 ANE의 컴퓨트 코어(compute core)가 여전히 트랜스포머에 유용하지만, 다른 데이터 흐름 내에서 그렇다는 것을 확인시켜 줍니다.

그럼에도 불구하고, 16개의 각 컴퓨트 코어 내부의 데이터 경로(datapath)는 다음과 같습니다:

┌────────────────────── core ─────────────────────┐
│ ┌───────── 256× MACs ─────────┐ ┌────────────┐ │
│ │ MAD ─► add ─► accumulator │─►│ activation │ │
...

ANE은 16개의 병렬 컴퓨트 코어를 가지고 있습니다. 각 컴퓨트 코어는 128개의 FP16 (또는 256개의 INT8) 병렬 곱셈-누산기(MAC) 레인으로 구성되어 있습니다. 각 MAC 레인은 다음의 재귀를 수행합니다:
[ s hanks{s} imes a imes b ]
두 피연산자 $a$와 $b$를 곱한 후, 그 결과를 누적 합계기(accumulator)에 더하는 것입니다.

MAC 연산을 T 사이클 동안 반복하면 T-항 내적을 계산합니다:
[ s_T=s_0 + \sum_{t=0}^{T-1} a_t , b_t. ]
따라서 MAC 레인은 시간에 걸쳐 스칼라 축소(scalar reduction)를 수행합니다. 16코어 ANE은 2048개의 병렬 MAC 레인을 가지므로,

각 사이클마다 2048개의 병렬 축소(reduction)가 공간적으로 이루어지며, 시간만이 유일한 축소 축입니다:

개별 MAC 레인은 자신이 행렬이나 텐서의 어떤 차원에 걸쳐 축소를 수행하는지 알지 못합니다. 피연산자들이 코어에 어떻게 매핑되고 스케줄링되느냐에 따라 내적(dot product)이 될지, 행렬 곱셈(matrix multiplication)이 될지, 합성곱(convolution)이 될지가 발생한다는 점을 주목하는 것이 중요합니다. ANE 코어는 (나중에 논의될 커널 메모리를 제외하고) 4채널 CNN 계층을 하드웨어에 인코딩하지 않습니다.

내부적으로 MAC 데이터패스는 승산기(multiplier), 가산기(adder), 그리고 32비트 누적기 레지스터로 구성됩니다. 매 사이클마다 가산기는 새로운 승산기 출력에 이전 합계를 더하고, 이 값이 새로운 실행 합계가 됩니다.

operand a ──┐ ┌────────────┐ p[31:0] ┌──────────────┐ s_next[31:0] ┌─────────────┐
├──►│ MULTIPLIER │────────────►│ 32-BIT ADDER │────────────────►│ ACCUMULATOR │
operand b ──┘ └────────────┘ └──────▲───────┘ └──────┬──────┘
...

이 피드백 경로는 부분 합계(partial sum)를 MAC 레인에 로컬 메모리에 유지하므로, MAC 사이클 사이에 멀리 떨어진 외부 메모리에서 가져올 필요가 없습니다.

해상도와 관련하여, 읽기 시점에서는 FP16을 사용한 고정 소수점 감소(fixed-point reduction)를 수행합니다. 승산기는 16비트이며, Q16.16 형태로 32비트 레지스터에 누적된 후, 부호 확장 등을 통해 FP16으로 읽어 나갑니다. 정수형 (hex) FP16 표현으로 작동하여 누적기 범위를 테스트하기 위해 모든 (1)로 구성된 벡터와 내적(dot product)을 계산하는 CoreML ANE 프로그램을 구축했습니다. 그러면 각 승산기는 유한한 v를 결과로 얻지만, 누적기의 실행 합계는 계속 증가합니다:
[ s=\sum_{i=0}^{255}v=256v.]

(v)	CPU hex	CPU value	ANE hex	CoreML value
`127.9375`	`0x77ff`	`32752`	`0x77ff`	`32752`
`128`	`0x7800`	`32768`	`0x7c00`	`+∞`
`−128`	`0xf800`	`−32768`	`0xf800`	`−32768`
`−128.125`	`0xf801`	`−32800`	`0xfc00`	`−∞`

32768 자체가 유효한 FP16 워드(0x7800)이므로, ANE의 0x7c00은 FP16 출력 오버플로우가 아닙니다. 클램프는 누적기 내부에서, $2^{15}$에서 발생합니다. 따라서 누적기는 $2^{15}$에서 포화(saturates)하며, 이는 16비트 소수점 분수를 가진 부호 있는 32비트 고정 소수점 값의 정확한 범위입니다.

융합 계층(fused layer)의 경우, ANE는 다음을 계산합니다:

[[ y = f(\sum_k x_k w_k + b) ]] 중요하게도, 완성된 MAC 합계는 중간 메모리 왕복(memory round-trip)을 피하고 포스트-MAC 활성화 블록으로 직접 공급됩니다. 이는 활성화 함수가 점별(pointwise)이기 때문에 가능합니다. 즉, 스칼라 축소(scalar reduction)가 완료되면 그 활성화 값은 해당 스칼라에만 의존하며 즉시 적용될 수 있습니다.

ANE가 tanh()를 어떻게 구현하는지 알아보기 위해, 단일 TANH 활성화 레이어를 포함하는 CoreML 모델을 컴파일하고 결과로 생성된 컴파일된 하드웨어 레지스터 파일(hwx)을 검사합니다. 계수 영역에는 0x4288부터 시작하여 33개의 연속적인 FP16 워드가 있습니다:

00004270: 3120 3001 0000 0000 0000 0000 0000 0000
00004280: 0000 0044 0000 003c 0000 f52f d633 bc35 # 0.000000 0.124329 0.244873 0.358398
00004290: 6537 7038 1539 a239 183a 793a c93a 0a3b # 0.462158 0.554688 0.635254 0.704102 0.761719 0.809082 0.848145 0.879883
...

이 33개의 FP16 워드는 α anh(x)α의 33개 IEEE LE FP16 양자화 샘플과 일치합니다:
[ T_i=\operatorname{round}_{16}!(\tanh(i/8))], − i=0,1,…,32. !Core ML tanh overlaid with double-precision tanh, followed by signed error

이제 RELU 활성화 레이어로 전환합니다:

activation program	`NonlinearMode`	lookup coefficients
identity	0	none
ReLU	1	none
tanh	2	33 FP16 words

따라서 모드 2는 사용자 정의된 33개 항목의 조회 테이블(lookup table)을 선택합니다. 33개의 지점은 32개의 구간을 정의합니다. αR=3α일 때, 매듭(knots)은
[ x_i=\frac{i}{8},\qquad i=0,…,32,]
이며, 이는 간격이 1/8인 [0,4]를 커버합니다. 입력은 “u=2^R|x|”로 테이블에 매핑되므로, R은 매듭 간격을 설정합니다. 해상도는 33개 빈(bin)보다 부드럽습니다. 저는 인접한 항목들이 선형 보간(linearly interpolated)된다고 의심합니다. 테스트하기 위해 단일 스파이크를 가진 임펄스 LUT를 만듭니다:
[ T_8=1,\qquad T_k=0\ ext{for }k\ne 8,\qquad R=3. !One nonzero lookup-table entry produces two straight-line segments on the ANE

그런 다음 입력이 T_8 주변의 두 셀을 가로지르도록 스윕합니다. 측정된 출력은 삼각형 형태를 이룹니다: 크기가 |x|=7/8에서 0부터 시작하여 |x|=1에서 1까지 선형적으로 증가한 후, |x|=9/8에서 0까지 선형적으로 감소합니다.

따라서 모드 2는 33개의 항목으로 구성된 조각별 선형 LUT(Piecewise-linear LUT)를 구현한다는 것을 알 수 있습니다. $R$은 입력을 LUT 좌표로 스케일링하며,
[ u=2^R|x|, ] 따라서 노트 간격은 $\Delta x=2^{-R}$입니다. $\lfloor u\rfloor$와 $\lceil u\rceil$는 인접한 항목을 선택하고, $\alpha=u-\lfloor u\rfloor$는 그 사이의 보간 가중치를 제공합니다.

CoreML은 또한 선형 스케일링 및 바이어스 $ax + b$ 변환을 지원합니다. 저는 $ax + b$가 모드 2의 선형 보간 하드웨어를 공유할 수 있다고 의심했습니다. 이를 확인하기 위해, 상수 스케일과 오프셋이 적용된 ReLU를 가진 CoreML 모델을 구성했습니다:
[ z=4x-2, y=\operatorname{ReLU}\[rac{z}{2}+1],]만약 컴파일러가 상수 스케일과 오프셋을 컨볼루션에 통합(fold)한다면:
[ W'=\frac12W=2, b'=\frac12b+1=0, ][ y=\operatorname{ReLU}(2x). ]

model.espresso.weights를 디코딩하면 ReLU에 대한 정확히 이러한 통합 변환을 확인할 수 있습니다:

authored convolution: W = 4, b = -2
activation affine: s = 0.5, c = 1
compiled convolution: W' = 2, b' = 0

Core ML folds constant scale and offset into convolution weights and bias before ReLU

그리고 레지스터 파일 hexdiff는 바이어스와 활성화가 컴파일 시간에 동일한 post-MAC 경로로 어떻게 융합(fused)되는지를 보여줍니다:

Probe	Tasks	BiasMode	PostScaleMode	NonlinearMode
Plain convolution	1	0	0	0
Explicit Core ML `Bias`
1	1	0	0
Bias + ReLU	1	1	0	1
Bias + tanh	1	1	0	2

극도로 기괴한 아이디어: 비선형 보간을 사용하여 추가 커널 패스를 계산하거나, int8을 int4 가중치로 양자화하는 것.

ANE 드라이버 소스 코드는 실망스러울 정도로 지루합니다. 이 드라이버는 ANE에게 실행할 CONV, MATMUL, 또는 RELU 오피코드를 절대 제공하지 않습니다. 모든 신경 연산은 이미 작업 설명자(task descriptors, TDs)의 명령어 스트림으로 컴파일되어 있으며, 드라이버 소프트웨어는 단순히 작업을 메모리에 로드하고, (TM_ADDR, TM_SIZE)를 통해 불투명한 작업 블롭에 포인터를 설정한 다음, 도어벨을 울려 스테이징된 작업을 제출합니다.

하드웨어는 완료될 때까지 제출된 작업을 소유하며, 작업이 완료되면 ARM64 코어에 인터럽트를 발생시킵니다.

static void ane_tm_handle_irq(struct ane_device *ane)
{
int line;
...

이 (지루한) 명령어 제출 프론트엔드는 GPU의 것과 유사하며, NVIDIA의 pushbuffer/PBDMA를 생각해보세요. 소프트웨어는 메모리에 상주하는 명령어 스트림을 제출하고, GPU의 명령어 프로세서는 명령 스트림을 따라가며 명령어를 디스패치하는데, 이 과정에서 해당 명령어가 무엇을 실행할지는 알지 못합니다.

TM_ADDR과 TM_INFO는 전역 스테이징 레지스터이며, TM_PUSH는 그 스테이징된 런처 상태를 원자적으로 커밋합니다. 이는 TM_PUSH가 쓰여질 때까지 아무 일도 일어나지 않기 때문입니다 ("마법"). 특히 TM_INFO에는 공급된 스트림의 전체 디스크립터 개수가 저장됩니다:

TM_INFO[31:16] = descriptor_dwords - 1
TM_INFO[15:0] = descriptor_count

TM_INFO 레지스터는 자연스럽게 하드웨어 카운터에 매핑됩니다:

if (fetch) begin
if (word_ctr == descriptor_dwords_minus_1) begin
word_ctr <= 0;
...

왜 "-1"일까요? length - 1을 인코딩하는 것은 제로 기반 카운터를 임계 경로(critical path) 밖에서 종료시키는 RTL 친화적인 방법입니다. 하지만 GPU 명령어가 링버퍼 내의 가변 길이 디스크립터 스트림을 파싱하는 것과 비교했을 때, ANE는 전체 개수만을 받으며 디스크립터가 고정 크기라는 것을 나타내는 점에 주목하세요.

한 단계 더 깊이 들어가서, 작업 관리자(Task Manager)가 선택하는 작업 큐(TQ)에는 무엇이 있을까요?

+------------------+
CPU / 드라이버 ------>| 작업 관리자 (Task Manager) |
| |
...

동일한 레지스터 블록의 복사본이 8개 있습니다 (qid(0…7)으로 인덱싱). 구조는 다음과 같습니다:

TQ[qid] + 0x000 STATUS
0x010 PRIORITY
0x014 VACANT
...

각 TQ에는 다음이 포함됩니다:

(1) Per-TQ 스케줄링 상태 (status, priority, 그리고 vacancy)
(2) TQ당 두 세트의 명령어 스트림 디스크립터(ADDR1/ADDR2, SIZE1/SIZE2, NID1/NID2, 그리고 32 BARs). 이 두 슬롯은 한 슬롯이 실행되는 동안 소프트웨어가 다른 슬롯을 수정할 수 있도록 하는 명백한 ping-pong 스테이징 방식입니다.

TM_PUSH가 TQ_ADDR/TQ_SIZE에 참조된 작업을 qid를 첨부하여 실행하는 것을 주목하세요:

tm_write32(ane, TM_PUSH, TQ_PRTY_TABLE[qid] | (qid & 7) << 8); // magic

자연스러운 해석은 디스크립터 스트림이 어떤 작업을 실행할지 지정하고, qid가 디스크립터가 실행되는 런칭 컨텍스트를 선택한다는 것입니다. 상주하는 TQ 컨텍스트(BAR, NID)는 GPU 하드웨어 채널과 매우 유사합니다.
여기에 제가 하나의 TQ에 로드하여 실행시키는 드라이버 코드가 있습니다:

int ane_tm_enqueue(struct ane_device *ane, struct ane_request *req)
{
int qid = req->qid;
...

사후 분석을 통한 Apple Neural Engine 역공학

요약

핵심 포인트

사후 분석을 통한 Apple Neural Engine 역공학

댓글

사후 분석을 통한 Apple Neural Engine 역공학