O fluxo de trabalho de desenvolvimento do XLA geralmente é centrado no
IR do HLO, que representa a computação funcional isolada
fornecida ao compilador. A XLA vem com várias ferramentas de linha de comando
(descritas abaixo) que consomem e executam a HLO ou fornecem um
estágio de compilação intermediário. O uso dessas ferramentas é muito importante para um ciclo de iteração
compile->modify->run rápido, já que a HLO pode ser visualizada e
hackeada, e a mudança e a execução iterativas geralmente são a maneira mais rápida de
entender e corrigir o desempenho ou o comportamento de uma XLA.
A maneira mais fácil de acessar o HLO de um programa compilado com XLA geralmente é
usando 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.
Snippets de HLO em execução: run_hlo_module
A ferramenta run_hlo_module opera em HLO de pré-otimização e, por padrão,
reúne a compilação de pacotes, a execução e a comparação com a implementação
do interpretador de referência. Por exemplo, a invocação usual 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
Como em todas as ferramentas, --help pode ser usado para acessar a lista completa de opções.
Como executar snippets de HLO com suporte ao SPMD: multihost_hlo_runner
O HLO runner multihost é uma ferramenta muito semelhante, com a ressalva de que ele oferece suporte ao SPMD, incluindo a comunicação entre hosts. Consulte Multi-Host HLO Runner (link em inglês) para mais detalhes.
Repetição de HLO múltipla
A invocação com vários módulos tem suporte para run_hlo_module e
hlo_runner_main, o que geralmente é conveniente para reproduzir todos os módulos em um diretório
de despejo:
$ hlo_runner_main /dump/*before_optimizations*
Execução de passagens/etapas da compilação de HLO: hlo-opt
Ao depurar ou entender o funcionamento do compilador, muitas vezes é útil extrair 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 (estável).
O hlo-opt oferece suporte a vários estágios de saída: PTX, HLO após otimizações,
LLVM IR antes das otimizações ou TritonIR. O conjunto exato de estágios com suporte
depende da plataforma (por exemplo, o PTX é específico da NVIDIA) e pode ser visto usando
o comando --list-stages:
$ hlo-opt --platform=CUDA --list-stages
hlo
llvm
ptx
Depois de selecionar uma fase, o usuário pode gravar o resultado da conversão de uma plataforma específica para um determinado stream:
$ hlo-opt myinput.hlo --platform=CUDA --stage=llvm
que imprime o despejo no stdout (ou em um determinado arquivo se -o for especificado).
Uso sem dispositivo
O acesso a uma GPU não é necessário para a maior parte da compilação. Ao especificar uma especificação de GPU na linha de comando, podemos receber, por exemplo, a saída PTX sem acesso a um acelerador:
$ hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=(pwd)/tools/data/gpu_specs/a100_pcie_80.txtpb input.hlo
As especificações de GPUs conhecidas 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"
A compilação sem dispositivo pode ter problemas se a autoajuste for necessário. Felizmente, também podemos fornecer esses dados na linha de comando:
$ 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
o exemplo parecido com este:
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 autoajuste pode ser serializado usando
XLA_FLAGS=--xla_gpu_dump_autotune_results_t=<myfile.pbtxt>
Como executar uma única passagem do compilador
As flags de XLA_FLAGS também são compatíveis, portanto, a ferramenta pode ser usada para testar
a execução de uma única passagem:
$ hlo-opt --platform=CUDA --stage=hlo --passes=algebraic_simplifer input.hlo