Giải mã Apple Neural Engine: Phân tích kiến trúc M1 và lý do NPU dần thoái trào

10 tháng 8, 2026

(5089 từ)

Tôi đã ngừng làm việc với driver Apple Neural Engine (ANE) được kỹ thuật đảo ngược (reverse-engineered) từ ba năm trước, sau một nhận thức nhỏ đầy buồn bã rằng khối ANE thực sự không hữu dụng đến thế, và tôi có thể làm những việc có ích hơn, nên tôi đã chuyển sang tập trung vào việc upstream các khối khác hữu ích hơn. Kiến trúc của ANE quá cứng nhắc để có thể xây dựng một nền tảng tăng tốc đa năng xung quanh nó, và một driver Linux mở quyền truy cập API phần cứng ANE cũng không thể mở rộng các loại tác vụ mà nó có thể thực hiện. Ngay cả macOS cũng chỉ thường xuyên sử dụng ANE của riêng họ để tạo các bản xem trước hình ảnh đã được lấy mẫu lại (upsampled) trong Finder.

https://github.com/eiln/ane/tree/main

Tính năng chủ đạo của M5 (2025) là "hiệu năng LLM", và họ cũng đã khéo léo tích hợp các nhân ANE vào bên trong các nhân GPU — tôi biết điều này sẽ xảy ra, nhưng nó thực sự mang lại cảm giác như khởi đầu của sự kết thúc đối với NPU độc lập. Vì vậy, để tưởng nhớ sự ra đi rõ ràng của ANE, chúng ta sẽ làm một việc thậm chí còn vô dụng hơn: quay lại và kỹ thuật đảo ngược ANE trên M1, hoàn thành những gì chúng ta đã bắt đầu. Đã ba năm rồi (chết tiệt), và tôi nên biết nhiều hơn so với lúc mới bắt đầu làm việc này.

Nếu mục tiêu ba năm trước là làm cho ANE trở nên hữu dụng bằng cách chạy các phép toán trên đó; thì lần này, mục tiêu thiên về việc lập bản đồ toàn bộ kiến trúc nội bộ — tính toán, đường dẫn dữ liệu (datapath), bộ lập lịch, bộ nhớ và mô hình thực thi — bởi vì những quyết định thiết kế nội bộ đó tiết lộ các giả định về khối lượng công việc ML mà Apple đã sẵn sàng cam kết đưa vào silicon lần đầu tiên trên A11 Bionic (2017), và điều đó nói lên điều gì về sự chuyển dịch từ các NPU thời CNN sang các GPU chạy khối lượng công việc transformer ngày nay.

1. Tính toán (Compute)

16 nhân tính toán có lẽ là phần ít thú vị nhất của ANE. Apple ban đầu nhắm đến các khối lượng công việc CNN xử lý hình ảnh dày đặc, bao gồm các phép giảm tensor dày đặc với khả năng tái sử dụng có thể dự đoán được. Nhân tính toán ANE trên M1 là một mảng song song lớn các đơn vị nhân-tích lũy (MAC), nhưng chỉ riêng điều đó gần như không nói lên gì về việc nó được thiết kế cho các tác vụ nào và nó vượt trội ở đâu.

Một lớp tích chập (convolutional layer) thực hiện phép nhân chấm (dot product) giữa cửa sổ kích hoạt và trọng số kernel đã học, và cơ chế attention thực hiện phép nhân chấm giữa vector truy vấn (query) và khóa (key). Một phép nhân chấm là một phép nhân chấm, và một đơn vị MAC thực hiện chính xác điều đó. Điều làm nên sự chuyên biệt của ANE đối với các mô hình CNN năm 2017 không phải là MAC, mà là luồng dữ liệu (dataflow) xung quanh các MAC: khi nào và ở đâu các đầu vào và đầu ra của MAC đi vào, ở lại, di chuyển. Giả định mà các transformer đã phá vỡ, đặc biệt là với giải mã tự hồi quy (autoregressive decode), chính là các mô hình tái sử dụng có thể dự đoán được, vốn được ANE khai thác để kiến trúc một luồng dữ liệu đủ hiệu quả để chạy trên điện thoại. Quyết định trên M5 xác nhận rằng nhân tính toán của ANE vẫn hữu ích cho các transformer, nhưng nằm trong một luồng dữ liệu khác.

Tuy nhiên, đây là đường dẫn dữ liệu bên trong mỗi nhân trong số 16 nhân tính toán:

Nhân-Tích lũy (Multiply-Accumulate)

ANE có 16 nhân tính toán song song. Mỗi nhân tính toán có 128 làn nhân-tích lũy (MAC) song song FP16 (hoặc 256 INT8). Mỗi làn MAC thực hiện phép lặp:

Nhân hai toán hạng \(a\) và \(b\), sau đó cộng tích số vào tổng đang chạy (bộ tích lũy).

Lặp lại thao tác MAC qua T chu kỳ sẽ tính toán một phép nhân chấm T-số hạng:

Do đó, một làn MAC thực hiện phép giảm vô hướng (scalar reduction) theo thời gian. Một ANE 16 nhân có 2048 làn MAC song song,

Vì vậy, mỗi chu kỳ thực hiện 2048 phép giảm song song theo không gian, với thời gian là trục giảm duy nhất:

Một làn MAC riêng lẻ không biết nó đang giảm theo chiều nào của ma trận hoặc tensor. Cần lưu ý rằng phép nhân chấm so với phép nhân ma trận so với tích chập nảy sinh từ cách các toán hạng được ánh xạ và lập lịch lên nhân. Nhân ANE (ngoại trừ bộ nhớ kernel, sẽ thảo luận sau) không mã hóa một lớp CNN 4 kênh vào phần cứng.

Về mặt nội bộ, đường dẫn dữ liệu MAC bao gồm một bộ nhân, bộ cộng và một thanh ghi tích lũy 32-bit. Mỗi chu kỳ, bộ cộng thêm đầu ra mới từ bộ nhân với tổng trước đó, sau đó trở thành tổng đang chạy mới.

Đường dẫn phản hồi này giữ tổng một phần trong bộ nhớ cục bộ của làn MAC, vì vậy nó không cần phải lấy từ bộ nhớ ngoài ở xa giữa các chu kỳ MAC.

Về độ phân giải, nó thực hiện giảm điểm cố định (fixed-point) với FP16 tại thời điểm đọc. Bộ nhân là 16-bit, được tích lũy trong thanh ghi 32-bit dưới dạng Q16.16, sau đó được đọc ra dưới dạng FP16 thông qua mở rộng dấu (sign-extend) và các bước khác. Khi làm việc với biểu diễn FP16 số nguyên (hex), để thăm dò phạm vi bộ tích lũy, hãy xây dựng một chương trình CoreML ANE tính toán phép nhân chấm với một vector toàn các số (1), vì vậy mỗi bộ nhân dẫn đến một giá trị \(v\) bị chặn, nhưng tổng đang chạy trong bộ tích lũy cứ tiếp tục tăng:

Vì 32768 bản thân nó là một từ FP16 hợp lệ (0x7800), giá trị 0x7c00 của ANE không thể là tràn đầu ra FP16, việc kẹp (clamp) xảy ra bên trong bộ tích lũy, tại \(2^{15}\). Do đó, bộ tích lũy bão hòa tại \(2^{15}\), chính xác là phạm vi của một giá trị điểm cố định 32-bit có dấu với 16 bit phân số.

Kích hoạt phi tuyến tính (Nonlinear Activation)

Đối với một lớp được hợp nhất (fused layer), ANE tính toán:

Quan trọng là, các tổng MAC đã hoàn thành được đưa trực tiếp vào khối kích hoạt sau MAC, tránh việc phải quay lại bộ nhớ trung gian. Điều này khả thi vì kích hoạt là theo điểm (pointwise): một khi phép giảm vô hướng hoàn tất, kích hoạt của nó chỉ phụ thuộc vào vô hướng đó và có thể được áp dụng ngay lập tức.

Để xác định cách ANE thực hiện tanh, hãy biên dịch một mô hình CoreML chứa một lớp kích hoạt TANH duy nhất và kiểm tra tệp thanh ghi phần cứng (hwx) đã biên dịch. Vùng hệ số chứa 33 từ FP16 liên tiếp bắt đầu tại 0x4288:

33 từ FP16 đó khớp với 33 mẫu định lượng IEEE LE FP16 của \(\tanh(x)\):

Bây giờ chuyển sang lớp kích hoạt RELU:

Như vậy, chế độ 2 chọn một bảng tra cứu (lookup table) tùy chỉnh gồm 33 mục. 33 điểm xác định 32 khoảng. Với \(R=3\), các nút là

bao phủ \([0,4]\) với khoảng cách \(1/8\). Đầu vào ánh xạ vào bảng dưới dạng \(u=2^R|x|\), vì vậy \(R\) thiết lập khoảng cách nút. Độ phân giải mượt mà hơn so với 33 bin của nó; tôi nghi ngờ rằng các mục liền kề được nội suy tuyến tính. Để kiểm tra, hãy xây dựng một LUT xung với một gai đơn lẻ:

Sau đó quét đầu vào qua hai ô xung quanh \(T_8\). Đầu ra đo được tạo thành một hình tam giác: độ lớn tăng tuyến tính từ \(0\) tại \(|x|=7/8\) lên \(1\) tại \(|x|=1\), sau đó giảm tuyến tính xuống \(0\) tại \(|x|=9/8\).

Như vậy, chúng ta biết rằng chế độ 2 thực hiện một LUT tuyến tính từng đoạn gồm 33 mục. \(R\) chia tỷ lệ đầu vào thành tọa độ LUT,

vì vậy khoảng cách nút là \(\Delta x=2^{-R}\). \(\lfloor u\rfloor\) và \(\lceil u\rceil\) chọn các mục liền kề, và \(\alpha=u-\lfloor u\rfloor\) cung cấp trọng số nội suy giữa chúng.

Chia tỷ lệ và Độ lệch (Scaling and Bias)

CoreML cũng hỗ trợ phép biến đổi chia tỷ lệ tuyến tính và độ lệch \(ax + b\). Sau đó, tôi nghi ngờ \(ax + b\) có thể chia sẻ phần cứng nội suy tuyến tính của chế độ 2. Để xác nhận, hãy xây dựng một mô hình CoreML với ReLU có tỷ lệ và độ lệch không đổi:

Nếu trình biên dịch hợp nhất tỷ lệ và độ lệch không đổi vào phép tích chập:

Giải mã model.espresso.weights xác nhận chính xác phép biến đổi được hợp nhất này trên ReLU:

Và tệp hexdiff của thanh ghi cho thấy cách độ lệch và kích hoạt được hợp nhất vào cùng một đường dẫn sau MAC tại thời điểm biên dịch:

Một ý tưởng cực kỳ "nguyền rủa": sử dụng nội suy phi tuyến tính để tính toán thêm một lượt kernel, hoặc định lượng trọng số int8 thành int4.

2. Bộ lập lịch (Scheduler)

Mã nguồn của driver ANE gây thất vọng vì quá nhàm chán. Driver không bao giờ cung cấp cho ANE một opcode CONV, MATMUL hoặc RELU để chạy. Tất cả các thao tác thần kinh đều đã được biên dịch thành một luồng lệnh gồm các bộ mô tả tác vụ (task descriptors - TDs), và phần mềm driver chỉ đơn giản là tải tác vụ vào bộ nhớ, đặt con trỏ đến blob tác vụ ẩn thông qua (TM_ADDR, TM_SIZE), và gửi tác vụ đã dàn dựng bằng cách rung chuông (TM_PUSH).