Arquitectura de XLA

XLA (Accelerated Linear Algebra) es un compilador de aprendizaje automático (AA) que optimiza el álgebra lineal y proporciona mejoras en la velocidad de ejecución y el uso de la memoria. En esta página, se proporciona una breve descripción general de los objetivos y la arquitectura del compilador de XLA.

Objetivos

En la actualidad, XLA admite varios frontends de framework de AA (incluidos PyTorch, TensorFlow y JAX) y forma parte del proyecto OpenXLA, un ecosistema de tecnologías de compilación de código abierto para el AA que se desarrollan de forma colaborativa por las organizaciones líderes de software y hardware 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 canalizadas para reducir la sobrecarga de la memoria y especializa las formas de tensor conocidas para permitir una propagación constante más agresiva.

  • Mejora el uso de la memoria. Analizar y programar el uso de la memoria, y eliminar muchos búferes de almacenamiento intermedios

  • Reduce la dependencia de las operaciones personalizadas. Para quitar la necesidad de muchas operaciones personalizadas, mejora el rendimiento de las operaciones de bajo nivel que se fusionan automáticamente para que coincida con el rendimiento de las operaciones personalizadas que se fusionaron de forma manual.

  • Mejora la portabilidad. Facilitar relativamente la escritura de un backend nuevo para hardware nuevo, de modo que una gran fracción de los modelos de AA puedan ejecutarse sin modificar en ese hardware Esto contrasta con el enfoque de especializar operaciones monolíticas individuales para hardware nuevo, que requiere que los modelos se reescriban a fin de 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 objetivos incluye los siguientes pasos:

  1. XLA realiza varios pases de análisis y optimización integrados en el grafo estable que son independientes del objetivo, como CSE, fusión de operaciones independiente 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 un dialecto HLO interno.

  2. 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 específica del destino y las necesidades en mente. Por ejemplo, el backend de la GPU puede realizar fusiones de operaciones que son beneficiosas específicamente 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 hacer coincidir determinadas operaciones o combinaciones de ellas con llamadas optimizadas a la biblioteca.

  3. Luego, el backend realiza la generación de código específico de destino. Los backends de CPU y GPU incluidos en XLA usan LLVM para IR, optimización y 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 este IR de LLVM.

En este proceso, el compilador de XLA es modular, en el sentido de que es fácil insertarlo en un backend alternativo para orientar una arquitectura de hardware novedosa. Actualmente, el backend de GPU es compatible con GPU de NVIDIA a través del backend de NVPTX de LLVM. El backend de la CPU admite varios ISA de CPU.