Como usar ferramentas do XLA

O fluxo de trabalho de desenvolvimento do XLA costuma ser centrado em IR HLO, que representa informações funcionais isoladas de computação fornecida ao compilador. O XLA vem com várias ferramentas de linha de comando (descrito abaixo) que consomem HLO e o executam ou fornecem na fase de compilação intermediária. O uso dessas ferramentas é inestimável para uma rápida compile->modify->run, já que o HLO pode ser visualizado e e mudá-lo e executá-lo de forma iterativa costuma ser a maneira mais rápida de entender e corrigir o desempenho ou comportamento de um XLA.

A maneira mais fácil de obter o HLO de um programa compilado com XLA é geralmente usa 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, junto com com muitos outros artefatos úteis.

Executando snippets HLO: run_hlo_module

A ferramenta run_hlo_module opera em um HLO de pré-otimização e, por padrão, compilação de pacotes, execução e comparação com o intérprete de referência implementação. Por exemplo, a invocação normal para executar um arquivo de entrada computation.hlo em uma GPU NVIDIA e para verificar a correção é:

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

Como acontece com todas as ferramentas, --help pode ser usado para acessar a lista completa de opções.

Execução de snippets HLO compatíveis com SPMD: multihost_hlo_runner

O runner HLO multihost é uma ferramenta muito parecida, com a ressalva de que oferece suporte SPMD, incluindo comunicação entre hosts. Consulte Confira detalhes em Executor HLO de vários hosts.

Repetição de multi-HLO

A invocação com vários módulos tem suporte para run_hlo_module e hlo_runner_main, que geralmente é conveniente para reproduzir todos os módulos em um despejo diretório:

$ hlo_runner_main /dump/*before_optimizations*

Passagens/estágios em execução da compilação HLO: hlo-opt

Ao depurar ou entender o funcionamento do compilador, muitas vezes isso é útil para conseguir a expansão de um hardware específico em um ponto específico do pipeline (HLO, HLO otimizado, TritonIR ou LLVM) para um determinado HLO (estável) entrada.

hlo-opt oferece suporte a vários estágios de saída: PTX, HLO após otimizações, IR do LLVM antes das otimizações ou TritonIR. O conjunto exato de estágios compatíveis depende da plataforma (por exemplo, o PTX é específico da NVIDIA) e pode ser visto usando use o comando --list-stages:

$ hlo-opt --platform=CUDA --list-stages
hlo
llvm
ptx

Após selecionar uma etapa, o usuário pode escrever o resultado da conversão para uma plataforma para um determinado fluxo:

$ hlo-opt myinput.hlo --platform=CUDA --stage=llvm

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

Uso sem dispositivo

O acesso a uma GPU não é necessário para a maior parte da compilação e, ao especificar uma a especificação da GPU na linha de comando. Por exemplo, saída PTX sem acesso a uma 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 populares são enviadas com o compilador, e o arquivo fornecido é serialização de strings 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 poderá ter problemas se o ajuste automático for necessário. Felizmente, também podemos fornecê-las 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 exemplo, que será 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_t=<myfile.pbtxt>

Como executar um passe de compilador único

As flags de XLA_FLAGS também são compatíveis. Portanto, a ferramenta pode ser usada para testar executando um único cartão:

$ hlo-opt --platform=CUDA --stage=hlo --xla-hlo-enable-passes-only=algebraic_simplifer input.hlo