> ## Content Index
> Fetch the complete content index at: https://blog.hacking.pub/llms.txt
> Use this file to discover other available public pages before exploring further.

# 用一个矩阵乘法看懂 ANE、GPU 与统一内存
- URL: https://blog.hacking.pub/2026/10/25/001/
- Published: 2026-10-25T16:34:34.000Z
- Updated: 2026-10-25T17:10:30.000Z
- Author: chainhelen

最近看 Kapitsa 的 ANE 实现时，我想先弄明白一个很小的问题：如果只计算两个矩阵相乘，数据究竟怎样进入 Apple Neural Engine，答案又从哪里回来？

于是把例子收进一个 `main.mm` 文件：输入是两个 2×2 矩阵，计算图内嵌在文件里，乘法交给 ANE，最后打印结果。本文先把硬件和内存关系画清楚，再顺着这个例子读代码。

## 1\. ANE 是什么，它和 GPU 有什么区别？

**ANE（Apple Neural Engine）是苹果的神经网络加速器，可以归入 NPU。它与 GPU 是并列的计算单元，不是 GPU 内部的一种核心。**

它们都能做大量乘加运算，但面向的软件接口和硬件取舍不同：

| 计算单元 | 在这里怎样理解         | 常见工作                  |
| ---- | --------------- | --------------------- |
| CPU  | 运行程序、组织任务       | 控制流程、准备数据、提交计算        |
| GPU  | 可编程的并行计算单元      | 图形渲染、科学计算、神经网络训练与推理   |
| ANE  | 面向神经网络计算图的专用加速器 | 对支持的算子、形状和数据布局执行高能效计算 |

GPU 可以通过 Metal 编写计算程序；ANE 通常要先把计算图编译成它能执行的形式。不能把一段普通 C++ 函数直接丢给 ANE，让它代替 CPU 执行。[Apple 的 ANE 文章](https://machinelearning.apple.com/research/neural-engine-transformers?ref=blog.hacking.pub)讨论了这种面向算子和内存布局的优化，[Metal](https://developer.apple.com/metal/?ref=blog.hacking.pub)则提供图形和通用计算接口。

“都能用于神经网络”也不等于“都是为训练设计的”。训练还包括计算梯度和更新权重；本文只有一次矩阵乘法，没有训练过程。GPU 用于训练很常见，ANE 更常见于推理加速，但不能由此推导出 ANE 在硬件上绝对不能参与训练。

## 2\. 一张图看清统一内存

![CPU、ANE、GPU 与统一内存：本例 CPU 提交任务，ANE 读写 IOSurface，GPU 未使用](https://blog.hacking.pub/qiniu/content/images/2026/10/ane-unified-memory.png)

先只看图中的蓝色 CPU、绿色 ANE，以及下面的统一内存：

1. CPU 把输入矩阵放进输入缓冲区。
2. CPU 提交任务，告诉运行时使用哪些输入、把结果写到哪里。
3. ANE 读取输入，执行乘法，把答案写入输出缓冲区。
4. 等计算完成后，CPU 读取输出缓冲区，打印答案。

图里的 GPU 用来说明硬件关系，本例没有调用它。

**统一内存是这些计算单元共享的 DRAM 内存池，不是 ANE 或 GPU 内部的小容量缓存。**

在 Apple Silicon 上，不必把同一份数据分别放进传统独立显卡式的主存和显存池。这个共享内存池的设计可参考 [Apple 对 M1 统一内存架构的说明](https://www.apple.com/newsroom/2020/11/apple-unleashes-m1/?ref=blog.hacking.pub)。

但“共享同一份数据”仍然需要读取、计算和写回，也仍然受内存带宽影响。共享内存不保证任意指针都能直接交给硬件：缓冲对象、数据类型、布局和同步方式还要符合接口要求。

## 3\. 这个例子到底算什么？

输入和预期输出只有这么大：

```text
     A             B             C
┌       ┐     ┌       ┐     ┌         ┐
│ 1   2 │  ×  │ 5   6 │  =  │ 19   22 │
│ 3   4 │     │ 7   8 │     │ 43   50 │
└       ┘     └       ┘     └         ┘

C[0][0] = 1×5 + 2×7 = 19
C[0][1] = 1×6 + 2×8 = 22
C[1][0] = 3×5 + 4×7 = 43
C[1][1] = 3×6 + 4×8 = 50

```

为使用这份已经验证过的计算图，例子把两个小矩阵分别放到 32×32 缓冲区的左上角，其余位置补零。ANE 实际执行 32×32 的矩阵乘法，我们只取结果的左上角 2×2。

**32×32 是本例的布局选择，不代表所有 ANE 矩阵乘法都必须用这个尺寸，也不是把矩阵当图片缩放。**

补零后，左上角的四个结果仍然是上面的 19、22、43、50。

## 4\. 一个文件里有两种角色

`main.mm` 的 `.mm` 表示 Objective-C++ 源文件，可以混合使用 C++ 和 Objective-C 语法。它编译出的主程序在 CPU 上运行。

文件里的 `kMilSource` 是一段字符串，内容是 MIL 计算图。这段图描述“两个输入做矩阵乘法，然后输出”。程序运行时把它写到缓存目录中的 `model.mil`，再通过 ANE 私有框架编译和加载。

因此，只有一个源文件需要保存，但运行时仍会生成计算图文件。CPU 执行的是主程序，ANE 执行的是编译后的图。

从阅读角度，可以按下面的顺序理解：

```text
main：准备 A、B
  ↓
ane_multiply：本例的封装入口
  ↓
run：编译并加载计算图 → 准备缓冲区 → 绑定输入输出
  ↓
evaluateWithModel：提交并等待 ANE 执行
  ↓
读取输出缓冲区 → 取左上角 2×2 → 打印

```

代码看起来长，是因为它还包含私有接口检查、资源释放和错误处理。真正描述乘法的地方，是 MIL 里的 `matmul(...)`。

## 5\. 为什么要把矩阵复制进 IOSurface？

本例的 ANE 请求通过 IOSurface 对象引用输入输出。IOSurface 是可用于共享数据的缓冲对象；CPU 可以取得它的内存地址，ANE 请求也能引用对应缓冲区。接口背景见 [Apple IOSurface 文档](https://developer.apple.com/documentation/iosurface?ref=blog.hacking.pub)。

普通的 `float a[2][2]` 在本例里不能直接替代它：类型是 FP32，大小是 2×2，而计算图要的是 FP16、32×32 的特定布局。因此需要这一步转换：

```cpp
s->data()[row*32+col] = (_Float16)input[row][col];

```

`s->data()` 取得缓冲区首地址；`row*32+col` 找到对应元素；赋值把输入值转成 FP16 并写进去。第一行从下标 0 开始，第二行从下标 32 开始：

```text
原始 A（2×2）            输入 IOSurface（32×32，示意）
                        第 0 行：[1  2  0  0  ...  0]
[1  2]       写入 →      第 1 行：[3  4  0  0  ...  0]
[3  4]                  第 2 行：[0  0  0  0  ...  0]
                        ...

实际写入的下标：0、1、32、33

```

一块缓冲区的逻辑数据是 `32×32×2 = 2048` 字节，代码申请的 IOSurface 分配大小则是 16 KiB。两者不能混为一谈。

如果上游一开始就按要求把数据写进这块 IOSurface，就有机会省掉从另一个数组复制过来的步骤。但本例保留普通数组作为输入，便于修改和理解，所以它不是一份“全程零复制”的实现。

## 6\. request 是任务对象，不是数组

最容易误读的是 Objective-C 的方括号：

```objc
[对象 方法名:参数]

```

这是方法调用。数组字面量则写成 `@[元素1, 元素2]`，前面有一个 `@`。

例子里的 `request` 是创建好的任务对象；`id` 表示 Objective-C 对象引用。它记录输入、输出和程序入口，本身还没有执行乘法。

下面是一整个方法的名字，冒号也是名字的一部分：

```text
requestWithInputs:inputIndices:outputs:outputIndices:weightsBuffer:perfStats:procedureIndex:sharedEvents:transactionHandle:

```

调用时，每个冒号后面放对应参数。`inputIndices:` 是这个方法名中的一段，不是参数类型，也不是另外一次调用。完整源码中保留了每个参数的注释。

而 `(Class<DemoRequest>)NSClassFromString(@"_ANERequest")` 是先根据名字找到类，再作类型转换，让编译器知道它有哪些方法；方法调用的接收者就是这个类对象。

输入编号也不是凭空增加的仪式：程序需要把编译后的图中名为 `a`、`b` 的输入，对应到正确的两块缓冲区，所以查询编译结果中的符号编号，再建立绑定。

## 7\. 答案是在哪里出现的？

提交请求之前，输出缓冲区就已经分配好了，并且绑定到了请求里。执行这句时，ANE 才开始按图计算：

```objc
check([client evaluateWithModel:model options:@{} request:request qos:QOS_CLASS_DEFAULT error:&error], "ANE evaluation");

```

这个例子使用同步执行接口。成功返回后，程序锁定输出缓冲区，读取 ANE 写入的 FP16 数值并转成 `float`：

```cpp
for (int i=0;i<1024;i++) result[i]=float(output.data()[i]);

```

最后才是之前容易让人困惑的这句：

```cpp
answer[row][col] = result[row*32+col];

```

**这里没有算乘法。它只是把已经算好的答案，从 32×32 结果里取出来，复制到调用者的 2×2 数组。**

完整过程是：提前指定输出位置 → ANE 写入答案 → CPU 读取答案。指针负责定位，计算发生在前面的执行步骤。

## 8\. 编译、运行与验证范围

把文末代码完整保存为 `main.mm`，在有 Xcode Command Line Tools 的 Apple Silicon Mac 上执行：

```sh
xcrun clang++ -std=c++20 -O2 -fobjc-arc -Wall -Wextra main.mm \
  -framework Foundation -framework IOSurface -o ane-demo
./ane-demo

```

本例已在 **Apple M4 Pro、macOS 15.7.1（24G231）** 上编译运行，并从一个空目录启动生成的可执行文件，确认不依赖旁边另放的 `.mil` 或其他源文件。运行结果的最后几行是：

```text
Compile MIL ...
Evaluate on ANE ...
结果：
[19 22]
[43 50]

```

这验证的是调用链和默认输入的结果，没有做性能对比，也没有采集硬件计数器。两个 2×2 矩阵的计算量太小，不能拿这个例子论证 ANE 比 CPU 或 GPU 快。

代码通过私有 `AppleNeuralEngine.framework` 调用 ANE，没有编写 CPU/GPU 乘法回退路径；准备数据、类型转换和打印仍由 CPU 完成。私有接口会随系统版本变化，运行前的 ABI 检查只能检查所列方法的签名，不能保证所有系统都兼容。

## 9\. 回头看 Kapitsa 的做法

这个小例子帮助理解的是 ANE 调用链。实际推理引擎还需要把它接到 GPU 计算流程中。

以这里参考的 [Kapitsa ANE 实现](https://github.com/nace-ai/kapitsa/blob/f8e3c24274be4fa3ce79d42ddf5ecc97e63acccb/src/ane/ane.mm?ref=blog.hacking.pub)为例，它可以把 IOSurface 的内存包装成 GPU 可访问的缓冲区，让 GPU 准备 ANE 输入，再消费 ANE 输出。在支持的 [FFN prefill 路径](https://github.com/nace-ai/kapitsa/blob/f8e3c24274be4fa3ce79d42ddf5ecc97e63acccb/src/ane/ffn%5Fsplit.h?ref=blog.hacking.pub)中，这样可以减少为了衔接两个计算单元而经过 CPU 中转的工作。

真正值得优化的，是让前一步直接产出后一步能消费的数据布局，并复用缓冲区。统一内存提供了共享数据的基础，算子布局、数据转换和同步仍要由程序设计处理。

## 附：完整单文件示例

下面保留已验证的 `main.mm` 原文。建议先读 `main`、`kMilSource` 和 `run` 的五个编号步骤，再看接口签名与异常处理。私有接口声明参考 Kapitsa，源码注释保留了对原始接口项目 incoai/splash 的说明。

```objc
// 单文件 ANE 矩阵乘法例子：输入、运行层和 MIL 计算图都在本文件。
// 编译：xcrun clang++ -std=c++20 -O2 -fobjc-arc -Wall -Wextra main.mm -framework Foundation -framework IOSurface -o ane-demo
// 运行：./ane-demo
// 阅读顺序：main（输入和输出）→ kMilSource（计算方法）→ run（准备、执行、读回）。
// Standalone ANE exercise. No Metal or Core ML dependency and no CPU/GPU inference fallback.
// Private selectors/type encodings follow Kapitsa src/ane/ane.mm (original interface: incoai/splash).
#import <Foundation/Foundation.h>
#import <IOSurface/IOSurface.h>
#import <objc/runtime.h>
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <dlfcn.h>
#include <fcntl.h>
#include <stdexcept>
#include <string>
#include <sys/file.h>
#include <sys/qos.h>
#include <unistd.h>
#include <exception>

// 本例自己的封装函数，定义在文件末尾。
void ane_multiply(const float a[2][2], const float b[2][2], float output[2][2]);

int main() {
  // ① 准备两个矩阵；你可以先改这里的数字。
  const float a[2][2] = {
    {1, 2},
    {3, 4}
  };
  const float b[2][2] = {
    {5, 6},
    {7, 8}
  };
  float result[2][2] = {};

  // ② 让 ANE 按下方 kMilSource 描述的算法计算，把答案写进 result。
  try {
    ane_multiply(a, b, result);
  } catch (const std::exception& error) {
    std::fprintf(stderr, "计算失败：%s\n", error.what());
    return 1;
  }

  // ③ 打印答案。
  std::printf("结果：\n[%g %g]\n[%g %g]\n",
              result[0][0], result[0][1], result[1][0], result[1][1]);
}

// 计算图内嵌为字符串；运行时写到独立缓存目录，供 ANE 编译器加载。
// 改这个字符串后需重新编译 main.mm。
static constexpr char kMilSource[] = R"MIL(program(1.3)
{
    func main<ios18>(
        tensor_buffer<fp16, shape=[1, 1, 32, 32], strides=[1024, 1024, 32, 1], interleave_factors=[1, 1, 1, 1]> a,
        tensor_buffer<fp16, shape=[1, 1, 32, 32], strides=[1024, 1024, 32, 1], interleave_factors=[1, 1, 1, 1]> b
    ) {
        tensor<fp16, [1, 1, 32, 32]> at = tensor_buffer_to_tensor<ios17>(input = a);
        tensor<fp16, [1, 1, 32, 32]> bt = tensor_buffer_to_tensor<ios17>(input = b);
        tensor<fp16, [1, 1, 32, 32]> product = matmul(transpose_x = bool(false), transpose_y = bool(false), x = at, y = bt);
        tensor_buffer<fp16, shape=[1, 1, 32, 32], strides=[1024, 1024, 32, 1], interleave_factors=[1, 1, 1, 1]> output = tensor_to_tensor_buffer<ios17>(input = product, interleave_factors = tensor<uint8, [4]>([1, 1, 1, 1]), strides = tensor<int64, [4]>([1024, 1024, 32, 1]));
    } -> (output);
}
)MIL";

@protocol DemoModel
+ (id)modelAtURL:(NSURL*)url key:(NSString*)key;
- (NSDictionary*)modelAttributes;
@end
@protocol DemoClient
+ (id)sharedConnection;
- (BOOL)compileModel:(id)model options:(NSDictionary*)options qos:(unsigned)qos error:(NSError**)error;
- (BOOL)loadModel:(id)model options:(NSDictionary*)options qos:(unsigned)qos error:(NSError**)error;
- (BOOL)unloadModel:(id)model options:(NSDictionary*)options qos:(unsigned)qos error:(NSError**)error;
- (BOOL)evaluateWithModel:(id)model options:(NSDictionary*)options request:(id)request qos:(unsigned)qos error:(NSError**)error;
@end
@protocol DemoSurface
+ (id)objectWithIOSurface:(IOSurfaceRef)surface;
@end
@protocol DemoRequest
+ (id)requestWithInputs:(NSArray*)inputs inputIndices:(NSArray*)inputIndices outputs:(NSArray*)outputs
          outputIndices:(NSArray*)outputIndices weightsBuffer:(id)weights perfStats:(id)stats
         procedureIndex:(NSNumber*)procedure sharedEvents:(id)events transactionHandle:(NSNumber*)transaction;
@end

static void require(bool ok, const std::string& message) {
  if (!ok) throw std::runtime_error(message);
}
static std::string error_text(NSError* error) {
  return error ? error.description.UTF8String : "no error details";
}

// 1. 加载苹果的私有框架，确认所需方法及参数类型与本机系统匹配。
static void check_api() {
  void* framework = dlopen("/System/Library/PrivateFrameworks/AppleNeuralEngine.framework/AppleNeuralEngine", RTLD_NOW);
  require(framework, "Cannot load AppleNeuralEngine.framework");
  struct Signature { const char* cls; bool is_class; const char* selector; const char* types; };
  const Signature signatures[] = {
      {"_ANEModel", true, "modelAtURL:key:", "@@:@@"},
      {"_ANEModel", false, "modelAttributes", "@@:"},
      {"_ANEClient", true, "sharedConnection", "@@:"},
      {"_ANEClient", false, "compileModel:options:qos:error:", "B@:@@I^@"},
      {"_ANEClient", false, "loadModel:options:qos:error:", "B@:@@I^@"},
      {"_ANEClient", false, "unloadModel:options:qos:error:", "B@:@@I^@"},
      {"_ANEClient", false, "evaluateWithModel:options:request:qos:error:", "B@:@@@I^@"},
      {"_ANEIOSurfaceObject", true, "objectWithIOSurface:", "@@:^{__IOSurface=}"},
      {"_ANERequest", true, "requestWithInputs:inputIndices:outputs:outputIndices:weightsBuffer:perfStats:procedureIndex:sharedEvents:transactionHandle:", "@@:@@@@@@@@@"},
  };
  for (const auto& s : signatures) {
    Class cls = NSClassFromString(@(s.cls));
    SEL sel = sel_registerName(s.selector);
    Method method = cls ? (s.is_class ? class_getClassMethod(cls, sel) : class_getInstanceMethod(cls, sel)) : nullptr;
    require(method, std::string("Missing private selector: ") + s.selector);
    std::string actual = method_getTypeEncoding(method);
    std::erase_if(actual, [](char c) { return c >= '0' && c <= '9'; });
    require(actual == s.types, std::string("Private ABI changed: ") + s.selector + " " + actual);
  }
  puts("Private ABI: 9 signatures match.");
}

// 2. 共享缓冲区：CPU 可通过 data() 读写，ANE 请求通过 ref 引用同一块缓冲。
struct Surface {
  IOSurfaceRef ref;
  Surface() {
    ref = IOSurfaceCreate((__bridge CFDictionaryRef)@{
      (id)kIOSurfaceWidth:@32, (id)kIOSurfaceHeight:@32,
      (id)kIOSurfaceBytesPerElement:@2, (id)kIOSurfaceBytesPerRow:@64,
      (id)kIOSurfaceAllocSize:@16384, (id)kIOSurfacePixelFormat:@0x4c303068
    });
    require(ref, "IOSurfaceCreate failed");
    if (IOSurfaceGetBytesPerRow(ref) != 64 || IOSurfaceGetAllocSize(ref) < 2048) {
      CFRelease(ref);
      throw std::runtime_error("Unexpected IOSurface layout");
    }
  }
  ~Surface() { CFRelease(ref); }
  Surface(const Surface&) = delete;
  _Float16* data() { return static_cast<_Float16*>(IOSurfaceGetBaseAddress(ref)); }
};

static void run(const float input_a[2][2], const float input_b[2][2], float answer[2][2]) {
  // Share Kapitsa's model-run lock; return immediately if another benchmark holds it.
  const char* lock_path = getenv("BENCH_GPU_LOCK");
  int lock_fd = open(lock_path ? lock_path : "/tmp/kapitsa-gpu.lock", O_CREAT | O_RDWR, 0600);
  require(lock_fd >= 0, "Cannot open shared model-run lock");
  struct Close { int fd; ~Close() { close(fd); } } close_lock{lock_fd};
  require(flock(lock_fd, LOCK_EX | LOCK_NB) == 0, "Another model run holds the lock; retry later.");

  // 3. 将内嵌 MIL 写入新缓存目录，并编译、加载；新 key 避免误用旧计算图。
  NSString* key = NSUUID.UUID.UUIDString;
  NSString* cache = [NSFileManager.defaultManager URLsForDirectory:NSCachesDirectory inDomains:NSUserDomainMask].firstObject.path;
  NSString* directory = [[cache stringByAppendingPathComponent:@"ane-handwritten-example"] stringByAppendingPathComponent:key];
  NSError* error = nil;
  auto check = [&](BOOL ok, const char* action) { require(ok, std::string(action) + ": " + error_text(error)); };
  check([NSFileManager.defaultManager createDirectoryAtPath:directory withIntermediateDirectories:YES attributes:nil error:&error], "Create directory");
  NSString* mil = [NSString stringWithUTF8String:kMilSource];
  check([mil writeToFile:[directory stringByAppendingPathComponent:@"model.mil"]
              atomically:YES encoding:NSUTF8StringEncoding error:&error], "Write embedded MIL");
  printf("Program directory: %s\n", directory.UTF8String);
  id<DemoModel> model = [(Class<DemoModel>)NSClassFromString(@"_ANEModel") modelAtURL:[NSURL fileURLWithPath:directory isDirectory:YES] key:key];
  id<DemoClient> client = [(Class<DemoClient>)NSClassFromString(@"_ANEClient") sharedConnection];
  require(model && client, "Cannot create ANE model/client");
  NSDictionary* options = @{@"kANEFModelType":@"kANEFModelMIL", @"kANEFNetPlistFilenameKey":@"model.mil"};
  puts("Compile MIL ...");
  check([client compileModel:model options:options qos:QOS_CLASS_DEFAULT error:&error], "ANE compile");
  check([client loadModel:model options:@{} qos:QOS_CLASS_DEFAULT error:&error], "ANE load");
  @try {
    try {
      Surface a, b, output;
      for (Surface* s : {&a, &b, &output}) {
        require(IOSurfaceLock(s->ref, 0, nullptr) == kIOReturnSuccess, "IOSurfaceLock failed");
        std::memset(s->data(), 0, 2048);
        if (s == &a || s == &b) {
          const float (*input)[2] = s == &a ? input_a : input_b;
          for (int row=0;row<2;row++)
            for (int col=0;col<2;col++) s->data()[row*32+col] = (_Float16)input[row][col];
        }
        if (s == &output) for (int i=0;i<1024;i++) s->data()[i] = (_Float16)NAN;
        require(IOSurfaceUnlock(s->ref, 0, nullptr) == kIOReturnSuccess, "IOSurfaceUnlock failed");
      }

      // 4. 查询编译后的输入编号，将 MIL 中的 a、b 对应到我们的两个输入缓冲区。
      NSDictionary* description = [model modelAttributes][@"ANEFModelDescription"];
      NSNumber* procedure = description[@"kANEFModelProcedureNameToIDMapKey"][@"main"];
      require(procedure, "Compiled program has no main procedure");
      NSDictionary* entry = nil;
      for (NSDictionary* e in description[@"ANEFModelProcedures"])
        if ([e[@"ANEFModelProcedureID"] isEqual:procedure]) entry = e;
      NSArray* symbols = description[@"kANEFModelInputSymbolsArrayKey"];
      NSArray* indices = entry[@"ANEFModelInputSymbolIndexArray"];
      NSArray* out_indices = entry[@"ANEFModelOutputSymbolIndexArray"];
      require(indices.count == 2 && out_indices.count == 1, "Expected two inputs and one output");
      Class<DemoSurface> sc = (Class<DemoSurface>)NSClassFromString(@"_ANEIOSurfaceObject");
      NSMutableArray* inputs = [NSMutableArray array];
      for (NSNumber* index in indices) {
        NSString* name = symbols[index.unsignedIntegerValue];
        require([name isEqual:@"a"] || [name isEqual:@"b"], "Unexpected input symbol");
        [inputs addObject:[sc objectWithIOSurface:[name isEqual:@"a"] ? a.ref : b.ref]];
      }
      // 外层 [...] 是一次类方法调用；冒号前的片段共同组成方法名，冒号后是参数值。
      // (Class<DemoRequest>) 是类型转换，NSClassFromString(...) 找到被调用的类。
      id request = [(Class<DemoRequest>)NSClassFromString(@"_ANERequest")
          requestWithInputs:inputs                         // 1. NSArray*：输入缓冲区列表
          inputIndices:indices                             // 2. NSArray*：输入编号列表
          outputs:@[[sc objectWithIOSurface:output.ref]]    // 3. NSArray*：包装输出缓冲区，再放入数组
          outputIndices:out_indices                        // 4. NSArray*：输出编号列表
          weightsBuffer:nil                                // 5. id：额外权重对象，本例留空
          perfStats:nil                                    // 6. id：性能统计对象，本例留空
          procedureIndex:procedure                         // 7. NSNumber*：程序入口编号
          sharedEvents:nil                                 // 8. id：同步事件对象，本例留空
          transactionHandle:nil];                          // 9. NSNumber*：事务句柄，本例留空
      require(request, "Cannot construct ANE request");

      // 5. 提交并等待 ANE 执行；成功返回后，output 缓冲区里才有计算结果。
      puts("Evaluate on ANE ...");
      check([client evaluateWithModel:model options:@{} request:request qos:QOS_CLASS_DEFAULT error:&error], "ANE evaluation");
      require(IOSurfaceLock(output.ref, kIOSurfaceLockReadOnly, nullptr) == kIOReturnSuccess, "Output lock failed");
      // 读取 ANE 刚写入的 fp16 输出，转为 float；每行有 32 个数。
      float result[1024];
      for (int i=0;i<1024;i++) result[i]=float(output.data()[i]);
      IOSurfaceUnlock(output.ref, kIOSurfaceLockReadOnly, nullptr);
      for (float value : result) require(std::isfinite(value), "ANE produced an unwritten or non-finite element");
      // 取左上角 2×2，写回调用者的结果数组；这里仅复制结果，不执行乘法。
      for (int row=0;row<2;row++)
        for (int col=0;col<2;col++) answer[row][col] = result[row*32+col];
    } catch (const std::exception& e) {
      @throw [NSException exceptionWithName:@"DemoError" reason:[NSString stringWithUTF8String:e.what()] userInfo:nil];
    }
  } @finally {
    [client unloadModel:model options:@{} qos:QOS_CLASS_DEFAULT error:&error];
  }
}

void ane_multiply(const float a[2][2], const float b[2][2], float output[2][2]) {
  setvbuf(stdout, nullptr, _IONBF, 0);
  // Bound a private-service hang; the OS closes the process's lock on exit.
  struct Deadline { Deadline() { alarm(60); } ~Deadline() { alarm(0); } } deadline;
  @autoreleasepool {
    @try {
      check_api();
      run(a, b, output);
    } @catch (NSException* e) { throw std::runtime_error(e.description.UTF8String); }
  }
}

```

## 补充：Kapitsa 怎样处理更大的矩阵与更长的输入？

上面的 32×32 是演示程序选定的计算图尺寸。Kapitsa 会按模型生成自己的 MIL 程序，并结合 token 分批、图内矩阵分段和 GPU / ANE 分工来执行。

下面只讨论 **Qwen 稠密模型在 prefill（读取输入阶段）的 FFN ANE 路径**，不代表整个模型都在 ANE 上运行。代码固定到 2026-10-26 核对的提交 `919b1b2`，便于以后对照。

![kapitsa-ane-ffn-flow.png](https://blog.hacking.pub/qiniu/content/images/2026/10/kapitsa-ane-ffn-flow.png)

图中展示的是单路批次的概念流程。实际实现还可以把一个批次拆成两部分，跨层交错安排 GPU 和 ANE，后面会说明。

### 先区分两种“大小”

一个批次的输入可以理解为“若干个 token，每个 token 有若干个特征”。token 数量和特征维度是不同的轴：

| 遇到的问题              | Kapitsa 在这条路径中的做法  |
| ------------------ | ------------------ |
| 输入 token 太多        | 上层引擎分批处理           |
| 当前 token 数与入口规格不一致 | 选择能容纳它的入口，多余位置补零   |
| 矩阵的特征维度很大          | 在计算图内分段相乘，再把部分结果相加 |

### 1\. 提前准备多种计算入口，调用时选合适的

当最小批次设为 512、且没有手动指定入口列表时，代码每隔 128 个 token 准备一个入口，直到 2048：

```text
ffn512  → 处理 512 个 token 的规格
ffn640  → 处理 640 个 token 的规格
ffn768  → 处理 768 个 token 的规格
……
ffn2048 → 处理 2048 个 token 的规格

```

这些入口放在同一个 MIL 程序中，初始化时编译或从缓存加载，后续复用。一次提交有 600 个 token 时，入口选择函数会找到 `ffn640`，GPU 打包输入时把多出的 40 个位置补零。这里的 512、640 是 token 轴的长度，并不是 512×512、640×640 的方阵。

源码：[生成入口规格](https://github.com/nace-ai/kapitsa/blob/919b1b2a6782ec205e23725dee76ee5f57de2857/src/ane/ffn%5Fsplit.cpp?ref=blog.hacking.pub#L101)、[编译程序并绑定缓冲区](https://github.com/nace-ai/kapitsa/blob/919b1b2a6782ec205e23725dee76ee5f57de2857/src/ane/ffn%5Fsplit.cpp?ref=blog.hacking.pub#L188)、[选择入口](https://github.com/nace-ai/kapitsa/blob/919b1b2a6782ec205e23725dee76ee5f57de2857/src/ane/ffn%5Fsplit.cpp?ref=blog.hacking.pub#L525)、[打包时补零](https://github.com/nace-ai/kapitsa/blob/919b1b2a6782ec205e23725dee76ee5f57de2857/src/kernels/ane%5Fffn.metal?ref=blog.hacking.pub#L101)。

### 2\. token 超过 2048，由上层分批

例如没有命中前缀缓存、需要处理 5000 个 token，且每批设置为 2048 时，基本分批方式是：

```text
5000 个 token → 第一批 2048 + 第二批 2048 + 第三批 904

```

这里的 2048 是这条 ANE 路径接受的批次上限，不是模型的上下文上限。分批由上层引擎组织，ANE 接口不会自动把任意长输入切好。假如配置使某个批次本身超过 2048，这个批次会走 GPU；太小、不满足启用条件或 ANE 尚未就绪的批次，也可能走 GPU。

源码：[上层分批规则](https://github.com/nace-ai/kapitsa/blob/919b1b2a6782ec205e23725dee76ee5f57de2857/src/engine/engine.cpp?ref=blog.hacking.pub#L680)、[ANE 批次条件](https://github.com/nace-ai/kapitsa/blob/919b1b2a6782ec205e23725dee76ee5f57de2857/src/ane/ffn%5Fsplit.cpp?ref=blog.hacking.pub#L525)、[超过 2048 时使用 GPU](https://github.com/nace-ai/kapitsa/blob/919b1b2a6782ec205e23725dee76ee5f57de2857/src/model/qwen35%5Fane.cpp?ref=blog.hacking.pub#L216)。

### 3\. 特征维度很大，在图内拆成几次乘法

以隐藏维度为 5120 的模型为例，gate / up 投影会沿这个维度拆成两段，每段 2560。省略量化与缩放细节，核心关系是：

```text
W = [ W0  W1 ]       X = [ X0 ]
                         [ X1 ]

W × X = W0 × X0 + W1 × X1

```

其中 W0、W1 各有 2560 列，X0、X1 各有 2560 行。两次乘法与相加都写在 ANE 计算图里，不需要 CPU 把每块结果取回来再相加。down 投影也会沿中间通道分段后求和。

这与“把 5000 个 token 分成几批”是两回事：前者拆的是特征维度，后者拆的是 token 维度。2560 是这份实现的分段选择；代码还会检查模型形状是否支持，不能理解成任意形状都能直接套用。

源码：[模型形状检查与分段](https://github.com/nace-ai/kapitsa/blob/919b1b2a6782ec205e23725dee76ee5f57de2857/src/ane/ffn%5Fsplit.cpp?ref=blog.hacking.pub#L79)、[生成 matmul 与求和节点](https://github.com/nace-ai/kapitsa/blob/919b1b2a6782ec205e23725dee76ee5f57de2857/src/ane/ffn%5Fsplit.cpp?ref=blog.hacking.pub#L363)。

### 4\. GPU 准备数据，ANE 计算自己的份额，GPU 合并结果

Kapitsa 按中间通道分配 FFN 工作：ANE 负责分配到的部分，GPU 负责其余部分；具体比例由芯片、模型、校准结果或配置决定，某些配置也可能把全部 FFN 通道交给 ANE。

GPU 将输入转换成 ANE 需要的布局并量化，写进共享 IOSurface。ANE 读取这些输入，执行编译后的图，把结果写入共享输出。GPU 等待完成信号后读取输出，应用缩放并合并结果。共享内存减少了经 CPU 中转的需要，但量化、重排、DRAM 读写和同步仍然存在。

实际代码默认启用的 wavefront 调度，还会在满足条件时把一个 token 批次拆成前后两部分，让一部分的 ANE FFN 与另一部分的 GPU 计算交错进行。因此，一个 2048-token 批次不一定只调用一次 `ffn2048`；实际入口按每次提交的 token 数选择。上图为了看清分工，省略了这层跨层流水线。

源码：[选择 ANE 份额与初始化](https://github.com/nace-ai/kapitsa/blob/919b1b2a6782ec205e23725dee76ee5f57de2857/src/model/qwen35%5Fane.cpp?ref=blog.hacking.pub#L99)、[共享缓冲区分配](https://github.com/nace-ai/kapitsa/blob/919b1b2a6782ec205e23725dee76ee5f57de2857/src/ane/ffn%5Fsplit.cpp?ref=blog.hacking.pub#L131)、[两部分交错调度](https://github.com/nace-ai/kapitsa/blob/919b1b2a6782ec205e23725dee76ee5f57de2857/src/model/qwen35%5Fane.cpp?ref=blog.hacking.pub#L340)。

如果 ANE 执行失败，引擎会恢复这个批次的状态，改用 GPU 重跑，而不是继续消费不完整的输出。相关代码见[批次恢复与重跑](https://github.com/nace-ai/kapitsa/blob/919b1b2a6782ec205e23725dee76ee5f57de2857/src/engine/engine.cpp?ref=blog.hacking.pub#L3990)。

这里介绍的是实现策略，没有把它当作本机性能测试结果。512、640、2048、2560 都应结合这条代码路径理解，不能当作苹果对所有 ANE 程序规定的统一尺寸上限。