Utilizzare gli strumenti XLA

Il flusso di lavoro per lo sviluppo XLA è solitamente incentrato HLO IR, che rappresenta una funzione isolata calcolo fornito al compilatore. XLA include vari strumenti a riga di comando (descritti di seguito) che utilizzano l'HLO e lo eseguono oppure di compilazione intermedia. L'uso di questi strumenti è inestimabile per una rapida compile->modify->run, poiché l'HLO è sia visualizzabile sia è soggetto a pirateria informatica e la modifica e la gestione iterativa è spesso il modo più rapido per comprendere e correggere un rendimento o un comportamento XLA.

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

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

che archivia tutti i file HLO precedenti all'ottimizzazione nella cartella specificata, insieme con molti altri artefatti utili.

Snippet HLO in esecuzione: run_hlo_module

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

$ 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 di snippet HLO con supporto SPMD: multihost_hlo_runner

Il runner HLO Multihost è uno strumento molto simile, con l'avvertenza che supporta SPMD, inclusa la comunicazione tra host. Consulta Esecutore HLO multi-host per maggiori dettagli.

Ripetizione multi-HLO

La chiamata con più moduli è supportata sia per run_hlo_module che hlo_runner_main, che spesso è conveniente riprodurre tutti i moduli in un dump directory:

$ hlo_runner_main /dump/*before_optimizations*

Pass correnti/fasi della compilazione HLO: hlo-opt

Spesso è utile eseguire il debug o comprendere il funzionamento del compilatore per ottenere l'espansione per un particolare hardware in un determinato punto della pipeline (HLO, HLO ottimizzato, TritonIR o LLVM), per un determinato HLO (stabile) di testo.

hlo-opt supporta più fasi di output: PTX, HLO dopo le ottimizzazioni, LLVM IR prima delle ottimizzazioni, o TritonIR. L'esatto insieme 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
hlo
llvm
ptx

Dopo aver selezionato una fase, l'utente può scrivere il risultato della conversione per una determinata piattaforma a uno stream specifico:

$ hlo-opt myinput.hlo --platform=CUDA --stage=llvm

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

Utilizzo senza dispositivo

L'accesso a una GPU non è necessario per la maggior parte della compilazione e specificando un Specifiche GPU nella riga di comando, ad esempio Uscita PTX senza accesso 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 per le GPU più utilizzate vengono fornite insieme al compilatore e il file fornito è serializzazione della stringa 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 causare problemi se è necessaria l'ottimizzazione automatica. Per fortuna possiamo specificarli 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 testuale di autotune_results.proto, con un esempio 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>

Eseguire un pass per compilatore singolo

Sono supportati anche i flag di XLA_FLAGS, quindi lo strumento può essere usato per testare eseguendo un solo passaggio:

$ hlo-opt --platform=CUDA --stage=hlo --xla-hlo-enable-passes-only=algebraic_simplifer input.hlo