Esta página foi traduzida pela API Cloud Translation.
Switch to English

Código MLIR CodeGen for XLA

O XLA opera no HloInstruction e realiza muitas otimizações nessa representação, compartilhando muitas delas entre dispositivos de HloInstruction . Em algum momento, uma programação linear é calculada e o buffer de memória é atribuído estaticamente a cada valor. O codegen específico do dispositivo opera atravessando essa sequência e chamando "emissores" para gerar uma representação adequada para o dispositivo (por exemplo, uma única função LLVM por computação XLA na CPU, ou uma sequência de "thunks" que encapsula as operações da GPU e possivelmente gera PTX quando GPU de segmentação).

Como etapa de preparação, estamos atualmente no processo de interceptar o processo logo após o XLA concluir a fase de atribuição do buffer e emitir um módulo MLIR no dialeto lhlo . A partir daí, executamos o codegen usando componentes MLIR (principalmente Linalg, affine e dialeto GPU), dependendo do dispositivo.

Abaixo está o plano de registro para migrar XLA / GPU de forma lhlo usando lhlo como entrada de codegen.

Tarefas

Hospedeiro Dispositivo
Formato de entrada HloInstruction * (Tarefa 1) HloInstruction * (Tarefa 1)
Formato de saída xla :: Thunk (tarefa 2) RI LLVM (tarefa 3)
  • A tarefa 1 altera o formato de entrada do host e do dispositivo de HloInstruction * para LHLO.
  • A tarefa 2 altera o formato de saída do host de thunks para "alguma plataforma de aterrissagem para o host" (veja abaixo).
  • A tarefa 3 migra a saída do dispositivo do LLVM IR para alguma forma de MLIR. É opcional para este projeto e consulte a seção "Migrando dispositivo LLVM IR" para obter detalhes.

Este projeto prioriza ter modelos executáveis ​​de ponta a ponta com emissores LHLO habilitados o máximo possível. Isso implica que a seguinte lista de objetivos por ordem de prioridade:

  • Torne o XLA / GPU executável com emissores LHLO, com Thunks e emissores existentes não modificados.
  • Elimine as referências ao HloInstruction * no LHLO, caso a caso:
    • Alterne um emissor herdado para um emissor baseado em MLIR (por exemplo, Linalg) ou
    • Traduza mecanicamente o emissor existente para obter a representação MLIR (migre para o padrão com o GPU Dialect).

Migrando Thunks (Tarefa 2)

xla :: gpu :: Thunk é uma estrutura de dados que:

  • Pode ser chamado a partir do host (xla :: gpu :: Thunk :: ExecuteOnStream ()).
  • Carrega vários dados em suas subclasses.
  • Interage com BufferAllocation :: Slice e StreamExecutor.
  • Lança kernels
  • Chamadas em todas as bibliotecas de tempo de execução.

O custo disso inclui:

  • Representando dados de configuração específicos da operação (por exemplo, configurações de convolução).
  • Migrando formas op e formas de operando.
  • Representando uma árvore de thunks (while, condition, etc).

O trabalho de migração é independente da migração LHLO / emissor. Com recursos limitados, é priorizado por trás da migração de LHLO / emissor.

Temos várias opções de como baixar a parte do lado do host do LHLO:

  • TFRT
    • (Pro) ótimos wrappers CUDA e HIP para uso.
    • (Pro) fácil de implementar chamadas de biblioteca (cuDNN, cuBLAS, cuFFT, etc), pois as operações TFRT são interpretadas pelo código C ++.
    • O lado (host) do host está em desenvolvimento e não foi testado.
  • Código da CPU emitida
    • (Pro) grande capacidade inferior. Crie alguns loops e condições e pronto.
    • (Con) GPUDialect ainda não modela cadeias / fluxos / assincronicidade / alocação de dispositivo.
    • (Con) O suporte ao tempo de execução CUDA / HIP é mínimo (caminho do kit de ferramentas, versão, carregamento dinâmico, etc.).
  • Tempo de execução XLA existente (interpretação)

Decisão: adote o TFRT, mas também ofereça suporte ao código da CPU no TFRT.

Migrando dispositivo LLVM IR (tarefa 3)

Um emissor elementar gera um alvo op preenchendo-o elemento por elemento. Cada elemento de saída depende de um conjunto de elementos dos operandos. Todos os elementos são descritos combinando o buffer com índices dinâmicos. É suficiente descrever quase todas as operações "matemáticas", mas por razões de desempenho, apenas um grande subconjunto de operações "matemáticas" é implementado diretamente no ElementalIrEmitter (Cpu | Gpu).

ElementalIrEmitter é único em que:

  • Uma grande parte do código é compartilhada entre XLA / GPU e CPU.
  • Representa uma grande parte das operações vistas nos modelos, incluindo todas as operações em elementos.
  • A maioria das fusões depende apenas do ElementalIrEmitter.
  • É estruturalmente simples, pois descreve um DAG de dependência de dados entre elementos op e elementos operand.
  • É principalmente portátil e de alto nível (por exemplo, ao contrário do GPU kReduce e GPU kCopy).
  • O suporte a formas dinâmicas é fácil para, pelo menos, operações por elementos.

Agora, para todas as operações, emitidas ou não elementarmente, existem vários tipos de estado final de cada operação XLA:

  1. O código do dispositivo permanece como LLVM IR.
  2. Refatore o emissor antigo para ser como LHLO -> MLIR LLVM Dialect:
    • (Custo) Será um trabalho descartável, se quisermos migrar para o Standard.
    • (Benefício) É fácil e mecânico. Pode ser feito em um curto período.
    • (Benefício) Não beneficia mais em comparação com (1).
  3. Refatore os emissores antigos para serem como LHLO -> GPU MLIR + Padrão + Loops:
    • (Custo) A elevação dos emissores existentes para o padrão apresenta alguns desafios. Ponteiros e GEPs precisam ser convertidos em MemRefs e SubViews. Garantir a integridade do amdgpu é outro.
    • (Custo) XLA / GPU depende muito dos metadados do LLVM:
      • range para índices de bloco / rosca.
      • align , dereferenceable , invariant.load , alias.scope , noalias para carga / armazenamento.
      • llvm.loop.unroll.disable , llvm.loop.unroll.full , llvm.loop.vectorize.enable para loops seqüenciais.
    • (Benefício) Pode ser de longo prazo. Mais portátil.
  4. Refatore emissores antigos para serem LHLO -> Linalg e escreva novos emissores Linalg
    • (Custo) Isso é caso a caso. Comparado às opções anteriores, uma nova implementação que corresponda ao desempenho do XLA precisa passar pelo benchmark <-> otimizar o fluxo de trabalho, o que pode ser um custo significativo para algumas operações.
    • Pilha unificada (benefício); suporte da comunidade; portabilidade; mais potenciais de otimização.

Conclusões:

  • Não vá para (2). (1) ou (3) são apenas melhores que (2). (2) custa mais que (1), pois exige muita refatoração mecânica. Com (1), ainda podemos alcançar o objetivo de permitir que o XLA pegue emissores MLIR. Isso é feito ao LHLO -> LLVM IR -> executar emissores de dispositivos herdados.
  • As operações de ElementalIrEmitter são válidas para (4), mas não incrementalmente. Não há como fazê-lo op por op, porque todas as operações emitidas por elementos estão conectadas no mesmo gráfico. Este trabalho também pode servir como um ponto de unificação de várias forças em andamento (xla / service / mlir_gpu, o gerador de kernel, Linalg).
  • Todas as outras operações vão para (1). Como uma meta estendida, eles podem ser migrados para (3) ou (4).

Priorização

Embora todas as três tarefas mencionadas acima sejam paralelizáveis, com recursos limitados, elas precisam ser serializadas. A priorização se concentra nos resultados visíveis para a conclusão de cada tarefa.

A priorização é: Tarefa1 (LHLO para emissores herdados)> Tarefa 2 (Thunks)> Tarefa 3 (emissores MLIR).

No final da Tarefa 1, os usuários do XLA podem gerar um LHLO (por exemplo, gerador de kernel) e executá-los. O formato de compilação não será MLIR serializável.

No final da Tarefa 2, o LHLO reduz para o MLIR adequado e serializável. Isso permite a compilação offline.

No final da Tarefa 3, todos os emissores XLA são baseados em MLIR em sua implementação.

Projeto detalhado

Etapa 1: (Tarefa 1) Conclua o LHLO e faça com que os emissores herdados obtenham o LHLO

Esta etapa faz com que todos os emissores XLA / GPU existentes interajam com as operações MLIR. Esta etapa é pura refatoração e NFC.

Esta etapa é principalmente mecânica, mas vale a pena observar as seguintes discrepâncias entre um HloComputation não aninhado e um LHLO:

  • Cada HloInstruction tem acesso direto aos seus operandos (um DAG de fluxo de dados). Pelo contrário, cada operação LHLO só tem acesso aos seus buffers de operando (um bipartido entre ops e buffers). As operações da LHLO precisam passar por cadeias de uso-definição para acessar suas operações.
  • Emissores legados não aninhados empiricamente quase nunca acessam seus operandos. A única exceção é o kReduce.
  • Os emissores legados não aninhados acessam o BufferAssignment apenas para obter fatias, não para acessar estruturas de dados auxiliar como dataflow_analysis () ou alias_analysis (). O llvm_ir cria seu próprio alias_analysis () com base nas informações da fatia.

A conclusão é que o LHLO deve se encaixar corretamente, sem grandes problemas.

Etapa 2: (opcional) suporte a criação de perfil

Essa etapa é necessária apenas se começarmos a descartar parte da lógica do XLA Thunk (consulte a próxima etapa).

Antes de realmente ativar quaisquer emissores baseados em MLIR, precisamos criar um perfil para emissores baseados em MLIR.

Atualmente, o XLA executa seu próprio perfil chamando o timer do StreamExecutor. O timer sob o capô insere dois eventos antes e após o lançamento do kernel e mede o tempo de sincronização entre esses dois eventos.

Existem aproximadamente três abordagens para oferecer suporte à criação de perfil no MLIR:

  • Execute um criador de perfil de ponta a ponta
  • Adicione um op de perfil para cada op no LHLO, usando um criador de perfil injetado.

A abordagem "de ponta a ponta" é transparente ao MLIR, mas sofre o mesmo problema que faz com que o XLA não o use: as chamadas de biblioteca coletadas por um criador de perfil (nvprof / ...) não podem se relacionar facilmente ao HLO ops. Por exemplo, o cuDNN lança vários kernels para cada HLO, e é difícil dizer quais kernels correspondem a qual HLO.

A abordagem "injected profiler" requer:

  • LHLO para usar um perfilador como parâmetro.
  • inserindo profile.start / profile.end antes e depois de cada op.
  • uma passagem desse perfil inferior. {start, end} para uma implementação em C ++.

O perfil exato não pode ser feito facilmente para operações geradas por MLIR, pois:

  • O MLIR não possui um timer nem depende do TFRT / StreamExecutor.
  • O MLIR não chama facilmente as funções C com parâmetros complicados.

Etapa 3: (Tarefa 2) Migrando Thunks

Como nota, existem aproximadamente três tipos de thunks:

  • KernelThunk, que lança um kernel.
  • Controle os thunks de fluxo, que possuem lógica de fluxo de controle do host (condicional, while, for, sequence) e inicie os kernels do corpo.
  • Thunks de biblioteca: cuDNN, cuBLAS, cuFFT, NCCL, etc.

O plano é:

  • Tornar o Thunks (des) serializável.
  • Ajude a melhorar o TFRT para um estado em que ele possa suportar essas semânticas.
  • À medida que o estado melhora, migre os thunks individuais de forma incremental.

Esses itens de ação são solicitados apenas parcialmente. A ordem de execução real / paralelismo de engenharia deve ser avaliada à medida que avança.

Etapa 4: (Tarefa 3) Migrated ElementalIrEmitter

Quando o perfil estiver pronto, podemos concluir e ajustar todos os emissores baseados em ElementalIrEmitter no MLIR. Em seguida, ativamos-os por padrão, assumindo que todos esses emissores baseados em MLIR usem um único fluxo.

Observe que também é benéfico migrar o ElementalIrEmitter do XLA / CPU, pois eles compartilham uma grande parte do código.

Com todos os testes de comparação e desempenho realizados (TODO: definir paridade de desempenho), ativamos o novo emissor elementar baseado em MLIR e excluímos o ElementalIrEmitter herdado.

Esta etapa também fornece transições de fusão fáceis (operações aninhadas) para a migração posterior.

Etapa 5: Suporte para Multi-Stream ou Drop

Não podemos excluir alguns dos emissores até que o suportemos no MLIR, ou eliminamos o recurso. É uma quantidade relativamente grande de trabalho no MLIR e um pequeno ganho no XLA. Devemos investigar os usuários atuais de usuários XLA / GPU de fluxo múltiplo e tentar excluir esse recurso, se for razoável.

Etapa 6: (Tarefa 3) Operações de dispositivo migradas

Esta etapa migra todas as operações não aninhadas, e podemos excluir todos os emissores não aninhados.

Isso requer uma reescrita / refatoração para o kCopy e o kReduce. O kReduce já é trabalhado em abundância, portanto, a quantidade real de trabalho que precisa ser feita permanece para ser vista.