Il flusso di lavoro di sviluppo XLA è generalmente incentrato sull'IR di HLO, che rappresenta il calcolo funzionale isolato fornito al compilatore. XLA è dotato di diversi strumenti a riga di comando (descritti di seguito) che utilizzano HLO e lo eseguono oppure forniscono una fase di compilazione intermedia. L'utilizzo di questi strumenti è inestimabile per un ciclo di iterazione compile->modify->run veloce, dal momento che HLO può essere visualizzato e compromesso, e modifica ed eseguirlo in modo iterativo è spesso il modo più rapido per comprendere e correggere un comportamento o prestazioni XLA.
Il modo più semplice per ottenere l'HLO per un programma compilato con XLA è
di solito utilizzare la variabile di ambiente XLA_FLAGS:
XLA_FLAGS=--xla_dump_to=/tmp/myfolder ./myprogram-entry-point
che archivia tutti i file HLO prima dell'ottimizzazione nella cartella specificata, insieme a molti altri artefatti utili.
Esecuzione degli snippet HLO in esecuzione: run_hlo_module
Lo strumento run_hlo_module opera sull'HLO precedente all'ottimizzazione e, per impostazione predefinita, raggruppa la compilazione, l'esecuzione e il confronto con l'implementazione dell'interprete di riferimento. Ad esempio, la solita chiamata per eseguire un file di input computation.hlo su una GPU NVIDIA e verificare che sia corretta è:
run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo
Come con tutti gli strumenti, è possibile utilizzare --help per ottenere l'elenco completo delle opzioni.
Esecuzione degli snippet HLO con supporto SPMD: multihost_hlo_runner
L'esecuzione HLO multihost è uno strumento molto simile, ma è necessario chiarire che supporta SPMD, inclusa la comunicazione tra host. Una chiamata tipica ha il seguente aspetto:
hlo_runner_main --device_type=gpu --use_spmd_partitioning=true --num_partitions=4 --num_replicas=1 --hlo_file=computation.hlo
Pass/fasi in esecuzione della compilazione HLO: hlo-opt
Durante il debug o la comprensione del funzionamento del compilatore, spesso è utile ottenere l'espansione di un particolare hardware in un determinato punto della pipeline (ad esempio HLO, HLO ottimizzato, TritonIR o LLVM), per un determinato input HLO (stabile).
hlo-opt supporta più fasi di output: PTX, HLO dopo le ottimizzazioni, IR VM prima delle ottimizzazioni o TritonIR. L'insieme esatto di fasi supportate dipende dalla piattaforma (ad es. 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 flusso di dati:
$ hlo-opt myinput.hlo --platform=CUDA --stage=llvm
in modo da stampare il dump su stdout (o su un determinato file se -o è stato specificato).
Utilizzo senza dispositivo
L'accesso a una GPU non è necessario per la maggior parte della compilazione e specificando una specifica GPU nella riga di comando è possibile ottenere, ad esempio, output PTX senza accesso a un acceleratore:
$ hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=(pwd)/tools/data/gpu_specs/a100_80.txtpb input.hlo
Le specifiche per le GPU più utilizzate 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 fornire anche quelli nella riga di comando:
hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=gpu_specs/a100_80.txtpb --xla_gpu_load_autotune_results_from=results.textpb input.hlo
Il file di ottimizzazione automatica è la serializzazione testuale di autotune_results.proto, con un esempio simile al seguente:
version: 2
results {
device: "sm_8.0 with 42331013120B RAM, 108 cores, 1410000KHz clock, 1215000KHz mem clock, 41943040B L2$"
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 dell'ottimizzazione automatica può essere serializzato utilizzando XLA_FLAGS=--xla_gpu_dump_autotune_results_t=<myfile.pbtxt>
Esecuzione di un pass per compilatore singolo
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 --xla-hlo-enable-passes-only=algebraic_simplifer input.hlo