Do HLO aos Thunks

Este documento descreve a jornada de um módulo do otimizador de alto nível (HLO) do XLA do estado inicial até um executável final. Às vezes, omitimos o "módulo" e nos referimos a ele apenas como "HLO".

HLO pré-otimizado

Começamos com o módulo HLO de pré-otimização. O HLO pré-otimizado não contém operações (ops) consideradas internas ao XLA, como fusion ou bitcast. As operações não têm um layout nesta etapa ou, se tiverem, ele será ignorado. O HLO pré-otimizado geralmente é produzido por frameworks de nível superior, como TensorFlow e JAX. Ao usar a flag XLA -xla_dump_to, o HLO pré-otimizado é despejado em um arquivo com o sufixo "before_optimizations.txt".

Otimizar o módulo HLO

O pipeline XLA:GPU transforma o HLO pré-otimizado em HLO otimizado executando uma sequência de transmissões. Os comandos podem ser agrupados semanticamente e executados na seguinte ordem:

Fragmentação de cartões relacionados

Isso inclui transmissões como o Shardy Partitioner ou aquelas para fragmentação SPMD.

Passagens de otimização

Isso pode incluir passes de legalização e simplificação.

Passagens de otimização coletiva

Semelhante a passes de otimização, mas com foco em operações coletivas.

Passagens de atribuição de layout

Cada operação HLO recebe um layout, que faz parte da forma da instrução. O layout controla como o tensor é organizado fisicamente na memória.

Exemplo de uma forma com layout:

f32[10,20,30]{2,0,1}

Depois do tipo de elemento, há as dimensões lógicas da forma, seguidas pela permutação de layout em ordem crescente. Neste exemplo, a dimensão mais secundária é 30, a segunda mais secundária é 10 e a principal é 20.

O objetivo da atribuição de layout é minimizar o número de transposições físicas necessárias usando uma estratégia greedy. Ele começa com determinadas restrições de layout (por exemplo, as bibliotecas cuDNN/cuBLAS esperam dimensões consecutivas) e propaga layouts "para baixo" e "para cima" no gráfico HLO. Ao final da propagação do layout, algumas instruções podem ter layouts conflitantes, um propagado de um operando e outro de um usuário. Para resolver esse conflito, uma instrução HLO copy é inserida, mudando o layout do operando para o layout da instrução.

Transmissões de normalização de layout

Como é um pouco difícil descobrir o formato físico, a normalização do layout tenta reescrever o formato para que ele use o layout padrão {rank-1, rank-2, …, 0}. No exemplo acima, a forma normalizada seria f32[20,10,30]{2,1,0}. As operações de cópia que mudam layouts são reescritas como uma combinação de transpose e bitcast. Como não podemos normalizar todas as operações no momento, ainda há algumas que podem ter layouts não padrão, principalmente gather e dot. Nos limites entre operações normalizadas e não normalizadas, há operações bitcast que representam uma transposição, ou seja, uma transposição com um layout atribuído que a torna uma operação nula fisicamente.

A normalização de layout também torna explícitas algumas transposições implícitas, o que é importante porque a geração de código pode processar transposições explícitas com um emissor dedicado. Por exemplo, uma mudança de formato pode ter um layout físico diferente entre o operando e o resultado (por exemplo, devido a classificações diferentes). A transmissão ReshapeDecomposer, que é executada como parte das transmissões de normalização de layout, transforma uma mudança de formato em uma sequência de transpose, mudança de formato bitcast e transpose.

Passagens de otimização de atribuição de layout pós-

As transmissões mais importantes aqui são fusões do Triton (fusões GEMM + fusões Softmax/Layernorm) ou reescritas para chamadas de biblioteca. O ajuste automático também é executado nessa etapa, em que o XLA escolhe entre diferentes emissores, seleciona o melhor algoritmo para convoluções ou pontos, encontra o melhor bloco para fusões processadas pelo emissor do Triton etc.

Cartões de fusão

As duas transmissões principais são a fusão PriorityFusion e Multi-Output.

No PriorityFusion, formamos fusões guiadas pelo modelo de custo. Ao fazer a fusão, permitimos a duplicação de operações com vários usuários se a operação puder ser fundida em todos eles. Também permitiríamos estender as fusões Softmax do Triton atuais, se possível.

A fusão Multi-Output é uma transmissão separada que permite fundir operações/fusões que compartilham um operando. Ele também pode fundir operandos/fusões de operandos em usuários sem duplicação adicionando mais saídas, para que outros usuários da operação a ser fundida possam ser redirecionados para essas saídas. Essa transmissão precisa ter cuidado para não introduzir ciclos no gráfico HLO.

Depois da fusão de várias saídas, a eliminação de subexpressões comuns (transmissão HloCSE) é executada, possivelmente mesclando novamente operações duplicadas anteriormente se elas acabarem na mesma fusão.

Várias transmissões pós-fusão

Várias transmissões relacionadas a coletivos (como transformá-los em assíncronos ou impor uma determinada ordem relativa de coletivos).

Por fim, executamos CopyInsertion, em que cópias são adicionadas para garantir que as operações no local não substituam dados que ainda são necessários em outro lugar.

Ao final da otimização, o HLO otimizado é despejado se a flag -xla_dump_to for usada em um arquivo com o sufixo "after_optimizations.txt". Se você quiser despejar o HLO após passes intermediários que realmente mudam o HloModule, use a flag -xla_dump_hlo_pass_re=.* ou uma expressão regular específica para restringir a determinados passes.

Programação

Um módulo HLO sem programação ainda tem algum grau de liberdade na ordem em que as operações são processadas. Qualquer classificação topológica que respeite as relações de operando/resultado e as dependências de controle é válida. A programação determina qual ordem específica usar. A principal preocupação nesta etapa é o consumo máximo de memória, que depende do tempo de vida dos tensores. Em uma etapa inicial, tentamos diferentes algoritmos de agendamento e escolhemos o que minimiza o consumo máximo de memória. Neste momento, ainda não trabalhamos com buffers físicos (isso vai acontecer em "Atribuição de buffer") e simulamos o uso de memória.

Em seguida, a transmissão LatencyHidingScheduler é executada e tenta maximizar a sobreposição de computação e comunicação. Mas isso pode aumentar o uso da memória novamente.

Por fim, caso o consumo máximo de memória seja maior do que a quantidade disponível, executamos HloRematerialization. Essa transmissão tenta reduzir o uso de memória ao custo do desempenho. Por exemplo, algumas fusões podem ser divididas e algumas operações podem ser duplicadas para reduzir o tempo de vida do buffer. Se a rematerialização ocorrer, pode ser útil investigar maneiras de reduzir os requisitos de memória no lado do modelo (por exemplo, usando tamanhos de lote menores).

Atribuição de buffer

Imediatamente antes de reduzir para LLVM IR, executamos as transmissões de atribuição de buffer que atribuirão slices de buffer a cada instrução no gráfico HLO. A atribuição de buffer é executada em várias etapas:

1. HloDataflowAnalysis atribui HloValues (essencialmente buffers lógicos) a instruções. Para operações no local, o HloValue de um operando pode ser reutilizado. Uma operação pode definir mais de um HloValue (por exemplo, com um formato de resultado de tupla).

2. HloAliasAnalysis tenta combinar buffers para operações de alias e calcula um mapeamento de HloValue para HloBuffer.

3. O BufferAssignment calcula um mapeamento de HloBuffers para intervalos de buffer dentro de um buffer grande de forma que o mesmo intervalo de buffer não seja usado para diferentes HloBuffers com tempos de vida sobrepostos. Para operações que podem ter alias, é aceitável que haja uma pequena sobreposição (o horário de término de um HloBuffer pode coincidir com o horário de início de outro HloBuffer). Ao usar a flag -xla_dump_to, algumas informações sobre a atribuição de buffer são despejadas em um arquivo com o sufixo "after_optimizations-buffer-assignment.txt".

Thunks

Depois que um gráfico HLO é otimizado e programado, ele é reduzido a uma sequência linear de thunks para um backend específico (CPU ou GPU).

No XLA, um thunk é uma abstração de uma unidade de trabalho independente que o tempo de execução executa. Pode ser um lançamento de kernel compilado, uma operação específica, uma chamada de biblioteca, uma construção de fluxo de controle, uma comunicação coletiva etc. Uma sequência de thunk representa todo o executável de um back-end específico.

Emissão de thunk

O processo de conversão de um cálculo de HLO programado em uma sequência de thunks é chamado de "emissão de thunk". Isso é processado por uma classe de emissor dedicada em cada backend.

Para o back-end de GPU, isso é processado por IrEmitterUnnested. EmitHloComputation itera pela lista programada de instruções HLO em uma computação e envia para um método Emit... especializado (por exemplo, EmitFusion, EmitConvolutionThunk, EmitWhile). Cada um desses métodos cria os objetos Thunk apropriados e os anexa à sequência de thunks.

Para o back-end da CPU, o ThunkEmitter desempenha essa função e é organizado de maneira semelhante. O ThunkSequence final é incorporado ao CpuExecutable.

Cada instrução no cálculo de entrada de um módulo HLO pode corresponder a nenhum (kTuple, kConstant, ..), um ou vários (por exemplo, instrução de classificação) thunks na sequência final de thunks.

Buffers de comando: como otimizar a execução na GPU

O hardware de GPU moderno permite gravar uma sequência de operações de GPU (inícios de kernel, cópias de memória etc.) uma vez e reproduzir a sequência várias vezes com sobrecarga mínima de CPU. Essa é uma otimização de desempenho essencial, especialmente para cargas de trabalho com muitos kernels pequenos e de inicialização rápida. O XLA usa o buffer de comandos como uma abstração de gráficos CUDA ou HIP. A interface principal é definida em GpuCommandBuffer.

Um buffer de comando é representado em uma sequência de thunk por CommandBufferThunk.

O emissor não produz esse thunk diretamente das instruções HLO. Em vez disso, isso é feito por CommandBufferConversionPass, que é executado na própria ThunkSequence.

A transmissão identifica subsequências contíguas de thunks compatíveis (por exemplo, uma série de KernelThunks e GemmThunks). Em seguida, ele substitui a subsequência encontrada por um único CommandBufferThunk. O novo thunk encapsula a lógica dos thunks originais como uma lista de objetos CommandBufferCmd leves. Quando um CommandBufferThunk é executado pela primeira vez em um determinado fluxo de GPU, ele "grava" a sequência de comandos em um buffer de comandos de hardware. Em todas as execuções subsequentes, ele simplesmente emite um único comando para a GPU "reproduzir" a sequência gravada. Isso evita a sobrecarga da CPU ao iniciar cada kernel individual.

Executável

O produto final do pipeline de compilação do XLA é um executável independente e específico da plataforma. Esse objeto encapsula todas as informações necessárias para executar o programa compilado em um dispositivo de destino, como uma CPU ou GPU. É a ponte entre o compilador e o tempo de execução. Ambientes de execução modernos, como o PJRT, usam abstrações de nível um pouco mais alto (consulte PjRtExecutable), mas elas acabam encapsulando um executável específico do back-end.

Um Executable contém várias informações importantes geradas durante a compilação. Embora o conteúdo exato varie de acordo com o back-end, ele geralmente inclui:

• Código compilado: é o código de máquina de baixo nível que será executado no dispositivo. Para CPUs, geralmente são um ou mais arquivos de objeto. Para GPUs, esse é o código do dispositivo compilado no formato PTX ou HSACO, que é carregado na GPU durante a execução.

• Plano de execução (ThunkSequence): o núcleo da lógica de tempo de execução. Essa é uma sequência linear de objetos Thunk. Cada thunk representa uma única unidade de trabalho, como iniciar um kernel, chamar uma função de biblioteca (por exemplo, cuBLAS) ou processar o fluxo de controle. O tempo de execução executa o programa iterando por essa sequência.

• Layout de memória (BufferAssignment): essa parte essencial dos metadados, produzida pelo BufferAssigner, descreve o layout de memória completo para a computação. Ele especifica o tamanho de cada buffer e como a memória é alocada e reutilizada para parâmetros, saídas e valores temporários. O ambiente de execução usa isso para alocar memória do dispositivo e transmitir os ponteiros corretos para cada thunk.

• (opcional) Módulo HLO: para depuração e criação de perfis, o executável geralmente mantém uma referência ao HloModule final e otimizado de que foi compilado.

A criação do executável final é organizada pelo compilador para cada back-end específico. O método RunBackend de uma implementação do compilador é a etapa final no processo de compilação, que empacota todos os artefatos compilados em um objeto executável. GpuCompiler e CpuCompiler segmentam GPU e CPU, respectivamente.

Quando um usuário chama Execute... em um executável, o tempo de execução usa o BufferAssignment para alocar memória e invoca o ThunkSequence para iniciar as operações no dispositivo usando o código compilado.