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 e lo eseguono o forniscono una
fase di compilazione intermedia. L'utilizzo di questi strumenti è prezioso per un ciclo di iterazione rapido, poiché HLO è visualizzabile e modificabile e la sua modifica ed esecuzione iterativa è spesso il modo più rapido per comprendere e correggere le prestazioni o il comportamento di XLA.compile->modify->run
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 pre-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-ottimizzato 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 controllarne la correttezza è:
run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo
Esegui più moduli HLO
L'invocazione 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>
Multihost HLO runner è uno strumento molto simile, con l'avvertenza che supporta SPMD, inclusa la comunicazione tra host. Per i dettagli, consulta Multi-Host HLO Runner.
Esegui più moduli HLO con il supporto SPMD
Analogamente a run_hlo_module, multihost_hlo_runner supporta anche l'invocazione
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 esegui il debug o cerchi di capire il funzionamento del compilatore, spesso è utile ottenere l'espansione per un particolare hardware in un determinato punto della pipeline (che si tratti di 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. Il set esatto di fasi supportate
dipende dalla piattaforma (ad esempio PTX è specifico di 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
fornisce un modo per specificare le specifiche della GPU nella riga di comando
(--xla_gpu_target_config_filename) per le fasi in cui è necessario 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 delle GPU più comuni 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"
Ulteriori specifiche della GPU sono disponibili all'indirizzo /xla/tools/hlo_opt/gpu_specs
Ottimizzazione automatica
A volte la compilazione può comportare la sintonizzazione automatica in base a una compilation --stage.
Affinché la compilazione senza dispositivo funzioni, l'utente deve
disattivare la sintonizzazione automatica con --xla_gpu_autotune_level=0
o
caricare risultati di sintonizzazione 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 del 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] HLO Pass Development And Debugging
# 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 un singolo passaggio
indipendentemente dalle fasi di compilazione della piattaforma specificata. Questo isolamento aiuta a
eseguire rapidamente 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 della tessera.
hlo-opt --list-passes
Gli utenti possono creare una pipeline personalizzata specificando più di un passaggio
all'opzione --passes.
hlo-opt --passes=pass1,pass2,pass3 input.hlo
Assist New HLO Pass Development
- Innanzitutto, scrivi il tuo pass.
Registra la nuova tessera nel registro delle tessere dello strumento
hlo-opt.RegisterPass<FooPass>(FooPassInputOptions)In base al tipo di tessera, scegli una delle seguenti posizioni per la registrazione:
opt_lib.ccTessere indipendenti dall'hardware.
cpu_opt.ccPass specifici per la CPU.
gpu_opt.ccPass specifici per la GPU.
compiled_opt.ccPassaggi comuni a CPU, GPU e XPU.
Non dimenticare di aggiungere la dipendenza dalla build.Includi la registrazione della tessera nella tua PR(esempio) in modo che la tessera sia disponibile per l'utilizzo per tutti gli utenti di
hlo-opt.Ricostruisci lo strumento
hlo-opt, convalida la registrazione della tessera utilizzando l'opzione--list-passese poi utilizza l'opzione--passesper eseguire la tessera.$ hlo-opt --passes=foo-pass input.hloPer scrivere test delle unità per la pass, consulta la pagina https://openxla.org/xla/test_hlo_passes per ulteriori dettagli.
Misurazione del runtime della tessera
Per i modelli di grandi dimensioni, le esecuzioni di compilazione complete possono richiedere fino a qualche minuto, il che rende
difficile rilevare regressioni delle prestazioni impercettibili. Al contrario, le esecuzioni
individuali dei pass utilizzando hlo-opt consentono una misurazione precisa
delle prestazioni e il rilevamento semplice anche di piccoli aumenti del
tempo di esecuzione causati da nuove modifiche al codice.
time hlo-opt --passes=reduce-window-rewriter,scatter_simplifier
--xla_reduce_window_rewrite_base_length=128 input.hlo
[hlo-opt] Convertire 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] Modulo LLVM del compilatore fino a 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>