Resumen Técnico: Hacia Modelos de Mundo para Compiladores

Declaración del Problema

La optimización de programas de tensores es esencial para los sistemas modernos de aprendizaje automático, sin embargo, la identificación del programa más eficiente se ve obstaculizada por un enorme espacio de búsqueda combinatoria. Para evitar el costo prohibitivo de las mediciones en el dispositivo, los auto-programadores existentes (por ejemplo, Ansor, AutoTVM) dependen de modelos de costo aprendidos para clasificar candidatos de programas. Sin embargo, estos modelos suelen tratar a los candidatos como instantáneas de código estático (por ejemplo, el texto final de TensorIR).

Esta perspectiva estática introduce dos limitaciones críticas:

1. Pérdida de Contexto de Transformación: Ignora el proceso de transformación paso a paso que generó el programa. En consecuencia, los evaluadores no pueden comprender cómo una decisión de programación en un paso posterior depende del contexto establecido por decisiones anteriores.
2. Sensibilidad al Ruido Sintáctico: Los codificadores estáticos son fácilmente engañados por variaciones superficiales de código. Por ejemplo, dos programas estructuralmente equivalentes (por ejemplo, multiplicación de matrices con tamaños de bucle $N=1024$ frente a $N=1023$) pueden producir textos de TensorIR significativamente diferentes debido al manejo de límites (peeling de bucles), lo que provoca que el modelo clasifique erróneamente su potencial de rendimiento a pesar de tener un comportamiento de hardware idéntico.

Metodología

Los autores proponen un marco de Modelo de Mundo de Compilador que redefine la evaluación de la programación como dinámicas latentes condicionadas por acciones sobre estados de programas. En lugar de evaluar un artefacto de código estático final, el marco simula la evolución del estado del programa a través de un espacio latente continuo.

El marco consta de tres componentes aprendidos implementados dentro del AutoScheduler de TVM:

1. Representación del Estado del Programa de Tensores (Codificador)

• Entrada: Texto de TensorIR (estados pre-programación y post-programación).
• Arquitectura: Una arquitectura basada en CodeBERT entrenable y adaptada para entradas largas. Utiliza una ventana deslizante para tokenizar y codificar fragmentos superpuestos, agregándolos mediante un agrupamiento de media enmascarada (masked mean pooling) en un vector de representación denso $z = E(s)$.
• Entrenamiento: Utiliza aprendizaje contrastivo para mapear estados de TensorIR similares (del mismo origen de carga de trabajo) cerca en el espacio latente, mientras aleja los estados disímiles. Esto crea un espacio de estados donde la equivalencia semántica se preserva a pesar de las variaciones sintácticas.

2. Transición de Estado de Múltiples Pasos Condicionada por Acciones

• Objetivo: Predecir la representación del estado terminal inducida por una secuencia de acciones de programación sin materializar explícitamente los AST intermedios.
• Mecanismo: Partiendo de un embedding de estado inicial $z_0$, el modelo ejecuta el despliegue (rollout) del efecto de una secuencia de acciones de programación $a_{1:T}$ paso a paso.
  • Las acciones se serializan y se codifican en embeddings densos.
  • Un modelo de transición $F(\cdot)$, implementado mediante una arquitectura basada en TransH (original de la representación de grafos de conocimiento), predice el siguiente estado latente: $z_t = F(z_{t-1}, a_t)$.
  • Esto modela las acciones de programación como transformaciones geométricas en un hiperplano específico de la relación, rastreando eficazmente la procedencia de la decisión y resolviendo ambigüedades referenciales (por ejemplo, qué representa una variable de bucle después de una fusión).
• Entrenamiento: Utiliza una pérdida de despliegue de múltiples pasos que combina un término de coincidencia de estado (minimizando la distancia al estado futuro real) y un término de clasificación basado en márgenes (asegurando que el estado predicho esté más cerca de la verdad fundamental que de las muestras negativas).

3. Evaluación de Candidatos Basada en Clasificación

• Entrada: La representación del estado terminal predicho $\hat{z}_T$, características de acción estructuradas (por ejemplo, conteos de acciones, longitud de secuencia) y características de hardware opcionales.
• Salida: Un puntaje escalar para la priorización relativa de candidatos dentro de la misma carga de trabajo.
• Arquitectura: Un modelo de clasificación XGBoost entrenado con un objetivo LambdaRank.
• Estrategia de Entrenamiento: El modelo no predice el tiempo de ejecución absoluto. En su lugar, aprende a clasificar candidatos dentro de un grupo de carga de trabajo basándose en los tiempos de ejecución medidos, tratando el problema como una tarea de aprendizaje de clasificación (learning-to-rank).

Principales Contribuciones

1. Perspectiva de Modelo de Mundo para Compiladores: El artículo adapta el paradigma de modelo de mundo a la optimización de tensores, estableciendo un marco donde los esquemas candidatos se modelan a través de la evolución del estado latente condicionado por acciones en lugar de instantáneas estáticas.
2. Conjunto de Datos de Trayectoria Acción-Estado: Construcción de un conjunto de datos basado en TVM/TenSet que organiza los registros de ajuste en trayectorias alineadas de acción-estado (estado pre-programación, secuencia de acciones, estados intermedios, estado post-programación) para la predicción de múltiples pasos.
3. Integración de Búsqueda Eficiente: Implementación del marco en el AutoScheduler de TVM, demostrando que el modelado de la dinámica latente reduce significativamente los costos de auto-ajuste manteniendo una alta calidad de búsqueda.

Resultados Experimentales

El método fue evaluado en un CPU Intel Xeon Gold 6430 y una GPU NVIDIA RTX 4090 a través de siete modelos de redes neuronales (incluyendo ResNet, BERT, GPT y OPT) y 103 subgrafos extraídos.

• Rendimiento vs. Ansor: Bajo un presupuesto de 64 ensayos, el método mejora la latencia del subgrafo representativo sobre Ansor en 1.37× en GPU y 1.54× en CPU.
• Eficiencia de Muestreo: El método logra una calidad de búsqueda comparable a Ansor-10K (10,000 ensayos) utilizando solo 1,000 ensayos (10× menos mediciones). Específicamente, supera ligeramente a Ansor-10K en la media aritmética y se iguala con este dentro de un 2.2% en la media geométrica en subgrafos representativos.
• Aceleración de Extremo a Extremo: Comparado con PyTorch/PyTorch-opt (cuDNN), el método ofrece aceleraciones de media geométrica de 4.61× (CPU) y 3.67× (GPU), con aceleraciones pico que alcanzan hasta 58.61×.
• Estudios de Ablación: Eliminar el módulo de transición aprendido (revirtiendo a la codificación de estado inicial o final estática) resultó en aumentos significativos de latencia (mediana ~56–69%), confirmando que el modelado de la evolución del estado intermedio es el componente más crítico.

Significancia y Limitaciones

Significancia:
El artículo argumenta que tratar la optimización del compilador como un proceso dinámico de transiciones de estado, en lugar de un problema de evaluación de código estático, permite a los evaluadores capturar las dependencias semánticas entre las decisiones de programación. Este enfoque suprime eficazmente el ruido sintáctico y resuelve las ambigüedidades referenciales, lo que conduce a una búsqueda más eficiente con menos mediciones de hardware.

Limitaciones (según los autores):

1. Dependencia del Espacio de Búsqueda: El método mejora la evaluación de candidatos, no su generación. Su efectividad está limitada por la capacidad del programador subyacente para proponer esquemas sólidos; un mejor evaluador no puede recuperar esquemas que nunca fueron propuestos.
2. Acumulación de Errores: Los errores de predicción pueden acumularse en trayectorias de programación largas o complejas en el espacio latente, reduciendo potencialmente la fiabilidad en casos difíciles.
3. Clasificación vs. Predicción: El modelo está entrenado para la clasificación dentro de una carga de trabajo, no para la predicción de latencia absoluta. Esto lo hace ideal para la priorización en línea, pero menos adecuado para entornos que requieren estimaciones de rendimiento calibradas entre diferentes cargas de trabajo.