Utilizzare gli strumenti XLA

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