Il flusso di lavoro di sviluppo XLA in genere è incentrato sull'IR HLO, che rappresenta il calcolo funzionale isolato fornito al compilatore. XLA è dotato di diversi strumenti a riga di comando
(descritti di seguito) che utilizzano l'HLO e lo eseguono oppure forniscono una
fase di compilazione intermedia. L'utilizzo di questi strumenti è fondamentale per un ciclo di iterazionecompile->modify->run
rapido, poiché l'HLO è visualizzabile e hackabile e modificarlo ed eseguirlo in modo iterativo è spesso il modo più rapido per comprendere e correggere il rendimento o il comportamento di XLA.
Il modo più semplice per ottenere l'HLO per un programma compilato con XLA è solitamente utilizzare la variabile di ambiente XLA_FLAGS
:
$ XLA_FLAGS=--xla_dump_to=/tmp/myfolder ./myprogram-entry-point
che memorizza tutti i file HLO precedenti all'ottimizzazione nella cartella specificata, insieme a molti altri artefatti utili.
Snippet HLO in esecuzione: run_hlo_module
Lo strumento run_hlo_module
opera su HLO pre-ottimizzazione e per impostazione predefinita riunisce la compilazione, l'esecuzione e il confronto con l'implementazione dell'interprete di riferimento. Ad esempio, l'invocazione consueta per eseguire un file di inputcomputation.hlo
su una GPU NVIDIA e verificarne la correttezza è:
$ run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo
Come per tutti gli strumenti, puoi utilizzare --help
per ottenere l'elenco completo delle opzioni.
Esecuzione di snippet HLO con il supporto SPMD: multihost_hlo_runner
Il programma di esecuzione HLO multihost è uno strumento molto simile, con l'avvertenza che supporta la tecnologia SPMD, inclusa la comunicazione tra host. Per maggiori dettagli, consulta Esegui HLO su più host.
Riproduzione di più HLO
L'invocazione con più moduli è supportata sia per run_hlo_module
sia per
hlo_runner_main
, il che è spesso utile per riprodurre tutti i moduli in una directory
dump:
$ hlo_runner_main /dump/*before_optimizations*
Passaggi/fasi di compilazione HLO in esecuzione: hlo-opt
Durante il debug o la comprensione del funzionamento del compilatore, è spesso utile ottenere l'espansione per un determinato hardware in un determinato punto della pipeline (che si tratti di HLO, HLO ottimizzato, TritonIR o LLVM), per un determinato input HLO (stabile).
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
hlo
llvm
ptx
Dopo aver selezionato una fase, l'utente può scrivere il risultato della conversione per una determinata piattaforma in un determinato stream:
$ hlo-opt myinput.hlo --platform=CUDA --stage=llvm
che stampa il dump su stdout (o su un determinato file se è stato specificato -o
).
Utilizzo senza dispositivo
L'accesso a una GPU non è necessario per la maggior parte della compilazione e, specificando una spec GPU sulla riga di comando, possiamo ottenere ad esempio l'output PTX senza accedere a un acceleratore:
$ hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=(pwd)/tools/data/gpu_specs/a100_pcie_80.txtpb input.hlo
Le specifiche delle 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"
La compilazione senza dispositivo potrebbe riscontrare problemi se è richiesta l'ottimizzazione automatica. Fortunatamente, possiamo anche fornirli nella riga di 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
Il file di ottimizzazione automatica è la serializzazione del testo di autotune_results.proto
, con un example simile al seguente:
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_t=<myfile.pbtxt>
Esecuzione di un singolo passaggio del compilatore
Sono supportati anche i flag di XLA_FLAGS
, quindi lo strumento può essere utilizzato per testare l'esecuzione di un singolo passaggio:
$ hlo-opt --platform=CUDA --stage=hlo --passes=algebraic_simplifer input.hlo