XLA (Accelerated Linear Algebra) es un compilador de aprendizaje automático (AA) que optimiza el álgebra lineal, lo que proporciona mejoras en la velocidad de ejecución y el uso de la memoria. En esta página, se proporciona una descripción general breve de los objetivos y la arquitectura del compilador XLA.
Objetivos
En la actualidad, XLA admite varios frontends de framework de AA (incluidos PyTorch, TensorFlow y JAX) y es parte del proyecto OpenXLA, un ecosistema de tecnologías de compilación de código abierto para el AA que se desarrolla de forma colaborativa mediante las organizaciones líderes en hardware y software de AA. Antes de crear el proyecto OpenXLA, XLA se desarrolló dentro del proyecto de TensorFlow, pero los objetivos fundamentales son los mismos:
Mejora la velocidad de ejecución. Compila subgrafos para reducir el tiempo de ejecución de las operaciones de corta duración y eliminar la sobrecarga del entorno de ejecución, fusiona las operaciones de canalización para reducir la sobrecarga de la memoria y especializa las formas de tensores conocidas para permitir una propagación constante más agresiva.
Mejora el uso de memoria. Analiza y programa el uso de memoria, lo que elimina muchos búferes de almacenamiento intermedios.
Reduce la dependencia de las operaciones personalizadas. Quita la necesidad de muchas ops personalizadas; para ello, mejora el rendimiento de las ops de bajo nivel fusionadas automáticamente para que coincida con el rendimiento de las ops personalizadas que se fusionaron de forma manual en un principio.
Mejora la portabilidad. Facilitar la escritura de un backend nuevo para hardware nuevo, de modo que una gran fracción de modelos de AA pueda ejecutarse sin modificaciones en ese hardware Esto contrasta con el enfoque de especializar operaciones monolíticas individuales para hardware nuevo, que requiere que los modelos se reescriban para usar esas operaciones.
Cómo funciona
El compilador de XLA toma gráficos de modelos de frameworks de AA definidos en StableHLO y los compila en instrucciones de máquina para varias arquitecturas. StableHLO define un conjunto de operaciones con control de versiones (HLO = operaciones de alto nivel) que proporciona una capa de portabilidad entre los frameworks de AA y el compilador.
En general, el proceso de compilación que convierte el grafo del modelo en un ejecutable optimizado para el destino incluye los siguientes pasos:
XLA realiza varios pases de optimización y análisis integrados en el grafo estable que son independientes del objetivo, como CSE, fusión de operaciones independientes del objetivo y análisis de búfer a fin de asignar memoria del entorno de ejecución para el cálculo. Durante esta etapa de optimización, XLA también convierte el dialecto StableHLO en uno interno.
XLA envía el cálculo de HLO a un backend para realizar más optimizaciones a nivel de HLO, esta vez con la información y las necesidades específicas del objetivo en mente. Por ejemplo, el backend de la GPU puede realizar fusiones de operaciones que son beneficiosas de manera específica para el modelo de programación de GPU y determinar cómo particionar el cálculo en transmisiones. En esta etapa, los backends también pueden coincidir con patrones por determinadas operaciones o combinaciones de estas con llamadas a bibliotecas optimizadas.
Luego, el backend realiza la generación de código de destino específico. Los backends de CPU y GPU incluidos en XLA usan LLVM para la IR, la optimización y la generación de código de bajo nivel. Estos backends emiten la IR de LLVM necesaria para representar el cálculo de HLO de manera eficiente y, luego, invocan a LLVM para emitir código nativo desde esta IR de LLVM.
Dentro de este proceso, el compilador de XLA es modular, en el sentido de que es fácil insertar un backend alternativo para orientar una arquitectura de hardware novedosa. En la actualidad, el backend de GPU admite GPU de NVIDIA a través del backend de LLVM NVPTX. El backend de la CPU admite varios ISA de la CPU.