En este documento, se describe el recorrido de un módulo del optimizador de alto nivel (HLO) de XLA desde su estado inicial hasta un ejecutable final. A veces, omitimos la palabra "módulo" y nos referimos a él solo como "HLO".
HLO previo a la optimización
Comenzamos con el módulo de HLO previo a la optimización. El HLO previo a la optimización no contiene operaciones (ops) que se consideran internas para XLA, como fusion o bitcast. En esta etapa, las operaciones no tienen un diseño o, si lo tienen, se ignorará. El HLO previo a la optimización suele producirse con frameworks de nivel superior, como TensorFlow y JAX. Cuando se usa la marca de XLA -xla_dump_to, el HLO previo a la optimización se vuelca en un archivo con el sufijo "before_optimizations.txt".
Módulo Optimize HLO
La canalización de XLA:GPU convierte el HLO previo a la optimización en HLO optimizado ejecutando una secuencia de pases. Los pases se pueden agrupar de forma semántica y ejecutar en el siguiente orden:
Fragmentación de pases relacionados
Esto incluye pases como ShardyPartitioner o los de la fragmentación SPMD.
Pases de optimización
Esto puede incluir tanto pases de legalización como pases de simplificación.
Pases de optimización colectiva
Es similar a los pases de optimización, pero se enfoca en las operaciones colectivas.
Pases de asignación de diseño
A cada operación de HLO se le asigna un diseño que forma parte de la forma de la instrucción. El diseño controla cómo se dispone el tensor físicamente en la memoria.
Ejemplo de una forma con diseño:
f32[10,20,30]{2,0,1}
Después del tipo de elemento, se encuentran las dimensiones lógicas de la forma, seguidas de la permutación de diseño en orden de menor a mayor. En este ejemplo, la dimensión más secundaria es 30, la segunda dimensión más secundaria es 10 y la dimensión principal es 20.
El objetivo de la asignación de diseño es minimizar la cantidad de transposiciones físicas necesarias con una estrategia greedy. Comienza con ciertas restricciones de diseño (p.ej., las bibliotecas cuDNN/cuBLAS esperan dimensiones consecutivas) y propaga los diseños “hacia abajo” y, luego, “hacia arriba” en el gráfico de HLO. Al final de la propagación del diseño, es posible que algunas instrucciones tengan diseños en conflicto, uno propagado desde un operando y otro propagado desde un usuario. Para resolver este conflicto, se inserta una instrucción de HLO de copy que cambia el diseño del diseño del operando al diseño de la instrucción.
Pases de normalización del diseño
Dado que es algo difícil determinar la forma física, la normalización del diseño intenta reescribir la forma de modo que use el diseño predeterminado {rank-1, rank-2, …, 0}. En el ejemplo anterior, la forma normalizada sería f32[20,10,30]{2,1,0}. Las operaciones de copia que cambian los diseños se reescriben como una combinación de transpose y bitcast. Dado que actualmente no podemos normalizar todas las operaciones, aún hay algunas que pueden tener diseños no predeterminados, en particular gather y dot. En los límites entre las operaciones normalizadas y las no normalizadas, habrá operaciones bitcast que representen una transposición, es decir, una transposición con un diseño asignado que la convierta en una operación nula físicamente.
La normalización del diseño también hace que algunas transposiciones implícitas sean explícitas, lo que es importante porque la generación de código puede controlar las transposiciones explícitas con un emisor dedicado. Por ejemplo, técnicamente, se permite que un cambio de forma tenga un diseño físico diferente entre el operando y el resultado (p.ej., debido a un rango diferente). El pase ReshapeDecomposer que se ejecuta como parte de los pases de normalización del diseño convierte un cambio de forma en una secuencia de transpose, cambio de forma bitcast y transpose.
Pases de optimización posteriores a la asignación del diseño
Los pases más importantes aquí son las fusiones de Triton (fusiones de GEMM + fusiones de Softmax/Layernorm) o las reescrituras en llamadas de biblioteca. En este paso, también se ejecuta el ajuste automático, en el que XLA elige entre diferentes emisores, selecciona el mejor algoritmo para las convoluciones o los productos escalares, encuentra la mejor división en mosaicos para las fusiones controladas por el emisor de Triton, etcétera.
Pases de fusión
Los dos pases principales son la fusión de PriorityFusion y Multi-Output.
En PriorityFusion, formamos fusiones guiadas por el modelo de costos. Cuando se fusionan, permitimos duplicar operaciones con varios usuarios si la operación se puede fusionar en todos los usuarios. También permitiríamos extender las fusiones existentes de Softmax de Triton si fuera posible.
La fusión de Multi-Output es un pase independiente que permite fusionar operaciones o fusiones que comparten un operando. También puede fusionar operandos o fusiones de operandos en usuarios sin duplicación agregando salidas adicionales, de modo que otros usuarios de la operación que se fusionará puedan redireccionarse a estas salidas. Este paso debe tener cuidado de no introducir ciclos en el grafo de HLO.
Después de la fusión de múltiples resultados, se ejecuta la eliminación de subexpresiones comunes (pase HloCSE), lo que podría volver a fusionar las operaciones duplicadas previamente si terminaron en la misma fusión.
Varios pases posteriores a la fusión
Se realizaron varios pases relacionados con los colectivos (por ejemplo, convertirlos en asíncronos o aplicar un orden relativo determinado de los colectivos).
Por último, ejecutamos CopyInsertion, donde se agregan copias para garantizar que las operaciones in situ no sobrescriban datos que aún se necesitan en otro lugar.
Al final de la optimización, se vuelca el HLO optimizado si se usa la marca -xla_dump_to en un archivo que tiene el sufijo de nombre de archivo "after_optimizations.txt". Si deseas volcar el HLO después de los pases intermedios que realmente cambian el HloModule, puedes usar la marca -xla_dump_hlo_pass_re=.* (o una expresión regular específica para restringirlo a ciertos pases).
Programación
Un módulo de HLO sin una programación aún tiene cierto grado de libertad en el orden en que se procesan las operaciones. Cualquier ordenamiento topológico que respete las relaciones de operandos y resultados, y las dependencias de control es válido. La programación determina qué orden específico usar. La principal preocupación en esta etapa es el consumo máximo de memoria, que depende de la vida útil de los tensores. En un paso inicial, probamos diferentes algoritmos de programación y elegimos el programa que debería minimizar el consumo máximo de memoria. Ten en cuenta que, en este punto, aún no trabajamos con búferes físicos (eso sucederá en "Asignación de búferes") y simulamos el uso de memoria.
Luego, se ejecuta el pase LatencyHidingScheduler y se intenta maximizar la superposición entre el procesamiento y la comunicación. Sin embargo, esto puede aumentar el uso de memoria nuevamente.
Por último, en caso de que el consumo máximo de memoria sea superior a la cantidad de memoria disponible, ejecutamos HloRematerialization. Este paso intenta reducir el uso de memoria a costa del rendimiento, ya que, por ejemplo, algunas fusiones podrían dividirse y algunas operaciones podrían duplicarse para tener una vida útil más corta del búfer. Si se produce una rematerialización, puede ser beneficioso investigar formas de reducir los requisitos de memoria del modelo (p.ej., usar tamaños de lote más pequeños).
Asignación de búfer
Inmediatamente antes de que bajemos al LLVM IR, ejecutamos los pases de asignación de búfer que asignarán segmentos de búfer a cada instrucción del gráfico de HLO. La asignación de búferes se ejecuta en varios pasos:
HloDataflowAnalysisasignaHloValues(básicamente, búferes lógicos) a las instrucciones. En el caso de las operaciones in situ, se puede reutilizar elHloValuede un operando. Un op puede definir más de unHloValue(p.ej., con una forma de resultado de tupla).HloAliasAnalysisintenta combinar búferes para las operaciones de alias y calcula una asignación deHloValueaHloBuffer.BufferAssignmentcalcula una asignación deHloBuffersa segmentos de búfer dentro de un búfer grande de tal manera que no se use el mismo segmento de búfer para diferentesHloBufferscon períodos de vida superpuestos. En el caso de las operaciones que pueden tener alias, está bien que haya una ligera superposición (la hora de finalización de unHloBufferpuede coincidir con la hora de inicio del otroHloBuffer). Cuando se usa la marca-xla_dump_to, se vuelca información sobre la asignación de búferes en un archivo con el sufijo de nombre "after_optimizations-buffer-assignment.txt".
Thunks
Después de que se optimiza y programa un gráfico de HLO, se reduce a una secuencia lineal de thunks para un backend específico (CPU o GPU).
En XLA, un Thunk es una abstracción de una unidad de trabajo autónoma que ejecuta el tiempo de ejecución. Puede ser un inicio de kernel compilado, una operación específica, una llamada a la biblioteca, una construcción de flujo de control, una comunicación colectiva, etcétera. Una secuencia de thunk representa todo el ejecutable de un backend específico.
Emisión de Thunk
El proceso de convertir un cálculo de HLO programado en una secuencia de thunks se denomina "emisión de thunks". Esto se controla con una clase de emisor dedicada en cada backend.
En el caso del backend de GPU, IrEmitterUnnested controla esto.
EmitHloComputation itera a través de la lista programada de instrucciones de HLO en un cálculo y envía a un método Emit... especializado (p.ej.,
EmitFusion, EmitConvolutionThunk, EmitWhile). Cada uno de estos métodos construye los objetos Thunk adecuados y los agrega a la secuencia de thunks.
En el backend de CPU, ThunkEmitter cumple este rol y se organiza de manera similar. El ThunkSequence final se incorpora en el CpuExecutable.
Ten en cuenta que cada instrucción en el cálculo de entrada de un módulo de HLO puede corresponder a ningún (kTuple, kConstant, …), uno o varios (por ejemplo, instrucción de ordenamiento) thunks en la secuencia final de thunks.
Búferes de comandos: Optimización de la ejecución en la GPU
El hardware de GPU moderno permite registrar una secuencia de operaciones de GPU (lanzamientos de kernel, copias de memoria, etcétera) una vez y, luego, reproducir la secuencia varias veces con una sobrecarga mínima de la CPU. Esta es una optimización del rendimiento fundamental, en especial para las cargas de trabajo con muchos kernels pequeños de inicio rápido. XLA usa Command Buffer como una abstracción de los grafos de CUDA o los grafos de HIP. La interfaz principal se define en GpuCommandBuffer.
Un búfer de comandos se representa en una secuencia de thunk con CommandBufferThunk.
El emisor no produce este thunk directamente a partir de las instrucciones de HLO. En cambio, esto se realiza con CommandBufferConversionPass, que se ejecuta en la propia ThunkSequence.
El pase identifica sub-secuencias contiguas de thunks compatibles (p.ej., una serie de KernelThunk y GemmThunk). Luego, reemplaza la subsecuencia encontrada por un solo CommandBufferThunk. El nuevo thunk encapsula la lógica de los thunks originales como una lista de objetos CommandBufferCmd ligeros.
Cuando un CommandBufferThunk se ejecuta por primera vez en un flujo de GPU determinado, "registra" su secuencia de comandos en un búfer de comandos de hardware. En todas las ejecuciones posteriores, simplemente emite un solo comando a la GPU para "reproducir" la secuencia grabada. Esto evita la sobrecarga de la CPU que implica iniciar cada kernel individual.
Ejecutable
El producto final de la canalización de compilación de XLA es un ejecutable autónomo y específico de la plataforma. Este objeto encapsula toda la información necesaria para ejecutar el programa compilado en un dispositivo de destino, como una CPU o una GPU. Es el puente entre el compilador y el tiempo de ejecución. Los tiempos de ejecución modernos, como PJRT, usan abstracciones de un nivel ligeramente superior (consulta PjRtExecutable), pero, en última instancia, estas encapsulan un ejecutable específico del backend.
Un Executable contiene varios datos clave que se generan durante la compilación. Si bien el contenido exacto varía según el backend, generalmente incluye lo siguiente:
Código compilado: Es el código de máquina de bajo nivel que se ejecutará en el dispositivo. En el caso de las CPU, suelen ser uno o más archivos de objeto. En el caso de las GPUs, este es el código del dispositivo compilado en formato PTX o HSACO, que se carga en la GPU en el tiempo de ejecución.
Plan de ejecución (ThunkSequence): Es el núcleo de la lógica de tiempo de ejecución. Es una secuencia lineal de objetos Thunk. Cada thunk representa una sola unidad de trabajo, como iniciar un kernel, llamar a una función de biblioteca (p.ej., cuBLAS) o controlar el flujo. El tiempo de ejecución ejecuta el programa iterando a través de esta secuencia.
Diseño de memoria (BufferAssignment): Esta pieza fundamental de metadatos, producida por BufferAssigner, describe el diseño de memoria completo para el cálculo. Especifica el tamaño de cada búfer y cómo se asigna y reutiliza la memoria para los parámetros, las salidas y los valores temporales. El tiempo de ejecución usa esto para asignar memoria del dispositivo y pasar los punteros correctos a cada thunk.
(Opcional) Módulo de HLO: Para la depuración y la generación de perfiles, el ejecutable suele conservar una referencia al HloModule final y optimizado desde el que se compiló.
El compilador coordina la creación del ejecutable final para cada backend específico. El método RunBackend de una implementación de Compiler es el paso final en el proceso de compilación, que empaqueta todos los artefactos compilados en un objeto Executable.
GpuCompiler y CpuCompiler se orientan a la GPU y la CPU, respectivamente.
Cuando un usuario llama a Execute... en un ejecutable, el tiempo de ejecución usa BufferAssignment para asignar memoria y, luego, invoca ThunkSequence para iniciar las operaciones en el dispositivo con el código compilado.