Strumenti XLA

Il flusso di lavoro di sviluppo di XLA è in genere incentrato su HLO IR, che rappresenta il calcolo funzionale isolato fornito al compilatore. XLA include più strumenti a riga di comando (descritti di seguito) che utilizzano HLO ed eseguono o forniscono una fase di compilazione intermedia. L'utilizzo di questi strumenti è prezioso per un ciclo di iterazione rapido compile->modify->run, poiché HLO è sia visualizzabile sia modificabile e la modifica e l'esecuzione iterativa sono spesso il modo più rapido per comprendere e correggere un comportamento o un rendimento di XLA.

Il modo più semplice per ottenere l'HLO per un programma compilato con XLA è in genere utilizzare la variabile di ambiente XLA_FLAGS:

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

che memorizza tutti i file HLO prima dell'ottimizzazione nella cartella specificata, insieme a molti altri artefatti utili.

[`run_hlo_module`] Esegui moduli HLO

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

Lo strumento run_hlo_module opera su HLO pre-ottimizzazione e, per impostazione predefinita, raggruppa la compilazione, l'esecuzione e il confronto con l'implementazione dell'interprete di riferimento. Ad esempio, la chiamata abituale per eseguire un file di input computation.hlo su una GPU NVIDIA e verificarne la correttezza è:

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

Esegui più moduli HLO

La chiamata con più moduli HLO è supportata per run_hlo_module. Per eseguire tutti i moduli HLO da una directory:

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

[`multihost_hlo_runner`] Esegui moduli HLO con supporto SPMD

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

Il runner HLO multi-host è uno strumento molto simile, con l'avvertenza che supporta SPMD, inclusa la comunicazione tra host. Per maggiori dettagli, consulta Runner HLO multi-host.

Esegui più moduli HLO con supporto SPMD

Analogamente a run_hlo_module, multihost_hlo_runner supporta anche la chiamata con più moduli.

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

[`hlo-opt`] Compila modulo HLO

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

Quando si esegue il debug o si comprende il funzionamento del compilatore, è spesso utile ottenere l'espansione per un hardware specifico in un punto specifico della pipeline (HLO, HLO ottimizzato, TritonIR o LLVM), per un determinato input HLO o StableHLO.

hlo-opt supporta più fasi di output: PTX, HLO dopo le ottimizzazioni, LLVM IR prima delle ottimizzazioni o TritonIR. L'insieme esatto di fasi supportate dipende dalla piattaforma (ad esempio, PTX è specifico per NVIDIA) e può essere visualizzato utilizzando il comando --list-stages:

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

Dopo aver selezionato una fase, l'utente può scrivere il risultato della conversione per una determinata piattaforma in un determinato stream:

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

che stamperebbe il dump su stdout (o su un determinato file se è stato specificato -o).

Compilazione senza dispositivo per GPU

La compilazione senza dispositivo non richiede l'accesso a una GPU. La compilazione senza dispositivo consente di specificare le specifiche della GPU nella riga di comando (--xla_gpu_target_config_filename) per le fasi in cui è richiesto l'accesso alla GPU, eliminando la necessità di un dispositivo GPU.

Esempio: output PTX senza accesso a un dispositivo GPU:

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

Le specifiche per le GPU più diffuse vengono fornite con il compilatore e il file fornito è la serializzazione di stringhe di 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"

Altre specifiche della GPU si trovano in /xla/tools/hlo_opt/gpu_specs

Ottimizzazione automatica

A volte la compilazione può comportare l'ottimizzazione automatica in base a una --stage di compilazione. Affinché la compilazione senza dispositivo funzioni, l'utente deve
disattivare l'ottimizzazione automatica con --xla_gpu_autotune_level=0
oppure
caricare i risultati di ottimizzazione automatica preesistenti con --xla_gpu_load_autotune_results_from=<filename> (ottenuti con --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

Il file di ottimizzazione automatica è la serializzazione di testo di autotune_results.proto, con un esempio simile a:

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
    }
  }
}

Il database di ottimizzazione automatica può essere serializzato utilizzando XLA_FLAGS=--xla_gpu_dump_autotune_results_to=<myfile.pbtxt>

[`hlo-opt`] Sviluppo e debug dei passaggi 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>

Lo strumento hlo-opt consente l'esecuzione di singoli passaggi indipendentemente dalle fasi di compilazione della piattaforma specificata. Questo isolamento consente di eseguire rapidamente i passaggi sul modulo HLO di input e individuare la causa principale degli errori.

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

Lo strumento hlo-opt supporta anche DebugOptions XLA_FLAGS.

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

Utilizza l'opzione --list-passes per ottenere la stringa del nome del passaggio.

hlo-opt --list-passes

Gli utenti possono creare una pipeline personalizzata specificando più passaggi all'opzione --passes.

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

Assisti il nuovo sviluppo di passaggi HLO

Innanzitutto, scrivi il passaggio.
Registra il nuovo passaggio nel registro dei passaggi dello strumento hlo-opt.
```
RegisterPass<FooPass>(FooPassInputOptions)
```
In base al tipo di passaggio, scegli una delle seguenti posizioni per la registrazione:
opt_lib.cc Passaggi indipendenti dall'hardware.
cpu_opt.cc Passaggi specifici per la CPU.
gpu_opt.cc Passaggi specifici per la GPU.
compiled_opt.cc Passaggi comuni a CPU, GPU, XPU.
Non dimenticare di aggiungere la dipendenza dalla build.

Includi la registrazione del passaggio come parte della richiesta di pull(esempio) in modo che il passaggio sia disponibile per tutti gli utenti hlo-opt.
Ricompila lo strumento hlo-opt, convalida la registrazione del passaggio riuscita utilizzando l'opzione --list-passes e poi utilizza l'opzione --passes per eseguire il passaggio.
```
$ hlo-opt --passes=foo-pass input.hlo
```
Stai scrivendo test delle unità per il passaggio? Per maggiori dettagli, consulta https://openxla.org/xla/test_hlo_passes.

Misurazione del runtime dei passaggi

Per i modelli di grandi dimensioni, le esecuzioni di compilazione complete possono richiedere fino a qualche minuto, il che rende difficile rilevare regressioni sottili del rendimento. Al contrario, le esecuzioni di singoli passaggi utilizzando hlo-opt consentono una misurazione precisa del rendimento e il rilevamento semplice anche di piccoli aumenti del tempo di esecuzione causati da nuove modifiche del codice.

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

[`hlo-opt`] Converti i formati dei moduli HLO

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

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

Converti `HLO Text` -> `HLO Proto`

hlo-opt --emit-proto input.hlo

Converti `HLO Proto` o `HLO Proto Binary` -> `HLO Text`

hlo-opt input.pbtxt or input.pb

[`ptx-opt`] Compila il modulo LLVM in PTX

Lo strumento eseguirà la pipeline di ottimizzazione LLVMIR e poi chiamerà CompileToPtx.

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

Lo strumento può anche eseguire il dump di LLVMIR dopo ogni percorso.

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

[`isolate_hlo`] Istruzioni HLO problematiche isolate

Se hai un dump HLO di grandi dimensioni e sospetti che un'istruzione o una sezione specifica all'interno del modulo HLO stia causando l'arresto anomalo, puoi utilizzare lo strumento isolate_hlo.

Questo strumento estrae una singola istruzione HLO (e il relativo contesto necessario) in un nuovo modulo HLO più piccolo. Questo è estremamente utile per creare un riproduttore minimo a livello di compilatore.

Documentazione e origine: lo strumento isolate_hlo è disponibile nel repository OpenXLA. Consulta la directory xla/tools nel codice sorgente XLA.
Utilizzo: crea lo strumento dall'albero di origine XLA. In genere accetta un file di modulo HLO di input (testo o proto), il nome dell'istruzione da estrarre e il percorso del file di output.
```
# Example usage after building XLA:
# ./build/tools/isolate_hlo --input=module.hlo --instruction_name=fusion.123 \
#   --output=isolated_fusion.123.hlo --input_format=txt --output_format=long_txt
```
Consulta il messaggio di aiuto dello strumento (--help) per opzioni di formato e flag specifici.

Strumenti XLA Mantieni tutto organizzato con le raccolte Salva e classifica i contenuti in base alle tue preferenze.

[run_hlo_module] Esegui moduli HLO

Esegui più moduli HLO

[multihost_hlo_runner] Esegui moduli HLO con supporto SPMD

Esegui più moduli HLO con supporto SPMD

[hlo-opt] Compila modulo HLO