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

Resumen Técnico: SNN-MLIR

Planteamiento del Problema
El ecosistema de las Redes Neuronales de Picos (SNN, por sus siglas en inglés) se encuentra actualmente fragmentado en múltiples marcos de entrenamiento (por ejemplo, snnTorch, Lava, Norse), cada uno de los cuales utiliza formatos de modelo distintos. Si bien la Representación Intermedia Neuromórfica (NIR) se introdujo para resolver el problema del intercambio al proporcionar un formato de grafo independiente del marco de trabajo, esta se queda corta al intentar resolver el problema del despliegue. La NIR define qué es una red, pero no proporciona un camino para ejecutarla. En consecuencia, cada backend debe implementar de forma independiente estrategias de despliegue, que a menudo implican bucles de inferencia escritos a mano, cuantización específica para el objetivo y entornos de ejecución (runtimes) ad hoc. Existe una falta de una representación intermedia (IR) de compilador compartida y transformable que permita inspeccionar, optimizar y reducir progresivamente las SNN a código de máquina de manera consistente con las prácticas estándar de compilación.

Metodología
El artículo presenta snn-mlir, un dialecto de MLIR fuera de la ruta principal (out-of-tree) diseñado para cerrar la brecha entre la NIR y el código C para hardware base (bare-metal). El flujo de compilación procede de la siguiente manera:

1. Frontend: Una herramienta de Python lee cualquier grafo NIR y emite el IR del dialecto SNN. En modos cuantizados, inserta automáticamente operaciones snn.rescale para alinear las diferentes escalas de punto fijo entre las capas lineales y la dinámica de las neuronas, evitando errores silenciosos comunes en la cuantización manual.
2. Diseño del Dialecto: El dialecto se basa en tres principios fundamentales:
   • Basado en Memref: Los operandos son referencias de memoria (memref) en lugar de valores SSA, lo que refleja la naturaleza con estado e in-place de la dinámica de las neuronas a través de los pasos temporales (timesteps).
   • Tipado Explícito: Las operaciones portan explícitamente los tipos de elementos y las formas (shapes) en el IR.
   • Polimorfismo de Tipos: Un único conjunto de operaciones modela tanto datos de punto flotante (f32/f64) como cuantizados (i8/i32). El dialecto no mantiene conjuntos de operaciones paralelos para flotantes e enteros; en su lugar, las operaciones portan ambas representaciones de parámetros (atributos de valor real y atributos de punto fijo Q12), seleccionando el conjunto apropiado según el tipo de operando durante la reducción (lowering).
3. Reducción (Lowering): Un paso de reducción de referencia (SNNToLinalg) convierte las operaciones del dialecto en operaciones estándar de linalg y arith. Esto aprovecha el toolchain existente de MLIR y LLVM para generar LLVM IR, que luego se compila a código de máquina.
4. Generación de Código: El toolchain produce código C11 autónomo y sin dependencias, capaz de ejecutarse en cualquier CPU compatible con C o en un objetivo embebido.

Contribuciones Clave

1. Un Dialecto SNN Polimórfico de Tipos: Un conjunto compacto de operaciones que modela la familia de neuronas CUBA-LIF y las sinapsis lineales. El diseño asegura que la simulación (punto flotante) y el despliegue (cuantizado) compartan un único IR, eliminando el riesgo de divergencia semántica entre ambos caminos.
2. Alineación Automática de Escala de Cuantización: Un mecanismo que inserta operaciones de reescalado en los límites sinapsis-neurona para alinear las escalas de los pesos con las escalas del estado de la neurona (Q12). Esto elimina una fuente común de error en la cuantización manual.
3. Reducción a CPU de Referencia: Un paso de conversión que reescribe las operaciones del dialecto en dialectos estándar de MLIR, permitiendo la generación de binarios C11 portátiles y sin dependencias.
4. Evaluación de Fidelidad y Portabilidad: El sistema fue evaluado en redes de alimentación hacia adelante (feedforward) y totalmente conectadas a través de objetivos x86 y RISC-V, verificando la corrección numérica y demostrando que el mismo IR puede ser re-dirigido sin modificaciones.

Resultados
La evaluación se centró en dos redes (snn_oxford y snntorch) compiladas tanto en modo de punto flotante como en modo cuantizado (int8):

• Fidelidad Numérica: En modo de punto flotante, la salida compilada es bit a bit idéntica con el marco de referencia. En modo cuantizado, los resultados son deterministas a través de diferentes conjuntos de instrucciones (x86 y RISC-V), con una precisión que coincide con la referencia dentro de un 99.58% y 92.00% respectivamente.
• Rendimiento: Los binarios C compilados corren de 157× a 266× más rápido que los baselines de los marcos de trabajo originales de Python en x86. Este aumento de velocidad se atribuye principalmente a la eliminación de la sobrecarga del intérprete de Python y de los costos de despacho de tiempo de ejecución por cada paso temporal.
• Portabilidad y Huella: El mismo grafo NIR se recompila sin cambios a un núcleo RISC-V X-HEEP embebido. La huella de código y de tiempo de ejecución es fija en aproximadamente 12.6 KB, independientemente del tamaño del modelo. La huella de datos está dominada por los pesos, que se reducen aproximadamente 4× mediante la cuantización, permitiendo el despliegue en dispositivos con RAM limitada.

Significado y Alcance
El artículo posiciona a snn-mlir no como una alternativa de alto rendimiento para la inferencia en GPU, sino como una vía de despliegue portátil y sin dependencias para las SNN. Establece la primera representación intermedia de compilador adecuada para redes de impulsos, tratando a la red como un objeto de compilación de primera clase en lugar de un grafo opaco.

Los autores reconocen explícitamente las limitaciones actuales: el sistema soporta actualmente solo topologías de alimentación hacia adelante, totalmente conectadas (grafos de cadena lineal) con un tamaño de lote (batch size) de uno y parámetros de neurona uniformes por capa. Las operaciones convolucionales, la ramificación y los grafos recurrentes aún no son compatibles. El trabajo se enmarca como un punto de partida fundacional, diseñado para crecer con la comunidad neuromórfica, con trabajos futuros dirigidos a operaciones convolucionales, soporte de grafos recurrentes y metodologías de cuantización estandarizadas en todo el ecosistema NIR. El proyecto se publica como código abierto bajo la licencia Apache-2.0.