¿Por qué creamos XLA?
Teníamos varios objetivos para que XLA funcionara con TensorFlow:
Mejora la velocidad de ejecución. Compile subgrafos para reducir el tiempo de ejecución de operaciones de corta duración para eliminar la sobrecarga del tiempo de ejecución de TensorFlow, fusionar operaciones canalizadas para reducir la sobrecarga de memoria y especializarse en formas de tensor conocidas para permitir una propagación constante más agresiva.
Mejora el uso de la memoria. Analice y programe el uso de la memoria, en principio eliminando muchos búferes de almacenamiento intermedios.
Reducir la dependencia de operaciones personalizadas. Elimine la necesidad de muchas operaciones personalizadas mejorando el rendimiento de las operaciones de bajo nivel fusionadas automáticamente para que coincida con el rendimiento de las operaciones personalizadas que se fusionaron a mano.
Reducir la huella móvil. Elimina el tiempo de ejecución de TensorFlow compilando el subgráfico con anticipación y emitiendo un par de archivos de objeto / encabezado que se pueden vincular directamente a otra aplicación. Los resultados pueden reducir la huella de la inferencia móvil en varios órdenes de magnitud.
Mejora la portabilidad. Haga que sea relativamente fácil escribir un nuevo backend para hardware novedoso, momento en el que una gran fracción de los programas de TensorFlow se ejecutarán sin modificaciones en ese hardware. Esto contrasta con el enfoque de especialización de operaciones monolíticas individuales para nuevo hardware, que requiere que los programas de TensorFlow se reescriban para hacer uso de esas operaciones.
¿Cómo actúa XLA?
El idioma de entrada a XLA se llama "HLO IR", o simplemente HLO (Operaciones de alto nivel). La semántica de HLO se describen en la operación semántica página. Es más conveniente pensar en HLO como IR compilador .
XLA toma gráficos ("cálculos") definidos en HLO y los compila en instrucciones de máquina para varias arquitecturas. XLA es modular en el sentido de que es fácil de ranura en un backend alternativa a apuntar algunos novedosa arquitectura HW . El backend de la CPU para x64 y ARM64, así como el backend de la GPU NVIDIA, se encuentran en el árbol de origen de TensorFlow.
El siguiente diagrama muestra el proceso de compilación en XLA:
XLA viene con varias optimizaciones y pases de análisis que son diana-independiente, como CSE , fusión de operación independiente de la diana, y el tampón de análisis para la asignación de memoria de tiempo de ejecución para el cálculo.
Después del paso independiente del objetivo, XLA envía el cálculo de HLO a un backend. El backend puede realizar más optimizaciones a nivel de HLO, esta vez teniendo en cuenta la información y las necesidades específicas del objetivo. Por ejemplo, el backend de XLA GPU puede realizar una fusión de operaciones beneficiosa específicamente para el modelo de programación de GPU y determinar cómo dividir el cálculo en flujos. En esta etapa, los backends también pueden hacer coincidir ciertas operaciones o combinaciones de las mismas con llamadas a bibliotecas optimizadas.
El siguiente paso es la generación de código específico para el objetivo. Los backends CPU y GPU incluyen con XLA uso LLVM para IR de bajo nivel, optimización y generación de código. Estos backends emiten el LLVM IR necesario para representar el cálculo de XLA HLO de manera eficiente, y luego invocan LLVM para emitir código nativo desde este LLVM IR.
El backend de la GPU actualmente admite GPU NVIDIA a través del backend LLVM NVPTX; el backend de la CPU admite múltiples ISA de CPU.