Deep dive revela o compilador oculto do Apple Neural Engine

O que é o texto

O autor, em continuidade a uma série de estudos sobre compiladores (LLVM, KGEN, MLIR), analisa o capítulo sobre o compilador do Apple Neural Engine (ANE) no artigo de Spencer Bryngelson, "Apple Neural Engine: Architecture, Programming, and Performance". O ANE é um alvo pouco documentado: a Apple o expõe publicamente via Core ML, mas o compilador, formato de programa, driver e firmware permanecem não documentados. O texto de Bryngelson é um trabalho independente de engenharia reversa, baseado em medições diretas e análise do software privado da Apple. O foco do post é o que acontece depois do grafo da rede neural até virar algo executável pelo acelerador.

Fusão: quando as fronteiras entre operações deixam de ser úteis

Um grafo com convolução, soma de bias e ReLU entra como três operações, mas o compilador do ANE pode absorver bias e ativação na própria convolução: o bias vai para a maquinaria de ganho-offset do hardware e a ativação ocupa um dos slots de ativação da operação. Padding e parte do caminho de quantização também podem ser incorporados. O resultado sai como uma única operação ANE fundida, não três despachos independentes.

O autor enfatiza que a fronteira entre as operações não era um erro: era útil para raciocinar sobre a estrutura da rede e desnecessária quando o hardware já expressa o comportamento combinado. A mesma lógica vale para transposes, batch normalization e dequantização, que podem desaparecer em operações vizinhas ou em caminhos dedicados.

Há limites claros: a soma com um bias constante cabe na maquinaria de ganho-offset; a soma entre duas saídas de convolução calculadas independentemente não cabe, pois ambos os operandos são valores vivos. Concat cria outra fronteira de fusão, atenção introduz segmentos mais difíceis, e restrições de capacidade podem até forçar o compilador a dividir algo antes fundido. A fusão não acontece porque "menos operações é melhor", mas quando o hardware tem uma representação capaz de expressar o comportamento combinado.

Validade não é legalidade

Um grafo matematicamente válido não garante que a máquina o execute naquela forma. O ANE tem um envelope: rank de tensor limitado, dimensões com limites, restrições específicas de kernels de convolução e coeficientes sujeitos a limites de memória. Operações fora dessas fronteiras precisam ser divididas (tiling), transformadas ou separadas em segmentos.

O compilador expõe validadores de formas, tipos, dimensões e capacidades do alvo, mas passar neles não garante que a geração de código no backend funcione. No M1, o artigo encontrou operações como top-k, sort, dynamic slice e convolução 3D que passam por partes da validação e ainda falham no lowering. Para o autor, "suportado" é uma palavra vaga: representável, válido, baixável e executável são fronteiras diferentes no stack.

Multiplicação de matrizes que vira convolução

Se o peso à direita de uma multiplicação de matrizes cabe no orçamento de memória on-chip do ANE, o compilador reescreve a operação como uma convolução residente; se não cabe, permanece uma multiplicação de matrizes com tiling. No M1, o limiar é de cerca de 2 MB. O autor conclui que o compilador preserva a computação, não o nome dado à operação: lowering é escolher entre realizações possíveis da mesma intenção.

Agendamento, memória e DMA

A fase de agendamento lineariza o grafo e o divide em partições de descritores de tarefa que respeitam o orçamento on-chip - no M1, um teto de cerca de 2 MB de memória estática. O particionamento acontece quando a próxima peça de trabalho empurraria a pressão de memória viva além do orçamento. Para o autor, isso mostra que, próximo ao hardware, o compilador vira também um sistema de gestão de recursos.

Na etapa de memória e DMA, a "seta" entre duas operações, vista como dependência gratuita no nível de grafo, torna-se tráfego de memória. O artigo descreve estados de alocação que distinguem dados residentes on-chip, transmitidos da memória dinâmica, em ring buffers, reescritos in-place ou encadeados entre operações: um produtor com um único consumidor pode manter o resultado on-chip. O compilador também reorganiza strides de buffers de destino para evitar conflitos de banco, incorpora padding e empacota streams de pesos. Quando a computação fica barata, mover dados passa a ser grande parte do problema.

Conclusão

Com um comentário de Chris Lattner sobre o custo de cada IR adicional, o autor sintetiza sua tese: representações existem enquanto respondem a perguntas; carregá-las além disso é overhead. Comparando LLVM, KGEN/MLIR e o ANE, o compilador aparece menos como tradutor mecânico e mais como o lugar onde várias descrições da mesma computação negociam com as restrições físicas da máquina - até que a forma de descrever a computação seja compatível com a forma de executá-la, preservando o resultado dado o mesmo input válido.