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
- Primeiro, escreva seu cartão.
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.ccCartões independentes de hardware.
cpu_opt.ccPassagens específicas da CPU.
gpu_opt.ccPassagens específicas da GPU.
compiled_opt.ccPassagens 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.Recompile a ferramenta
hlo-opt, valide o registro de cartão bem-sucedido usando a opção--list-passese use a opção--passespara executar o cartão.$ hlo-opt --passes=foo-pass input.hloPara 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>