No momento da escrita, o StableHLO está pronto para substituir o MHLO/HLO como interface do compilador. Ele pode ser produzido pelo TensorFlow, JAX e PyTorch, pode ser consumido pelo XLA e vários plug-ins PJRT de terceiros e tem todos os recursos públicos fornecidos pelo MHLO/HLO, além de outras funcionalidades.
Este documento descreve as próximas etapas do projeto StableHLO, categorizando o trabalho em andamento refletido no rastreador de problemas e organizando esse trabalho em entregas planejadas.
Marcos atuais
Nossos marcos atuais seguem duas tendências principais:
- Maximizar o benefício do StableHLO para toda a comunidade do OpenXLA.
- Unificar a experiência do desenvolvedor para todos os membros do OpenXLA.
- Descontinuação da MHLO: no 4º trimestre de 2024, começamos a explorar a descontinuação interna da MHLO, migrando passagens úteis, incluindo canonização e padrões de pastas, para a StableHLO. Quando o processo de migração for considerado trivial internamente, planejamos compartilhar um RFC com cronogramas para migrações externas para o StableHLO. Isso provavelmente acontecerá no primeiro trimestre de 2025, e planejamos dar tempo e suporte suficientes para que as equipes migrem para o StableHLO no primeiro semestre de 2025.
- Migrar otimizações independentes de hardware para o StableHLO: seguindo a tendência acima, queremos que o StableHLO seja o melhor lugar para consolidar simplificações de gráficos independentes de hardware, para que todos os plug-ins do PJRT, incluindo aqueles que convertem do StableHLO para um IR de compilador não XLA, possam ter os benefícios máximos. Parte desse objetivo envolve consolidar padrões usados no Google AI Edge, no projeto JAX-Enzyme e em outros projetos no repositório StableHLO. Parte dessa consolidação já começou, mas o fluxo de trabalho vai ser retomado e concluído no primeiro trimestre de 2025.
- Componentização do OpenXLA: começamos a criar componentes dedicados no openxla/xla para HLO, que se assemelha à configuração do repositório StableHLO (ref), além de começar a mover todos os back-ends do OpenXLA por trás dos plug-ins do PJRT. Além disso, estamos investindo na correção de problemas de UX importantes descobertos nesses plug-ins do PJRT, incluindo a comunicação precisa da versão do StableHLO nos plug-ins do StableHLO, para que os novos recursos possam ser usados imediatamente pelos novos plug-ins (ref).
- Fazer com que os compostos funcionem e2e: no terceiro trimestre de 2024, adicionamos compostos ao HLO, ativando o suporte completo da pilha do compilador para abstrações. No 4º trimestre de 2024, ensinamos o inliner do XLA sobre composições e adicionamos transmissões no HLO/StableHLO para inline composições desconhecidas com suas decomposições. Agora estamos investigando a adição de APIs JAX dedicadas para gerar compósitos do framework (as APIs PyTorch já existem), além de adicionar documentação do Colab sobre como usar corretamente os compósitos, que será concluída no 4º trimestre de 2024.
Marcos passados
No primeiro semestre de 2024, lançamos a StableHLO v1.0, que implementou melhorias de alta prioridade, incluindo o término da especificação de opset público, o interpretador de referência, a extensão da compatibilidade para frente e para trás para oferecer suporte à implantação no dispositivo com testes de unidade de compatibilidade, extensibilidade por operações compostas, suporte a dynamism especificado, lançamento de um conjunto completo de dados de teste com resultados ideais avaliados e muito mais.