Sumário Técnico: ANEForge

Definição do Problema

O Apple Neural Engine (ANE) é um acelerador neural de função fixa presente em todos os dispositivos Apple recentes. No entanto, o único método sancionado para acessar este hardware é através do framework CoreML da Apple. O CoreML trata o ANE como uma opção de escalonamento, e não como um alvo garantido; ele compila modelos e seleciona unidades de computação (CPU, GPU ou ANE) em tempo de execução com base em heurísticas. Consequentemente, nenhuma configuração força explicitamente a execução no ANE, e a CPU permanece como um fallback permitido. Além disso, o tempo de execução não reporta qual unidade executou uma chamada específica, tornando impossível garantir que uma carga de trabalho eficiente seja executada no hardware de maior eficiência energética.

A verdadeira superfície de programação do ANE reside em uma pilha de frameworks privados e não documentados da Apple (incluindo Espresso, aned e o driver do kernel) que o CoreML utiliza internamente. Esta pilha não é exposta ao código da aplicação, impedindo a caracterização direta das capacidades do motor e o desenvolvimento de fluxos de trabalho personalizados que exijam posicionamento determinístico no ANE.

Metodologia

ANEForge é um pacote Python projetado para contornar o CoreML e programar o ANE diretamente, fazendo interface com a pilha de frameworks privados. A metodologia envolve os seguintes componentes arquiteturais:

• Compilação de Grafo de Tensores Preguiçoso (Lazy Tensor Graph Compilation): Os usuários constroem um grafo de operadores preguiçoso usando uma superfície definida de 58 operadores fundidos e 19 operadores de ponte nativa. O compilador reduz todo este grafo em um único programa ANE usando a Linguagem Intermediária de Modelo (MIL).
• Despacho Direto via Símbolos Privados: O pacote utiliza um shim de Objective-C++ compacto (compilado na máquina do usuário) para acessar símbolos públicos de frameworks privados da Apple, especificamente a interface C e5rt exportada pelo runtime Espresso. Isso permite o despacho para o daemon do sistema aned e para o driver de kernel do ANE sem exigir permissões (entitlements) da Apple ou modificar a proteção de integridade do sistema (SIP).
• Superfície de Operadores:
  • Rota Fundida (Fused Route): Reduz 58 operadores (incluindo convoluções, multiplicações de matrizes, ativações, normalização e atenção) em um único programa.
  • Rota de Ponte Nativa (Native Bridge Route): Executa 19 operadores nativos (ex: atenção fundida, argmax, topk, camadas de geometria) como subprogramas via cortes de grafo, acessando camadas de hardware não expostas pelo toolchain público.
• Treinamento e Gerenciamento de Estado: O pacote implementa um módulo de autograd de modo reverso para executar passagens de ida (forward), passagens de volta (backward) e atualizações de otimizador inteiramente no motor. Ele utiliza o aliasing de buffers de saída-para-entrada para manter os caches de chave-valor (para decodificação autorregressiva) e estados de otimizador residentes no motor entre os passos, evitando a transferência de dados entre o host e o dispositivo.
• Compressão de Pesos: O sistema suporta o streaming de pesos int8, int4 (lookup-table) e esparsos diretamente do caminho de desquantização do motor, em vez de dobrá-los para meia precisão no momento da compilação.

Principais Contribuições

1. Acesso Livre ao CoreML: O ANEForge fornece a primeira interface Python para o ANE que não depende do CoreML para escalonamento, garantindo a execução determinística no acelerador de hardware.
2. Atenção Fundida Nativa: Expõe a camada dedicada de atenção fundida do motor (sdpa), que o compilador público do CoreML decompõe em operações padrão de multiplicação de matrizes e softmax.
3. Treinamento no Dispositivo: O pacote possibilita loops de treinamento completos (forward, backward, otimizador) no ANE, mantendo gradientes e estados de otimizador residentes no dispositivo.
4. Streaming de Pesos: Introduz um mecanismo para transmitir pesos comprimidos (int4, int8, esparsos) diretamente para o motor, reduzindo a pegada de memória em até 4x para int4 sem sacrificar o caminho de execução.
5. Métodos Numéricos: Estende a utilidade do ANE além de redes neurais para incluir álgebra linear numérica de iteração fixa (ex: solvers de Krylov) e métodos espectrais (FFT) como grafos de fluxo de dados estáticos.
6. Registro de Capacidade Verificado por Máquina: Uma suíte de testes de conformidade verifica se cada operador no registro compila e exece corretamente no hardware, servindo como um portão de regressão para novos lançamentos.

Resultados

O artigo valida o pacote contra frameworks de referência (NumPy, PyTorch, torchvision, transformers, diffusers) em Apple Silicon (M1 Max e M5 Pro):

• Performance: Um pequeno programa fundido completa uma chamada em aproximadamente 90 µs, aproximando-se do piso de despacho do motor de ~70 µs. Uma passagem de ida pré-treinada de ResNet-18 roda de ponta a ponta em 0,33 ms no ANE, comparado a 2,0 ms no GPU (MPS) e 6,0 ms na CPU.
• Acurácia: Modelos pré-treinados, incluindo ResNet-18, ViT-B/16 e MiniLM, alcançam similaridade de cosseno de 1,0000 contra referências de framework float32.
• Treinamento: Uma rede convolucional treinada no CIFAR-10 usando ANEForge alcançou 71,2% de acurácia de teste, dentro de 1,7 pontos percentuais de uma linha de base PyTorch (72,9%), com gradientes correspondendo a 1,0 de similaridade de cosseno.
• Otimização: Uma reescrita de otimizador que funde blocos de atenção em um único programa gera acelerações de 3,7–5,3x através de várias extensões de sequência.
• Energia: Durante o loop quente de ResNet-18, o ANE consome 4,5 W, demonstrando uso eficiente de energia comparado a alternativas de CPU/GPU.

Significância e Escopo

O artigo posiciona o ANEForge como uma ferramenta para pesquisa e desenvolvimento local, observando explicitamente que não pode ser distribuído via App Store devido à sua dependência de frameworks privados. Sua principal significância reside em permitir:

• Caracterização: Um censo verificado por máquina de quais operadores o ANE aceita e uma medição rigorosa de sua velocidade e eficiência energética sob condições controladas.
• Personalização Preservadora de Privacidade: Um substrato para treinamento e personalização no dispositivo onde dados, gradientes e estados de otimizador nunca deixam o dispositivo, utilizando o bloco de computação de menor consumo do SoC.
• Transparência de Hardware: Fornecer um caminho de execução determinístico que elimina a incerteza das heurísticas de tempo de execução encontradas no CoreML.

Os autores mantêm um escopo modesto, reconhecendo que o caminho de computação é limitado à meia precisão (fp16), o que restringe a computação científica geral e impede que loops de amostragem de difusão de guia livre de classificador (classifier-free guidance) se fechem totalmente devido a erros de acumulação. O pacote é projetado para ser extensível, permitindo que usuários adicionem operadores contra o portão de conformidade ou construam ferramentas de medição, mas permanece vinculado às capacidades específicas de hardware e versões de API privadas das versões verificadas do macOS.