Study Apple Core AI

参考文档: Apple Developer Documentation - Core AI

概述

Apple Core AI 是 Apple 推出的端侧 AI 推理框架，核心设计理念：

使用 .aimodel 文件格式打包模型资产，针对当前设备硬件进行优化
模型加载与推理完全异步，支持并发推理
与 Xcode 深度集成，可在 IDE 中直接查看模型结构、精度、算子分布、输入输出 shape 等信息
提供 Core AI Debugger 和 Debug Gauge 两套调试工具链

.aimodel 文件结构

.aimodel 是一个目录（类似 macOS bundle），不是单个文件：

文件	作用
`metadata.json`	程序元数据（函数列表、输入输出签名、张量 spec）
`main.mlirb`	MLIR 字节码（计算图 + 权重），大模型中占主要体积
`main.hash`	完整性校验哈希

极简例子（x + x）总大小约 446 字节，真实模型取决于权重体积。

metadata.json 示例：

{
  "functions": ["main"],
  "main": {
    "inputs": ["x"],
    "outputs": ["y"],
    "input_specs": {"x": "tensor<2x3xf32>"},
    "output_specs": {"y": "tensor<2x3xf32>"}
  }
}

权重存储：

常规权重直接编码在 .mlirb 里
特殊格式（float8/sub-byte/量化）作为外部资源附加到 .aimodel 目录

Python vs Swift API

两套 API 做同样的事，但风格不同：

操作	Swift (iOS/macOS App)	Python (服务端/工具链)
构建 .aimodel	N/A（通过 coreai-torch 等工具链）	`program.save_asset(path)`
加载模型（只读元数据）	`AIModelAsset(url)`	`AIModelAsset.load(path)`
特化编译	`AIModel(contentsOf:)`	`async with asset.executable() as model:`
获取函数	`model.loadFunction("main")`	`model.load_function("main")`
执行推理	`function.run(inputs:)`	`await function(inputs)`

注意： executable() 只在 Python API 中出现，Swift 里特化编译直接发生在 AIModel(contentsOf:)。

编译时机

.aimodel 文件本身只是序列化的 MLIR + 权重，不是预编译的设备代码。

save_asset()           → 序列化到磁盘（快，构建时）
     │
     ▼
AIModelAsset.load()    → 读 metadata.json（快，只解析签名）
     │
     ▼
executable() / AIModel → 编译特化（慢，运行时首次加载）
     │
     ▼
function.run()         → 纯执行（快，每次推理）

为什么 Runtime 编译：

硬件差异：同一个 .aimodel 要在 M1/M2/M3、不同 Neural Engine 版本上运行
部署灵活性：一次构建，多处部署，硬件自动适配
版本解耦：编译器升级后不需要重新导出模型

代价： 冷启动延迟。首次 executable() / AIModel(contentsOf:) 需要编译特化，是昂贵操作。优化方案：服务启动时预热，或提前指定目标硬件特化。

核心类型

AIModel

AIModel 代表一个经过设备特化的 .aimodel 资产。

创建方式:

let model = try await AIModel(contentsOf: modelURL)

异步加载，从磁盘读取 .aimodel 文件
加载时自动针对当前设备硬件进行特化（specialization），这是一个昂贵操作

主要职责:

创建特化模型实例
加载 InferenceFunction
模型特化（Specializing）
提供属性查询

AIModelAsset

AIModelAsset 用于在不特化的情况下检查模型结构和元数据。

通过提供磁盘上 .aimodel bundle 的 URL 创建
可以查询模型信息（结构、元数据、输入输出 shape 等）而无需执行特化
适用场景：在决定是否加载模型前，先检查其兼容性或大小

InferenceFunction

InferenceFunction 拥有推理所需的资源，包括模型权重和中间缓冲区。

从 AIModel 加载得到
一个模型可以有多个 InferenceFunction，相当于同一个模型可以提供多种不同功能的前向计算
支持在运行前检查 InferenceFunction 的输入输出格式
并发安全: 该类型是 Sendable，可以从多个任务并发调用
并发运行时框架会自动分配额外的中间缓冲区

推理调用:

let result = try await function.run(inputs: inputArray, states: states, outputViews: outputViews)

模型加载流程

异步加载模型 — 通过 AIModel(contentsOf:) 从 .aimodel 文件创建
加载 InferenceFunction — 从 AIModel 中加载具体的推理函数，一个模型可对应多个函数（不同功能/模式）
格式检查 — 可以在运行前检查 InferenceFunction 的格式是否匹配
Metal Toolchain — 需要添加 Metal 工具链支持，使模型能够利用 GPU/Neural Engine 加速

推理执行

输入/输出:

输入为 NDArray 类型（n 维数组）
需要区分是否使用 inplace 操作（mutable vs immutable），inplace 可节省内存但需注意数据竞争
从 output 中取出结果，通常也是 NDArray 格式

并发推理:

InferenceFunction 是 Sendable，支持多任务并发调用
框架自动管理中间缓冲区的分配，并发时每个任务获得独立的中间状态

Xcode 集成

模型资产查看

将 .aimodel 文件添加到 Xcode 项目后，IDE 可直接展示：

模型大小
精度信息（FP16/INT8 等）
算子（op）分布
输入输出的 shape

Core AI Debugger

调试工具，支持：

可视化模型计算图结构
查看中间层输出值（intermediate values），便于定位推理异常

Debug Gauge

Xcode 中集成的性能监控面板：

显示 CPU/GPU 时间占用
可直观看到推理过程中硬件资源的使用情况