Hoja de ruta de StableHLO

En el momento de escribir este artículo, StableHLO está listo para reemplazar MHLO/HLO como interfaz del compilador. Puede ser producido por TensorFlow, JAX y PyTorch, puede ser consumido por XLA y varios complementos PJRT de terceros, y tiene todas las funciones públicas que proporciona MHLO/HLO, así como funciones adicionales.

En este documento, se describen los próximos pasos del proyecto StableHLO, que consisten en categorizar el trabajo en curso que se refleja en el seguimiento de problemas y organizarlo en entregas planificadas.

Hitos actuales

Nuestros eventos importantes actuales siguen dos tendencias principales:

  1. Maximizar el beneficio de StableHLO para toda la comunidad de OpenXLA
  2. Unificar la experiencia del desarrollador para todos los miembros de OpenXLA
  • Dar de baja MHLO: En el cuarto trimestre de 2024, comenzamos a explorar la baja interna de MHLO y migramos los pases útiles, incluidos la canonización y los patrones de carpetas, a StableHLO. Una vez que el proceso de migración se demuestre como trivial de forma interna, planeamos compartir una RFC con cronogramas para las migraciones externas a StableHLO. Es probable que esto suceda en el 1ᵉʳ trimestre de 2025, y planeamos darles tiempo y asistencia suficientes a los equipos para que migren a StableHLO en el 1ᵉʳ trimestre de 2025.
  • Migra las optimizaciones independientes del hardware a StableHLO: Siguiendo la tendencia anterior, queremos que StableHLO sea el mejor lugar para consolidar las simplificaciones de gráficos independientes del hardware, de modo que todos los complementos de PJRT, incluidos los que convierten de StableHLO a un IR de compilador que no sea XLA, puedan ver los máximos beneficios. Parte de este objetivo implica consolidar los patrones que se usan en Google AI Edge, el proyecto JAX-Enzyme y otros proyectos en el repositorio de StableHLO. Parte de esta consolidación ya comenzó, pero el flujo de trabajo se reanudará y se completará en el 1ᵉʳ trimestre de 2025.
  • Componentización de OpenXLA: Comenzamos a crear componentes dedicados en openxla/xla para HLO que se asemejan a la configuración del repositorio de StableHLO (ref) y también comenzamos a mover todos los backends de OpenXLA detrás de los complementos de PJRT. Además, estamos invirtiendo en la corrección de problemas importantes de la UX que descubrimos en estos complementos de PJRT, como tener una comunicación precisa de la versión de StableHLO en los complementos de StableHLO, de modo que los complementos nuevos puedan usar funciones nuevas de inmediato (ref).
  • Hacer que los compuestos funcionen de extremo a extremo: En el tercer trimestre de 2024, agregamos compuestos a HLO, lo que habilita la compatibilidad completa de la pila del compilador para las abstracciones. En el cuarto trimestre de 2024, le enseñamos al intercalador de XLA sobre los compuestos y agregamos pases en HLO/StableHLO para intercalar compuestos desconocidos con sus descomposiciones. Ahora estamos investigando la adición de APIs de JAX dedicadas para generar compuestos desde el framework (ya existen APIs de PyTorch), así como la adición de documentación de Colab sobre cómo usar correctamente los compuestos, que se completará en el cuarto trimestre de 2024.

Eventos importantes pasados

En el 1ᵉʳ semestre de 2024, lanzamos StableHLO v1.0, que implementó mejoras de alta prioridad, como terminar la especificación pública del conjunto de operaciones, el intérprete de referencia, extender la compatibilidad con versiones anteriores y posteriores para admitir la implementación en el dispositivo con pruebas de unidades de compatibilidad, extensibilidad a través de operaciones compuestas, compatibilidad con el dinamismo especificado, lanzar un paquete completo de datos de prueba con resultados de referencia evaluados y mucho más.