Overview

This groundbreaking research project unlocks direct training capabilities on Apple's Neural Engine (ANE), bypassing official inference-only restrictions. By reverse-engineering private _ANEClient and _ANECompiler APIs, this project demonstrates a proof-of-concept for running backpropagation and custom compute graphs directly on ANE hardware. It offers a deep dive into ANE performance characteristics, including throughput, power consumption, and SRAM behavior, providing valuable insights for researchers and developers exploring edge AI optimization.

Key Features

• Direct ANE Training: Enables training neural networks, including custom compute graphs and backpropagation, directly on the Apple Neural Engine.
• Reverse-Engineered Private APIs: Leverages discovered private _ANEClient and _ANECompiler APIs, as well as the MIL format, for full ANE control.
• Performance Benchmarks: Provides detailed benchmarks on ANE throughput, power, and SRAM utilization, including comparisons between FP16 and INT8 W8A8 quantization.
• Dynamic Kernel Pipeline: Implements a dynamic pipeline for transformer models, allowing weights to be updated without recompilation, significantly improving efficiency.
• Optimized GPU-ANE Interoperability: Features a zero-copy pipeline between GPU and ANE via shared IOSurfaces for efficient inference workflows.
• INT8 Quantization Acceleration: Achieves up to 1.88x throughput improvement by leveraging INT8 activations and weights.
• Comprehensive Architecture Details: Outlines the kernel functions, CPU offloads, and key optimizations for various model architectures like MHA and GQA.

Typical Use Cases

• Research and Development: Ideal for researchers investigating the direct capabilities of NPUs beyond their advertised inference functions.
• Performance Analysis: Provides a platform for in-depth analysis of ANE hardware characteristics like throughput, power efficiency, and memory bandwidth.
• Exploring Edge AI Potential: Serves as a foundational reference for developers and enthusiasts looking to push the boundaries of on-device AI training and execution.
• Custom Compute Graph Exploration: Enables experimentation with custom compute graphs and novel neural network architectures on ANE hardware.