让Python直接调用苹果神经引擎,实现端侧高效推理与训练。
ANEForge: Python for direct computation on the Apple Neural Engine
- 通过编译张量图直接操控苹果神经引擎,跳过CoreML限制。
- ResNet-18端到端推理仅需0.33ms,单次调用约90us,接近硬件调度下限。
- 支持量化、稀疏权重及训练全过程,适合需要极致性能的开发者。
ANEForge是一个Python包,可直接编程苹果神经引擎(ANE),即所有近期苹果设备中的固定功能神经加速器,无需通过CoreML。在生产环境中,该引擎仅可通过CoreML访问,而后者将其视为调度选项:不强制使用ANE,模型可能无声地运行在CPU或GPU上。ANEForge将由58个融合操作符和19个原生桥接操作符构建的延迟张量图编译为单一ANE程序,并通过与苹果内部框架相同的ANE守护进程和内核驱动栈分发。除推理外,该工具还能调用引擎原生融合注意力机制,支持int8、int4及稀疏权重流,保持解码器与优化器状态跨步骤驻留,并在引擎上完成前向、反向传播及优化器更新。小型融合程序单次调用耗时约90μs,接近引擎70μs的每程序调度下限;预训练ResNet-18前向推理端到端耗时0.33ms。ResNet-18、句子编码器及视觉变换器在与框架参考结果对比中表现一致,Stable Diffusion U-Net的前向传播亦得到验证。ANEForge适用于macOS 14及更高版本的Apple Silicon。每次发布均经由记录的macOS与ANE编译器版本验证。
原文摘要 · Abstract (English)
ANEForge is a Python package that programs the Apple Neural Engine (ANE), the fixed-function neural accelerator on every recent Apple device, directly and without CoreML. In production the engine is reachable only through CoreML, which treats it as a scheduling option: no configuration requires the ANE, and a model can silently run on the CPU or GPU instead. ANEForge compiles a lazy tensor graph, built from 58 fused operators and 19 native bridge operators, into a single ANE program. The program is dispatched through the same ANE daemon and kernel-driver stack as Apple's internal framework. Beyond inference, the package reaches the engine's native fused attention, streams int8, int4, and sparse weights, keeps decoder and optimizer state resident across steps, and runs the forward pass, backward pass, and optimizer update of training on the engine. A small fused program completes a call in about 90us, near the engine's 70us per-program dispatch floor, and a pretrained ResNet-18 forward runs end-to-end in 0.33ms. ResNet-18, a sentence encoder, and a Vision Transformer run end-to-end against framework references, and a Stable Diffusion U-Net validates its forward pass. ANEForge targets Apple Silicon under macOS 14 and later. Each release is verified against a recorded macOS and ANE-compiler version.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。