SNN-MLIR: An MLIR Dialect for Compiling Neuromorphic SNNs from NIR to Bare-Metal C

Resumo Técnico: SNN-MLIR

Problema
O ecossistema de Redes Neurais de Impulsos (SNN) é atualmente fragmentado entre múltiplos frameworks de treinamento (ex: snnTorch, Lava, Norse), cada um utilizando formatos de modelo distintos. Embora a Representação Intermediária Neuromórfica (NIR) tenha sido introduzida para resolver o problema de intercâmbio, fornecendo um formato de grafo independente de framework, ela não chega a resolver o problema de implantação. A NIR define o que uma rede é, mas não fornece um caminho para executá-la. Consequentemente, cada backend deve implementar de forma independente estratégias de implantação, muitas vezes envolvendo loops de inferência escritos à mão, quantização específica para o alvo e runtimes ad-hoc. Há uma carência de uma representação intermediária (IR) de compilador compartilhada e transformável que permita que as SNNs sejam inspecionadas, otimizadas e progressivamente reduzidas ao código de máquina de maneira consistente com as práticas padrão de compiladores.

Metodologia
O artigo apresenta o snn-mlir, um dialeto MLIR out-of-tree projetado para preencher a lacuna entre a NIR e o código C de bare-metal. O fluxo de compilação procede da seguinte forma:

1. Frontend: Uma ferramenta Python lê qualquer grafo NIR e emite o IR do dialeto SNN. Em modos quantizados, ela insere automaticamente operações snn.rescale para alinhar diferentes escalas de ponto fixo entre camadas lineares e a dinâmica dos neurônios, prevenindo erros silenciosos comuns na quantização feita à mão.
2. Design do Dialeto: O dialeto é construído sobre três princípios fundamentais:
   • Baseado em Memref: Os operandos são referências de memória (memref) em vez de valores SSA, refletindo a natureza estatal e in-place da dinâmica dos neurônios através dos passos temporais (timesteps).
   • Tipagem Explícita: As operações carregam tipos de elementos e formas (shapes) explicitamente no IR.
   • Polimorfismo de Tipo: Um único conjunto de operações modela tanto dados de ponto flutuante (f32/f64) quanto quantizados (i8/i32). O dialeto não mantém conjuntos de operações paralelos para float e inteiro; em vez disso, as operações carregam ambas as representações de parâmetros (atributos de valor real e atributos Q12 de ponto fixo), selecionando o conjunto apropriado com base no tipo do operando durante a redução (lowering).
3. Redução (Lowering): Uma passagem de redução de referência (SNNToLinalg) converte as operações do dialeto em operações padrão dos dialetos linalg e arith. Isso aproveita o toolchain existente do MLIR e LLVM para gerar o IR do LLVM, que é então compilado para código de máquina.
4. Geração de Código: O toolchain produz código C11 autossuficiente e livre de dependências, capaz de rodar em qualquer CPU ou alvo embarcado compatível com C.

Principais Contribuições

1. Um Dialeto SNN Polimórfico de Tipo: Um conjunto compacto de operações que modelam a família de neurônios CUBA-LIF e sinapses lineares. O design garante que a simulação (ponto flutuante) e a implantação (quantizada) compartilhem um único IR, eliminando o risco de divergência semântica entre os dois caminhos.
2. Alinhamento Automático de Escala de Quantização: Um mecanismo que insere operações de reescalonamento nas fronteiras sinapse-neurônio para alinhar as escalas dos pesos com as escalas do estado do neurônio (Q12). Isso remove uma fonte comum de erro na quantização manual.
3. Redução para CPU de Referência: Uma passagem de conversão que reescreve as operações do dialeto em dialetos MLIR padrão, permitindo a geração de binários C11 portáteis e sem dependências.
4. Avaliação de Fidelidade e Portabilidade: O sistema foi avaliado em redes feedforward totalmente conectadas através de alvos x86 e RISC-V, verificando a correção numérica e demonstrando que o mesmo IR pode ser retargetado sem modificações.

Resultos
A avaliação focou em duas redes (snn_oxford e snntorch) compiladas tanto em modos de ponto flutuante quanto quantizados (int8):

• Fidelidade Numérica: No modo de ponto flutuante, a saída compilada é bit-a-bit idêntica ao framework de referência. No modo quantizado, os resultados são determinísticos entre diferentes conjuntos de instruções (x86 e RISC-V), com precisão correspondendo à referência em 99,58% e 92,00%, respectivamente.
• Desempenho: Os binários C compilados rodam de 157× a 266× mais rápido que as linhas de base dos frameworks Python originais em x86. Este ganho de velocidade é atribuído principalmente à remoção do overhead do interpretador Python e dos custos de despacho de runtime por passo temporal.
• Portabilidade e Pegada (Footprint): O mesmo grafo NIR é transcompilado sem alterações para um núcleo RISC-V X-HEEP embarcado. A pegada de código e runtime é fixa em aproximadamente 12,6 KB, independente do tamanho do modelo. A pegada de dados é dominada pelos pesos, que são reduzidos em cerca de 4× através da quantização, permitindo a implantação em dispositivos com RAM limitada.

Significância e Escopo
O artigo posiciona o snn-mlir não como uma alternativa de alto desempenho para inferência em GPU, mas como um caminho de implantação portátil e livre de dependências para SNNs. Ele estabelece a primeira IR de compilador adequada para redes de impulsos, tratando a rede como um objeto de primeira classe do compilador, em vez de um grafo opaco.

Os autores reconhecem explicitamente as limitações atuais: o sistema suporta atualmente apenas topologias feedforward e totalmente conectadas (grafos de cadeia linear) com um tamanho de lote (batch size) de um e parâmetros de neurônio uniformes por camada. Operações convolucionais, ramificações (branching) e grafos recorrentes ainda não são suportados. O trabalho é estruturado como um ponto de partida fundamental, projetado para crescer com a comunidade neuromórfica, com trabalhos futuros visando operações convolucionais, suporte a grafos recorrentes e metodologias de quantização padronizadas em todo o ecossistema NIR. O projeto é lançado como código aberto sob a licença Apache-2.0.