Se usó la API de Cloud Translation para traducir esta página.
Switch to English

MLIR CodeGen para XLA

XLA opera en HloInstruction y realiza muchas optimizaciones en esta representación, compartiendo muchas de ellas entre dispositivos específicos. En algún momento, se calcula un cronograma lineal y el búfer de memoria se asigna estáticamente a cada valor. El codegen específico del dispositivo opera atravesando esta secuencia y llamando a "emisores" para generar una representación adecuada para el dispositivo (por ejemplo, una sola función LLVM por cálculo de XLA en la CPU, o una secuencia de "thunks" que encapsulan las operaciones de GPU y posiblemente genera PTX cuando GPU dirigido).

Como paso intermedio, actualmente estamos en el proceso de interceptar el proceso justo después de que XLA complete la fase de asignación de búfer y, en su lugar, emita un módulo MLIR en el dialecto lhlo . A partir de ahí, realizamos el codegen usando componentes MLIR (Linalg, afín y dialecto GPU principalmente) dependiendo del dispositivo.

A continuación se muestra el plan de registro para migrar incrementalmente XLA / GPU usando lhlo como entrada de codegen.

Tareas

Anfitrión Dispositivo
Formato de entrada HloInstruction * (Tarea 1) HloInstruction * (Tarea 1)
Formato de salida xla :: Thunk (Tarea 2) LLVM IR (Tarea 3)
  • La tarea 1 cambia el formato de entrada del host y del dispositivo de HloInstruction * a LHLO.
  • La tarea 2 cambia el formato de salida del host de thunks a "alguna plataforma de aterrizaje para el host" (ver más abajo).
  • La Tarea 3 migra la salida del dispositivo de LLVM IR a alguna forma de MLIR. Es opcional para este proyecto, y consulte la sección "Migración del dispositivo LLVM IR" para más detalles.

Este proyecto prioriza tener modelos ejecutables de extremo a extremo con emisores LHLO habilitados tanto como sea posible. Esto implica que la siguiente lista de objetivos de orden por prioridad:

  • Haga que XLA / GPU sea ejecutable con emisores LHLO, con Thunks y emisores existentes sin modificar.
  • Elimine las referencias a HloInstruction * en LHLO, caso por caso:
    • Cambie un emisor heredado a un emisor basado en MLIR (por ejemplo, Linalg), o
    • Traduce mecánicamente el emisor existente para tomar la representación MLIR (migra a Estándar con Dialecto GPU).

Migrar Thunks (Tarea 2)

xla :: gpu :: Thunk es una estructura de datos que:

  • Se puede llamar desde el host (xla :: gpu :: Thunk :: ExecuteOnStream ()).
  • Lleva varios datos en sus subclases.
  • Interactúa con BufferAllocation :: Slice y StreamExecutor.
  • Lanza granos
  • Llamadas a todas las bibliotecas de tiempo de ejecución.

El costo de eso incluye:

  • Representar datos de configuración específicos de operación (por ejemplo, configuraciones de convolución).
  • Migración de formas operativas y formas de operandos.
  • Representando un árbol de troncos (while, condición, etc.).

El trabajo de migración es independiente de la migración de LHLO / emisor. Bajo recursos limitados, se prioriza detrás de la migración LHLO / emisor.

Tenemos varias opciones sobre cómo bajar la parte del host desde LHLO:

  • TFRT
    • (Pro) excelentes envoltorios CUDA y HIP para su uso.
    • (Pro) llamadas de biblioteca fáciles de implementar (cuDNN, cuBLAS, cuFFT, etc.), ya que las operaciones TFRT se interpretan mediante código C ++.
    • El lado de (Con) host está en desarrollo y no se ha probado.
  • Código de CPU jitted
    • (Pro) gran capacidad inferior. Crea algunos bucles y condiciones y listo.
    • (Con) GPUDialect aún no modela cadenas / transmisiones / asincronía / asignación de dispositivos.
    • (Con) El soporte de tiempo de ejecución CUDA / HIP es mínimo (ruta del kit de herramientas, versión, carga dinámica, etc.).
  • Tiempo de ejecución XLA existente (interpretación)

Decisión: adoptar TFRT, pero también es compatible con el código de CPU de jitting en TFRT.

Migración del dispositivo LLVM IR (Tarea 3)

Un emisor elemental genera una operación objetivo llenándola elemento por elemento. Cada elemento de salida depende de un conjunto de elementos de los operandos. Todos los elementos se describen combinando el búfer con índices dinámicos. Es suficiente describir casi todas las operaciones "matemáticas", pero por razones de rendimiento, solo un gran subconjunto de operaciones "matemáticas" se implementa directamente en (Cpu | Gpu) ElementalIrEmitter.

ElementalIrEmitter es único en eso:

  • Una gran parte del código se comparte entre XLA / GPU y la CPU.
  • Representa una gran parte de las operaciones que se ven en los modelos, incluidas todas las operaciones basadas en elementos.
  • La mayoría de las fusiones dependen únicamente de ElementalIrEmitter.
  • Es estructuralmente simple, ya que describe un DAG de dependencia de datos entre elementos operativos y elementos de operando.
  • Es principalmente portátil y de alto nivel (por ejemplo, a diferencia de GPU kReduce y GPU kCopy).
  • El soporte de forma dinámica es fácil para al menos operaciones con elementos.

Ahora, para todas las operaciones, emitidas elementalmente o no, hay varios sabores del estado final de cada operación XLA:

  1. El código del dispositivo permanece como LLVM IR.
  2. Refactorice el antiguo emisor para que sea como LHLO -> MLIR LLVM Dialecto:
    • (Costo) Será un trabajo descartable si finalmente queremos migrar a Standard.
    • (Beneficio) Es fácil y mecánico. Se puede hacer en poco tiempo.
    • (Beneficio) No se beneficia más en comparación con (1).
  3. Refactorice los emisores antiguos para que sean como LHLO -> MLIR GPU + Standard + Loops:
    • (Costo) Levantar los emisores existentes a Standard presenta algunos desafíos. Los punteros y los GEP deben convertirse a MemRefs y SubViews. Asegurar la integridad de amdgpu es otra.
    • (Costo) XLA / GPU depende en gran medida de los metadatos de LLVM:
      • range para índices de bloque / hilo.
      • align , dereferenceable , invariant.load , alias.scope , noalias para carga / tiendas.
      • llvm.loop.unroll.disable , llvm.loop.unroll.full , llvm.loop.vectorize.enable para bucles secuenciales.
    • (Beneficio) Puede ser a largo plazo. Más portátil
  4. Refactorice los emisores antiguos para que sean LHLO -> Linalg, y escriba nuevos emisores Linalg
    • (Costo) Esto es caso por caso. En comparación con las opciones anteriores, una nueva implementación que coincida con el rendimiento de XLA debe pasar por el punto de referencia <-> optimizar el flujo de trabajo, lo que puede ser un costo significativo para algunas operaciones.
    • (Beneficio) pila unificada; soporte comunitario; portabilidad; Más potenciales de optimización.

Conclusiones:

  • No vayas por (2). (1) o (3) son simplemente mejores que (2). (2) cuesta más que (1), ya que requiere mucha refactorización mecánica. Con (1) todavía podemos lograr el objetivo de permitir que XLA recoja los emisores MLIR. Esto es haciendo LHLO -> LLVM IR -> ejecutar emisores de dispositivos heredados.
  • Las operaciones de ElementalIrEmitter van para (4), pero no de forma incremental. No hay forma de hacerlo op por op, porque todas las operaciones emitidas elementalmente están conectadas en el mismo gráfico. Este trabajo también puede servir como un punto de unificación de varias fuerzas en curso (xla / service / mlir_gpu, el generador de kernel, Linalg).
  • Todas las demás operaciones van por (1). Como objetivo de estiramiento, podrían migrarse a (3) o (4).

Priorización

Si bien las tres tareas mencionadas anteriormente son paralelizables, bajo recursos limitados deben ser serializadas. La priorización se centra en resultados visibles para la finalización de cada tarea.

La priorización es: Tarea1 (LHLO para emisores heredados)> Tarea 2 (Thunks)> Tarea 3 (emisores MLIR).

Al final de la Tarea 1, los usuarios de XLA pueden generar un LHLO (por ejemplo, un generador de kernel) y ejecutarlos. El formato de compilación no será MLIR serializable.

Al final de la Tarea 2, LHLO baja a MLIR adecuado y serializable. Esto permite la compilación fuera de línea.

Al final de la Tarea 3, todos los emisores XLA están basados ​​en MLIR en su implementación.

Diseño detallado

Paso 1: (Tarea 1) Complete LHLO y haga que los emisores heredados tomen LHLO

Este paso hace que todos los emisores XLA / GPU existentes interactúen con operaciones MLIR. Este paso es pura refactorización y NFC.

Este paso es principalmente mecánico, pero vale la pena notar las siguientes discrepancias entre una HloComputation no verificada y LHLO:

  • Cada HloInstruction tiene acceso directo a sus operandos (un DAG de flujo de datos). Por el contrario, cada LHLO op solo tiene acceso a sus buffers de operandos (un bipartito entre ops y buffers). Las operaciones de LHLO tienen que pasar por cadenas use-def para acceder a sus operaciones de operandos.
  • Los emisores heredados no verificados empíricamente casi nunca acceden a sus operandos. La única excepción es kReduce.
  • Los emisores heredados no verificados acceden a BufferAssignment solo para obtener segmentos, no para acceder a estructuras de datos auxiliares como dataflow_analysis () o alias_analysis (). llvm_ir crea su propio alias_analysis () basado en información de corte.

La conclusión es que LHLO debería encajar perfectamente sin mayores problemas.

Paso 2: (Opcional) Soporte de perfiles

Este paso solo es necesario si comenzamos a descartar parte de la lógica XLA Thunk (consulte el siguiente paso).

Antes de encender los emisores basados ​​en MLIR, necesitamos perfiles para los emisores basados ​​en MLIR.

Actualmente, XLA realiza su propio perfil llamando al temporizador de StreamExecutor. El temporizador debajo del capó inserta dos eventos antes y después del lanzamiento del núcleo, y mide el tiempo de sincronización entre estos dos eventos.

Existen aproximadamente tres enfoques para admitir la creación de perfiles en MLIR:

  • Ejecute un generador de perfiles de extremo a extremo
  • Agregue una operación de perfil para cada operación en LHLO, utilizando un generador de perfiles inyectado.

El enfoque de "extremo a extremo" es transparente para MLIR, pero sufre el mismo problema que hace que XLA no lo use en primer lugar: las llamadas a la biblioteca recopiladas por un generador de perfiles (nvprof / ...) no pueden relacionarse fácilmente con HLO ops. Por ejemplo, cuDNN lanza múltiples núcleos para cada HLO, y es difícil saber qué núcleos corresponden a qué HLO.

El enfoque del "perfilador inyectado" requiere:

  • LHLO para tomar un perfilador como parámetro.
  • insertando profile.start / profile.end antes y después de cada operación.
  • un pase de ese perfil inferior. {inicio, fin} a una implementación de C ++.

El perfil exacto no se puede hacer fácilmente para operaciones generadas por MLIR, ya que:

  • MLIR no tiene un temporizador, ni depende de TFRT / StreamExecutor.
  • MLIR no llama fácilmente a funciones C con parámetros complicados.

Paso 3: (Tarea 2) Migración de Thunks

Como nota, hay aproximadamente tres tipos de thunks:

  • KernelThunk, que lanza un kernel.
  • Controle los thunks de flujo, que tiene lógica de flujo de control de host (condicional, while, for, secuencia) y lanza kernels del cuerpo.
  • Thunks de la biblioteca: cuDNN, cuBLAS, cuFFT, NCCL, etc.

El plan es:

  • Hacer Thunks (de) serializable.
  • Ayude a mejorar la TFRT a un estado en el que pueda soportar esta semántica
  • A medida que el estado mejora, migre thunks individuales de forma incremental.

Estos elementos de acción solo se ordenan parcialmente. El orden de ejecución real / paralelismo de ingeniería debe evaluarse a medida que avanza.

Paso 4: (Tarea 3) Migrated ElementalIrEmitter

Una vez que el perfil está listo, podemos completar y sintonizar todos los emisores basados ​​en ElementalIrEmitter en MLIR. Luego los activamos de manera predeterminada, suponiendo que todos estos emisores basados ​​en MLIR usan una sola secuencia.

Tenga en cuenta que también es beneficioso migrar ElementalIrEmitter de XLA / CPU, ya que comparten una gran parte del código.

Con toda la evaluación comparativa y la búsqueda de rendimiento realizada (TODO: definir la paridad de rendimiento), activamos el nuevo emisor elemental basado en MLIR y eliminamos el ElementalIrEmitter heredado.

Este paso también proporciona transiciones de fusión fáciles (operaciones anidadas) para la migración posterior.

Paso 5: Soporte de transmisión múltiple o soltar

No podemos eliminar algunos de los emisores hasta que lo admitamos en MLIR o dejemos de funcionar. Es una cantidad relativamente grande de trabajo en MLIR y una pequeña cantidad de ganancia para XLA. Deberíamos investigar a los usuarios actuales de usuarios de XLA / GPU de transmisión múltiple e intentar eliminar esta función si es razonable.

Paso 6: (Tarea 3) Operaciones de dispositivos migrados

Este paso migra todas las operaciones no activadas, luego podemos eliminar todos los emisores no activados.

Esto requiere una reescritura / refactorización para kCopy y kReduce. kReduce ya se ha trabajado en abundancia, por lo que aún queda por ver la cantidad real de trabajo que debe hacerse.