Resumen Técnico: DCC – Compilación Centrada en Datos de Kernels de Aprendizaje Automático para Arquitecturas de Procesamiento en Memoria

1. Declaración del Problema

Las arquitecturas de Procesamiento en Memoria (PIM) ofrecen una solución prometedora al cuello de botella de la memoria en el Aprendizaje Automático (ML) al colocar unidades de procesamiento cerca de los bancos de DRAM. Sin embargo, existe una incompatibilidad fundamental entre los requisitos de disposición de datos de los procesadores Host de alto rendimiento (por ejemplo, GPUs) y los núcleos PIM.

• Requisito del Host: Para explotar el paralelismo a nivel de banco y el alto ancho de banda, los procesadores Host requieren que los elementos de datos consecutivos se distribuyan entre múltiples bancos de DRAM.
• Requisito de PIM: Los núcleos PIM solo pueden acceder a datos de sus bancos locales. Para maximizar el ancho de banda local y habilitar un acceso eficiente a múltiples elementos, los elementos consecutivos deben colocarse dentro del mismo banco.

Esto hace necesarias reorganizaciones de datos (de entrada y salida) realizadas a través del bus de memoria del Host, lo que incurre en costos significativos de movimiento de datos que pueden dominar el tiempo de ejecución de extremo a extremo. Los enfoques de compilación actuales para sistemas PIM adolecen de dos limitaciones principales:

1. Falta de Optimización Sistemática: Las herramientas existentes suelen dirigirse a backends específicos (por ejemplo, UPMEM) o kernels específicos, careciendo de soporte para cargas de trabajo de ML diversas y múltiples dispositivos PIM (por ejemplo, Samsung HBM-PIM, SK Hynix GDDR6-AiM).
2. Ajuste Centrado en el Cálculo: Los compiladores existentes (por ejemplo, ATiM) optimizan los programas de cálculo independientemente de los costos de reorganización de datos. Los autores demuestran que las transformaciones de cálculo y las reorganizaciones de datos son interdependientes; optimizarlas de forma aislada produce un rendimiento subóptimo porque un programa de cálculo eficiente puede requerir reorganizaciones de datos costosas, mientras que un programa de cálculo menos eficiente podría permitir un movimiento de datos más barato.

2. Metodología: DCC

Los autores proponen DCC, el primer compilador de ML centrado en datos para sistemas PIM. DCC co-optimiza conjuntamente las estrategias de reorganización de datos y los programas de cálculo a nivel de bucle dentro de un proceso de ajuste unificado. El sistema consta de cuatro componentes clave:

2.1 Abstracción Multicapa

DCC introduce una abstracción genérica que mapea la jerarquía de memoria PIM en una jerarquía de cálculo, desacoplando al compilador de las semánticas específicas del backend. Esta abstracción opera en tres niveles:

• Nivel de Sistema: El xPU Host (por ejemplo, GPU), el espacio de memoria Host y los dispositivos PIM.
• Nivel de Grupo PIM: Grupos de núcleos PIM (y sus bancos locales) que pertenecen al mismo canal o rango de memoria. Esto permite optimizaciones de grano grueso y difusión de instrucciones a nivel de grupo.
• Nivel de Núcleo PIM: Unidades de procesamiento individuales y sus bancos locales, habilitando optimizaciones de grano fino para el acceso a memoria y la disposición de datos local.

2.2 Generador de Programas Centrados en Datos

A diferencia de los compiladores tradicionales que comienzan con transformaciones de bucles, DCC adopta una estrategia centrada en los datos:

1. Construcción del Conjunto de Dimensiones: Las dimensiones del tensor se correlacionan con variables de bucle para formar "conjuntos de dimensiones complejas".
2. Construcción de Fragmentos de Datos: El sistema genera exhaustivamente particiones de datos candidatas (fragmentos) a través de grupos y núcleos PIM.
3. Poda de Borradores: Tres reglas eliminan borradores poco prometedores:
   • Redundancia: Eliminar borradores con particiones de datos y mapeos idénticos.
   • Alineación SIMD: Podar borradores donde los tamaños de partición no se alinean con el ancho SIMD del backend (por ejemplo, 16 vías).
   • Equivalencia de Rendimiento: Mantener solo un borrador si múltiples borradores producen el mismo rendimiento por núcleo en el peor caso.
4. Mapeo: Los borradores restantes se mapean a fragmentos de cálculo, determinando los rangos de bucle y los índices de tensor para cada núcleo.

2.3 Optimizador de Código Específico de PIM

Este componente aplica optimizaciones adaptadas a las restricciones de PIM:

• Optimización de Reorganización de Datos: Explota la memoria controlable on-chip (por ejemplo, memoria compartida de GPU) para reorganizar los datos. Realiza lecturas en serie desde los canales de memoria del Host seguidas de escrituras paralelas a los canales PIM, maximizando el paralelismo de canal tanto para operaciones de lectura como de escritura.
• Optimización del Código de Cálculo: Utiliza la generación jerárquica de comandos. Prioriza comandos estilo SIMD a nivel de grupo (difundiendo una única instrucción a todos los núcleos de un grupo) para reducir la sobrecarga del bus de comandos, recurriendo a comandos a nivel de banco solo cuando sea necesario.

2.4 Predictor de Rendimiento Acoplado

Para evitar el costo prohibitivo de perfilar miles de borradores candidatos, DCC emplea un predictor basado en aprendizaje (XGBoost).

• Función: Estima conjuntamente los costos de reorganización de datos y los tiempos de ejecución de cálculo para seleccionar la configuración de mejor rendimiento de extremo a extremo.
• Entrenamiento: El modelo se entrena fuera de línea en un subconjunto diverso de borradores.
• Tiempo de Ejecución: Para tamaños de tensor estáticos, el mejor programa se pre-selecciona. Para tamaños dinámicos (por ejemplo, en LLM), el predictor estima el rendimiento sobre la marcha.

2.5 Interfaz de Programación

DCC proporciona una interfaz de Python compatible con PyTorch. Los usuarios anotan kernels con @DCC_kernel. El compilador analiza estas definiciones de alto nivel, genera la Representación Intermedia (IR) que involucra fragmentos de datos y cálculo, y emite instrucciones de backend de bajo nivel.

3. Resultados Clave

Los autores evaluaron DCC en dos backends PIM de última generación: Samsung HBM-PIM y AttAcc (un prototipo de investigación con unidades especializadas GEMV/Softmax), comparándolo con la ejecución solo en GPU y las implementaciones base de PIM.

• Kernels de ML Individuales:
  • HBM-PIM: DCC logra una aceleración de hasta 7.68× (2.21× promedio) sobre la ejecución en GPU.
  • AttAcc: DCC logra una aceleración de hasta 13.17× (3.92× promedio) sobre la ejecución en GPU.
  • Comparación con PIM Base: DCC mejora la implementación base de AttAcc en hasta 2.74× (1.75× promedio) y la de HBM-PIM en hasta 2.09× (1.53× promedio) para kernels específicos como GEMV y RED.
• Inferencia de LLM de Extremo a Extremo:
  • En AttAcc, DCC acelera la inferencia de GPT-3 y LLaMA-2 en un promedio de 4.52× (hasta 7.71× para LLaMA-2) sobre la ejecución solo en GPU.
  • DCC supera significativamente a la implementación base de AttAcc (que utiliza fragmentación fija) al habilitar la ejecución eficiente de las capas de generación y proyección QKV, que anteriormente estaban limitadas por la GPU.
• Eficiencia Energética: DCC reduce el consumo de energía en un promedio de 1.24× en comparación con el sistema base AttAcc debido a la reducción del movimiento de datos.
• Sobrecarga de Compilación: El tiempo de compilación fuera de línea es bajo (segundos), y el esquema de poda reduce el espacio de búsqueda en ~9×. El predictor acoplado logra una alta precisión (97.37% en LLM) y reduce el tiempo de compilación en ~3.8× en comparación con la búsqueda exhaustiva.

4. Significado y Afirmaciones

El artículo afirma que DCC es el primer compilador de ML centrado en datos para sistemas PIM que:

1. Optimiza Conjuntamente la reorganización de datos y los programas de cálculo, demostrando que estos dos factores son interdependientes y deben ajustarse juntos para evitar un rendimiento subóptimo.
2. Soporta Múltiples Backends a través de una abstracción multicapa, habilitando la portabilidad entre diversas arquitecturas PIM (por ejemplo, HBM-PIM, AttAcc) sin requerir un ajuste manual específico del backend.
3. Habilita Alto Rendimiento para kernels de ML intensivos en memoria e inferencia de LLM a gran escala, superando significativamente tanto a la ejecución solo en GPU como a los enfoques de compilación PIM existentes.
4. Mejora la Programabilidad al proporcionar una interfaz de Python de alto nivel que abstrae la complejidad de la gestión de disposición de datos de bajo nivel y la programación de instrucciones para dispositivos PIM.

Los autores concluyen que DCC proporciona un camino práctico y efectivo para la aceleración heterogénea de ML, equilibrando alto rendimiento con programabilidad en los ecosistemas PIM emergentes. El código es de código abierto para facilitar futuras investigaciones.