Ferramentas de XLA

O fluxo de trabalho de desenvolvimento do XLA geralmente é centrado na IR HLO, que representa a computação funcional isolada fornecida ao compilador. O XLA vem com várias ferramentas de linha de comando (descritas abaixo) que consomem HLO e o executam ou fornecem um estágio de compilação intermediário. Usar essas ferramentas é muito importante para um ciclo de iteração compile->modify->run rápido, já que o HLO pode ser visualizado e hackeado. Além disso, mudar e executar de forma iterativa é geralmente a maneira mais rápida de entender e corrigir um comportamento ou desempenho do XLA.

A maneira mais fácil de obter o HLO de um programa compilado com XLA geralmente é usar a variável de ambiente XLA_FLAGS:

$ XLA_FLAGS=--xla_dump_to=/tmp/myfolder ./myprogram-entry-point

que armazena todos os arquivos HLO antes da otimização na pasta especificada, além de muitos outros artefatos úteis.

[run_hlo_module] Executar módulos HLO

bazel run //xla/tools:run_hlo_module -- [flags] <filename>

A ferramenta run_hlo_module opera em HLOs pré-otimizados e, por padrão, agrupa a compilação, a execução e a comparação com a implementação do interpretador de referência. Por exemplo, a invocação comum para executar um arquivo de entrada computation.hlo em uma GPU NVIDIA e verificar se ele está correto é:

run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo

Executar vários módulos HLO

A invocação com vários módulos HLO é compatível com run_hlo_module. Para executar todos os módulos HLO de um diretório:

bazel run //xla/tools:run_hlo_module -- [flags] /dump/*before_optimizations*

[multihost_hlo_runner] Executar módulos HLO com suporte a SPMD

# Note: Binary name is `hlo_runner_main`.
bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- [flags] <filename>

O executor de HLO de vários hosts é uma ferramenta muito semelhante, com a ressalva de que ela oferece suporte a SPMD, incluindo comunicação entre hosts. Consulte Multi-Host HLO Runner para mais detalhes.

Executar vários módulos HLO com suporte a SPMD

Assim como run_hlo_module, multihost_hlo_runner também oferece suporte à invocação com vários módulos.

bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- [flags] /dump/*before_optimizations*

[hlo-opt] Compilar módulo HLO

bazel run //xla/tools:hlo-opt -- --platform=[gpu|cpu|...] [more flags] <filename>

Ao depurar ou entender o funcionamento do compilador, muitas vezes é útil receber a expansão de um hardware específico em um determinado ponto do pipeline (seja HLO, HLO otimizado, TritonIR ou LLVM), para uma determinada entrada de HLO ou StableHLO.

O hlo-opt é compatível com várias etapas de saída: PTX, HLO após otimizações, LLVM IR antes de otimizações ou TritonIR. O conjunto exato de etapas compatíveis depende da plataforma (por exemplo, PTX é específico da NVIDIA) e pode ser visto usando o comando "--list-stages":

hlo-opt --platform=CUDA --list-stages
buffer-assignment
hlo
hlo-backend
html
llvm
llvm-after-optimizations
llvm-before-optimizations
ptx

Depois de selecionar uma etapa, o usuário pode gravar o resultado da conversão para uma plataforma e um fluxo específicos:

hlo-opt --platform=cpu --stage=hlo input.hlo

que imprimiria o despejo para stdout (ou para um determinado arquivo se -o fosse especificado).

Compilação sem dispositivo para GPU

A compilação sem dispositivo não precisa de acesso a uma GPU. A compilação sem dispositivo oferece uma maneira de especificar a especificação da GPU na linha de comando (--xla_gpu_target_config_filename) para etapas em que o acesso à GPU é necessário, eliminando a necessidade de um dispositivo de GPU.

Exemplo: saída PTX sem acesso a um dispositivo de GPU:

hlo-opt  --platform=CUDA --stage=llvm  --xla_gpu_target_config_filename=/xla/tools/hlo_opt/gpu_specs/a100_pcie_80.txtpb input.hlo

As especificações das GPUs mais usadas são enviadas com o compilador, e o arquivo fornecido é a serialização de string de device_description.proto:

gpu_device_info {
  cuda_compute_capability {
    major: 8
    minor: 0
  }
  threads_per_block_limit: 1024
  threads_per_warp: 32
  shared_memory_per_block: 127152
  shared_memory_per_core: 65536
  threads_per_core_limit: 2048
  core_count: 6192
  fpus_per_core: 64
  block_dim_limit_x: 2147483647
  block_dim_limit_y: 65535
  block_dim_limit_z: 65535
  memory_bandwidth: 2039000000000
  l2_cache_size: 4194304
  clock_rate_ghz: 1.1105
  device_memory_size: 79050250240
}
platform_name: "CUDA"

Mais especificações de GPU estão em /xla/tools/hlo_opt/gpu_specs

Ajuste automático

Às vezes, a compilação pode envolver ajuste automático com base em um --stage de compilação. Para que a compilação sem dispositivo funcione, o usuário precisa
desativar o ajuste automático com --xla_gpu_autotune_level=0
ou
carregar resultados de ajuste automático preexistentes com --xla_gpu_load_autotune_results_from=<filename> (obtidos com --xla_gpu_dump_autotune_results_to=<filename>).

hlo-opt  --platform=CUDA --stage=llvm  --xla_gpu_target_config_filename=gpu_specs/a100_pcie_80.txtpb --xla_gpu_load_autotune_results_from=results.textpb input.hlo

O arquivo de ajuste automático é a serialização de texto de autotune_results.proto, com um exemplo assim:

version: 3
results {
  device: "CUDA: 8.0, Cores: 108, GPU clock: 1.41 GHz, Memory bandwidth: 1555 GB/s, L2 cache: 40 MB"
  hlo: "{\n  tmp_0 = f16[1,16,17,3]{3,2,1,0} parameter(0)\n  tmp_1 = f16[16,51]{1,0} bitcast(f16[1,16,17,3]{3,2,1,0} tmp_0)\n  tmp_2 = s8[16,17,3]{2,1,0} parameter(1)\n  tmp_3 = s8[51,16]{0,1} bitcast(s8[16,17,3]{2,1,0} tmp_2)\n  tmp_4 = f16[51,16]{0,1} convert(s8[51,16]{0,1} tmp_3)\n  tmp_5 = f16[16,16]{1,0} dot(f16[16,51]{1,0} tmp_1, f16[51,16]{0,1} tmp_4), lhs_contracting_dims={1}, rhs_contracting_dims={0}\n  ROOT tmp_6 = f16[1,16,16]{2,1,0} bitcast(f16[16,16]{1,0} tmp_5)\n}"
  result {
    run_time {
      nanos: 31744
    }
    triton {
      block_m: 32
      block_n: 32
      block_k: 32
      split_k: 1
      num_stages: 1
      num_warps: 4
    }
  }
}

O banco de dados de ajuste automático pode ser serializado usando XLA_FLAGS=--xla_gpu_dump_autotune_results_to=<myfile.pbtxt>

[hlo-opt] Desenvolvimento e depuração de transmissão HLO

# If you are working with hardware independent passes from the
# `xla/hlo/transforms/` directory, prefer light-weight version
# of the `hlo-opt` tool with fewer dependencies:

bazel run //xla/hlo/tools:hlo-opt -- [flags] <filename>

# Otherwise, for hardware independent and CPU, GPU passes use
# the same binary from "Compile HLO Modules" section above:

bazel run //xla/tools:hlo-opt -<- [flags>] filename

A ferramenta hlo-opt permite a execução de transmissões individuais independentemente das etapas de compilação da plataforma. Esse isolamento ajuda a executar rapidamente transmissões no módulo HLO de entrada e identificar a causa raiz das falhas.

hlo-opt --passes=schedule-aware-collective-cse input.hlo

A ferramenta hlo-opt também é compatível com DebugOptions XLA_FLAGS.

hlo-opt --passes=schedule-aware-collective-cse
--xla_gpu_experimental_collective_cse_distance_threshold=20 input.hlo

Use a opção --list-passes para receber a string do nome do cartão.

hlo-opt --list-passes

Os usuários podem criar um pipeline personalizado especificando mais de uma transmissão para a opção --passes.

hlo-opt --passes=pass1,pass2,pass3 input.hlo

Ajudar no desenvolvimento de novas transmissões de HLO

  1. Primeiro, escreva seu cartão.
  2. Registre o novo cartão no registro de cartões da ferramenta hlo-opt.

    RegisterPass<FooPass>(FooPassInputOptions)
    

    Com base no tipo de cartão, escolha um dos seguintes locais para registro:
    opt_lib.cc Cartões independentes de hardware.
    cpu_opt.cc Passagens específicas da CPU.
    gpu_opt.cc Passagens específicas da GPU.
    compiled_opt.cc Passagens comuns a CPU, GPU e XPU.
    Não se esqueça de adicionar a dependência de build.

    Inclua o registro do cartão como parte da sua RP(exemplo) para que ele fique disponível para todos os usuários do hlo-opt.

  3. Recompile a ferramenta hlo-opt, valide o registro de cartão bem-sucedido usando a opção --list-passes e use a opção --passes para executar o cartão.

    $ hlo-opt --passes=foo-pass input.hlo
    
  4. Para escrever testes de unidade para a transmissão, consulte https://openxla.org/xla/test_hlo_passes para mais detalhes.

Passagem de medição de tempo de execução

Para modelos grandes, as execuções de compilação completa podem levar até alguns minutos, o que dificulta a detecção de regressões de desempenho sutis. Por outro lado, execuções de transmissão individuais usando hlo-opt permitem uma medição precisa da performance e a detecção fácil de pequenos aumentos no tempo de execução causados por novas mudanças no código.

time hlo-opt --passes=reduce-window-rewriter,scatter_simplifier
--xla_reduce_window_rewrite_base_length=128 input.hlo

[hlo-opt] Converter formatos de módulo HLO

# Use the light weight version of the `hlo-opt` tool.

bazel run //xla/hlo/tools:hlo-opt -- [flags] <filename>

Converter HLO Text -> HLO Proto

hlo-opt --emit-proto input.hlo

Converter HLO Proto ou HLO Proto Binary -> HLO Text

hlo-opt input.pbtxt or input.pb

[ptx-opt] Módulo LLVM do compilador para PTX

A ferramenta vai executar o pipeline de otimização do LLVMIR e chamar o CompileToPtx.

bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 <filename>

A ferramenta também pode despejar LLVMIR após cada caminho.

bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 --xla_dump_to=<path> --xla_gpu_dump_llvmir <filename>