La cuestión de cómo compilar código más rápido y más pequeño surgió junto con el nacimiento de las computadoras modernas. Una mejor optimización del código puede reducir significativamente el costo de propiedad de las aplicaciones de grandes centros de datos. El tamaño del código compilado es más importante para los sistemas móviles e integrados o el software implementado en particiones de arranque seguras donde el binario compilado debe ajustarse a presupuestos de tamaño de código ajustados. A medida que el campo ha avanzado, el alcance se ha visto severamente limitado por heurísticas cada vez más complicadas, lo que dificulta el mantenimiento y otras mejoras.
Investigaciones recientes han demostrado que el aprendizaje automático (ML) puede desbloquear más posibilidades en la optimización del compilador al reemplazar heurísticas complicadas con pautas de ML. Sin embargo, la adopción de ML en compiladores industriales de propósito general sigue siendo un desafío.
Para abordar esto, presentamos MLGO: un marco de optimizaciones de compilador guiadas por aprendizaje automático, el primer marco de uso general de grado industrial para la integración sistemática de técnicas de ML en LLVM (una infraestructura de compilador industrial de código abierto que se utiliza para crear aplicaciones de misión crítica, software de alto rendimiento). MLGO usa Reinforcement Learning (RL) para entrenar redes neuronales para tomar decisiones que pueden reemplazar la heurística en LLVM. Describimos dos optimizaciones de MLGO para LLVM: 1) reducción del tamaño del código a través de la inserción; y 2) mejorar el rendimiento del código con la asignación de registros (regalloc). Ambas optimizaciones están disponibles en el repositorio LLVM y se implementaron en producción.
¿Cómo funciona MLGO? Con Inlining-for-Size como caso de estudio
La inserción ayuda a reducir el tamaño del código al tomar decisiones que permiten eliminar el código redundante. En el siguiente ejemplo, la función de llamada foo() llama a la función llamada bar()quien se llama a si misma baz(). Incrustar ambos sitios de llamadas devuelve un simple foo() Función que reduce el tamaño del código.
 |
| La inserción reduce el tamaño del código al eliminar el código redundante. |
En código real, hay miles de funciones que se llaman entre sí, formando un gráfico de llamadas. Durante la fase de inserción, el compilador recorre el gráfico de llamadas para todos los pares de llamador-llamado y toma decisiones sobre si inyectar o no un par de llamador-llamado. Es un proceso de toma de decisiones secuencial, ya que las decisiones de alineación anteriores cambian el gráfico de llamadas y afectan las decisiones posteriores y el resultado final. En el ejemplo anterior, el diagrama de llamada foo() → bar() → baz() necesita una decisión de «sí» en ambos bordes para causar la reducción del tamaño del código.
Antes de MLGO, la decisión en línea/no en línea se tomaba mediante una heurística que era cada vez más difícil de mejorar con el tiempo. MLGO reemplaza la heurística con un modelo ML. Mientras atraviesa el gráfico de llamadas, el compilador recibe consejos de una red neuronal sobre si debe insertar un par de llamador-llamado en particular mediante la alimentación de características relevantes (es decir, entradas) del gráfico, y ejecuta las decisiones secuencialmente hasta que se ejecuta todo el gráfico de llamadas. .
 |
| Ilustración de MLGO durante la inserción. «#bbs», «#users» y «callsite height» son funciones de ejemplo para pares de llamador-llamado. |
MLGO entrena la red de decisión (política) con RL utilizando algoritmos de estrategia de evolución y gradiente de política. Si bien no existe una verdad básica sobre las mejores decisiones, RL en línea itera entre la capacitación y la compilación continua con la política capacitada para recopilar datos y mejorar la política. Específicamente, dado el modelo actual en entrenamiento, el compilador consulta el modelo para la toma de decisiones en línea/no en línea durante la fase de integración. Una vez completada la compilación, se crea un registro del proceso secuencial de toma de decisiones (estado, acción, recompensa). Luego, el registro se pasa al entrenador para actualizar el modelo. Este proceso se repite hasta que conseguimos un modelo satisfactorio.
 |
Comportamiento del compilador durante el entrenamiento. El compilador compila el código fuente. foo.cpp a un archivo de objeto foo.o con una secuencia de pases de optimización, uno de los cuales es el pase en línea. |
Luego, la política entrenada se integra en el compilador para proporcionar decisiones en línea/no en línea durante la compilación. A diferencia del escenario de entrenamiento, la política no crea un registro. El modelo TensorFlow está integrado en XLA AOT, que convierte el modelo en código ejecutable. Esto evita la dependencia y la sobrecarga del tiempo de ejecución de TensorFlow, y minimiza la sobrecarga y los costos de memoria incurridos por la inferencia del modelo ML en tiempo de compilación.
 |
| Comportamiento del compilador en producción. |
Entrenamos la política de alineación por tamaño en un gran paquete de software interno con 30,000 módulos. La política entrenada es generalizable cuando se aplica para compilar otro software y logra una reducción de tamaño del 3% al 7%. Además de la generalización en todo el software, la generalización en el tiempo también es importante: tanto el software como el compilador evolucionan activamente, por lo que la política entrenada debe mantener un buen rendimiento durante un período de tiempo razonable. Verificamos el desempeño del modelo tres meses más tarde usando el mismo conjunto de software y solo encontramos un leve deterioro.
 |
| Porcentajes de reducción de tamaño de la política Inlining-for-Size. El eje X representa un software diferente y el eje Y representa la reducción porcentual del tamaño. El entrenamiento es el software en el que se entrenó el modelo e infra[1|2|3]“ son diferentes paquetes de software internos. |
La capacitación de MLGO en línea para el tamaño se brindó en Fuchsia, un sistema operativo de código abierto de propósito general diseñado para ejecutar un ecosistema diverso de hardware y software donde el tamaño binario es crítico. Aquí, MLGO mostró una reducción de tamaño del 6,3 % para las unidades de traducción de C++.
Registro de mapeo (para rendimiento)
Como marco general, usamos MLGO para mejorar el recorrido de asignación de registros, lo que mejora el rendimiento del código en LLVM. El mapeo de registros resuelve el problema de mapear registros físicos a áreas vivas (es decir, variables).
A medida que se ejecuta el código, diferentes secciones en vivo se completan en diferentes momentos, liberando registros para su uso en etapas posteriores de procesamiento. En el siguiente ejemplo, cada instrucción «sumar» y «multiplicar» requiere que todos los operandos y el resultado estén en registros físicos. La gama viva X se asigna al registro verde y se completa antes que las áreas vivas en el registro azul o en el registro amarillo. A X se completa, la pestaña verde pasa a estar disponible y se asigna al panel en vivo t.
 |
| Registrar ejemplo de mapeo. |
Cuándo es el momento de asignar el rango en vivo qno hay pestañas disponibles, por lo que el pase de mapeo de pestañas debe decidir qué área activa (si la hay) se puede «borrar» de su pestaña para dejar espacio q. Esto se llama el problema de «desalojo de rango vivo» y es la decisión que entrenamos al modelo para reemplazar las heurísticas originales. En este ejemplo particular, se elimina p.ej del registro amarillo y lo asigna q y la primera mitad p.ej.
Ahora miramos la segunda mitad no asignada del área en vivo p.ej. Tenemos otro conflicto, y esta vez es la sección en vivo t será expulsado y dividido, y la primera mitad t y la última parte de p.ej al final con el registro verde. La parte media de p.ej corresponde a las instrucciones q = t * jDónde p.ej no se utiliza, por lo que no se asigna a ningún registro y su valor se almacena en la pila del registro amarillo, que luego se recarga en el registro verde. Lo mismo sucede con t. Esto agrega instrucciones adicionales de carga/almacenamiento al código y degrada el rendimiento. El objetivo del algoritmo de mapeo de registros es reducir tales ineficiencias tanto como sea posible. Esto se utiliza como recompensa para liderar la capacitación en políticas de RL.
De manera similar a la política de alineación por tamaño, la política de asignación de registros (Regalloc-for-Performance) se entrena en un importante paquete de software interno de Google y se puede generalizar en diferentes software, con mejoras del 0,3 % al 1,5 % en consultas por segundo ( QPS) en varias aplicaciones internas de grandes centros de datos. La mejora de QPS continúa meses después de su inicio, lo que demuestra la generalización del modelo en el horizonte temporal.
Conclusión y trabajo futuro
Proponemos MLGO, un marco para integrar sistemáticamente técnicas de ML en un compilador industrial, LLVM. MLGO es un marco general que se puede ampliar para: 1) profundizar, p. B. agregar más funciones y aplicar mejores algoritmos de RL; y 2) más amplio al aplicarlo a más heurísticas de optimización más allá de inlining y regalloc. Estamos entusiasmados con las posibilidades que MLGO puede brindar al campo de la optimización del compilador y esperamos su adopción continua y futuras contribuciones de la comunidad de investigación.
Inténtalo tú mismo
Consulte la solución de capacitación y recopilación de datos de extremo a extremo de código abierto en Github y una demostración que utiliza el gradiente de la política para entrenar una política de alineación por tamaño.
Gracias
Nos gustaría agradecer a los contribuyentes y colaboradores de MLGO: Eugene Brevdo, Jacob Hegna, Gaurav Jain, David Li, Zinan Lin, Kshiteej Mahajan, Jack Morris, Girish Mururu, Jin Xin Ng, Robert Ormandi, Easwaran Raman, Ondrej Sykora, Maruf Zaber, Weiye Zhao. También nos gustaría agradecer a Petr Hosek, Yuqian Li, Roland McGrath y Haowei Wu por confiar en nosotros y usar MLGO en Fuchsia como los primeros clientes de MLGO; gracias a David Blaikie, Eric Christopher, Brooks Moses y Jordan Rupprecht por su asistencia en la implementación de MLGO en las aplicaciones internas de grandes centros de datos de Google; y gracias Ed Chi, Tipp Moseley por su apoyo de liderazgo.
