Outils XLA

Le workflow de développement XLA est généralement centré autour HLO IR, qui représente un calcul fonctionnel isolé donné au compilateur. XLA est fourni avec plusieurs outils de ligne de commande (décrits ci-dessous) qui consomment HLO et l'exécutent ou fournissent une étape de compilation intermédiaire. L'utilisation de ces outils est essentielle pour un cycle d'itération rapide compile->modify->run, car HLO est à la fois visualisable et piratable. Le modifier et l'exécuter de manière itérative est souvent le moyen le plus rapide de comprendre et de corriger les performances ou le comportement d'XLA.

Le moyen le plus simple d'obtenir le HLO pour un programme compilé avec XLA consiste généralement à utiliser la variable d'environnement XLA_FLAGS :

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

qui stocke tous les fichiers HLO avant optimisation dans le dossier spécifié, ainsi que de nombreux autres artefacts utiles.

[`run_hlo_module`] Exécuter des modules HLO

bazel run //xla/tools:run_hlo_module -- [flags] <filename>

L'outil run_hlo_module fonctionne sur le HLO avant optimisation et regroupe par défaut la compilation, l'exécution et la comparaison avec l'implémentation de l'interpréteur de référence. Par exemple, l'appel habituel pour exécuter un fichier d'entrée computation.hlo sur un GPU NVIDIA et vérifier son exactitude est le suivant :

run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo

Exécuter plusieurs modules HLO

L'appel avec plusieurs modules HLO est compatible avec run_hlo_module. Pour exécuter tous les modules hlo à partir d'un répertoire :

bazel run //xla/tools:run_hlo_module -- [flags] /dump/*before_optimizations*

[`multihost_hlo_runner`] Exécuter des modules HLO avec prise en charge de SPMD

# Note: Binary name is `hlo_runner_main`.
bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- [flags] <filename>

Le runner HLO multihôte est un outil très similaire, à la différence qu'il est compatible avec SPMD, y compris la communication entre hôtes. Pour en savoir plus, consultez la section Runner HLO multihôte.

Exécuter plusieurs modules HLO avec prise en charge de SPMD

Comme run_hlo_module, multihost_hlo_runner est également compatible avec l'appel avec plusieurs modules.

bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- [flags] /dump/*before_optimizations*

[`hlo-opt`] Compiler le module HLO

bazel run //xla/tools:hlo-opt -- --platform=[gpu|cpu|...] [more flags] <filename>

Lors du débogage ou de la compréhension du fonctionnement du compilateur, il est souvent utile d'obtenir l'expansion d'un matériel particulier à un moment donné du pipeline (qu'il s'agisse de HLO, de HLO optimisé, de TritonIR ou de LLVM), pour une entrée HLO ou StableHLO donnée.

hlo-opt est compatible avec plusieurs étapes de sortie : PTX, HLO après optimisation, LLVM IR avant optimisation ou TritonIR. L'ensemble exact des étapes compatibles dépend de la plate-forme (par exemple, PTX est spécifique à NVIDIA) et peut être consulté à l'aide de la commande --list-stages :

hlo-opt --platform=CUDA --list-stages
buffer-assignment
hlo
hlo-backend
html
llvm
llvm-after-optimizations
llvm-before-optimizations
ptx

Après avoir sélectionné une étape, l'utilisateur peut écrire le résultat de la conversion pour une plate-forme donnée dans un flux donné :

hlo-opt --platform=cpu --stage=hlo input.hlo

qui affiche le vidage sur stdout (ou dans un fichier donné si -o a été spécifié).

Compilation sans appareil pour GPU

La compilation sans appareil ne nécessite pas d'accès à un GPU. Elle permet de spécifier la spécification du GPU sur la ligne de commande (--xla_gpu_target_config_filename) pour les étapes où l'accès au GPU est requis, ce qui élimine le besoin d'un appareil GPU.

Exemple : sortie PTX sans accès à un appareil GPU :

hlo-opt  --platform=CUDA --stage=llvm  --xla_gpu_target_config_filename=/xla/tools/hlo_opt/gpu_specs/a100_pcie_80.txtpb input.hlo

Les spécifications des GPU courants sont fournies avec le compilateur, et le fichier fourni est une sérialisation de chaîne de 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"

D'autres spécifications de GPU se trouvent dans /xla/tools/hlo_opt/gpu_specs.

Réglage automatique

La compilation peut parfois impliquer un réglage automatique basé sur une --stage de compilation. Pour que la compilation sans appareil fonctionne, l'utilisateur doit
désactiver le réglage automatique avec --xla_gpu_autotune_level=0
ou
charger des résultats de réglage automatique préexistants avec --xla_gpu_load_autotune_results_from=<filename> (obtenus avec --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

Le fichier de réglage automatique est une sérialisation de texte de autotune_results.proto, avec un exemple comme celui-ci :

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
    }
  }
}

La base de données de réglage automatique peut être sérialisée à l'aide de XLA_FLAGS=--xla_gpu_dump_autotune_results_to=<myfile.pbtxt>.

[`hlo-opt`] Développement et débogage de pass HLO

# 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>

L'outil hlo-opt permet l'exécution de pass individuels indépendamment des étapes de compilation de la plate-forme donnée. Cette isolation permet d'exécuter rapidement des pass sur le module hlo d'entrée et d'identifier la cause première des échecs.

hlo-opt --passes=schedule-aware-collective-cse input.hlo

L'outil hlo-opt est également compatible avec DebugOptions XLA_FLAGS.

hlo-opt --passes=schedule-aware-collective-cse
--xla_gpu_experimental_collective_cse_distance_threshold=20 input.hlo

Utilisez l'option --list-passes pour obtenir la chaîne de nom de pass.

hlo-opt --list-passes

Les utilisateurs peuvent créer leur propre pipeline personnalisé en spécifiant plusieurs pass à l'option --passes.

hlo-opt --passes=pass1,pass2,pass3 input.hlo

Aider au développement de nouveaux pass HLO

Commencez par écrire votre pass.
Enregistrez le nouveau pass dans le registre de pass de l'outil hlo-opt.
```
RegisterPass<FooPass>(FooPassInputOptions)
```
En fonction du type de pass, choisissez l'un des emplacements suivants pour l'enregistrement :
opt_lib.cc Pass indépendants du matériel.
cpu_opt.cc Pass spécifiques au processeur.
gpu_opt.cc Pass spécifiques au GPU.
compiled_opt.cc Pass communs au processeur, au GPU et à l'XPU.
N'oubliez pas d'ajouter une dépendance de compilation.

Incluez l'enregistrement du pass dans votre demande d'extraction(exemple) afin que le pass soit disponible pour tous les utilisateurs hlo-opt.
Recompilez l'outil hlo-opt, validez l'enregistrement réussi du pass à l'aide de l'option --list-passes, puis utilisez l'option --passes pour exécuter le pass.
```
$ hlo-opt --passes=foo-pass input.hlo
```
Vous écrivez des tests unitaires pour le pass ? Consultez https://openxla.org/xla/test_hlo_passes pour en savoir plus.

Mesure de l'exécution du pass

Pour les grands modèles, les exécutions de compilation complètes peuvent prendre jusqu'à quelques minutes, ce qui rend difficile la détection de régressions de performances subtiles. En revanche, les exécutions de pass individuels à l'aide de hlo-opt permettent une mesure précise des performances et la détection facile, même de petites augmentations du temps d'exécution causées par de nouvelles modifications de code.

time hlo-opt --passes=reduce-window-rewriter,scatter_simplifier
--xla_reduce_window_rewrite_base_length=128 input.hlo

[`hlo-opt`] Convertir les formats de module HLO

# Use the light weight version of the `hlo-opt` tool.

bazel run //xla/hlo/tools:hlo-opt -- [flags] <filename>

Convertir `HLO Text` en `HLO Proto`

hlo-opt --emit-proto input.hlo

Convertir `HLO Proto` ou `HLO Proto Binary` en `HLO Text`

hlo-opt input.pbtxt or input.pb

[`ptx-opt`] Compiler le module LLVM en PTX

L'outil exécute le pipeline d'optimisation LLVMIR, puis appelle CompileToPtx.

bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 <filename>

L'outil peut également vider LLVMIR après chaque chemin.

bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 --xla_dump_to=<path> --xla_gpu_dump_llvmir <filename>

[`isolate_hlo`] Isoler les instructions HLO problématiques

Si vous disposez d'un vidage HLO volumineux et que vous pensez qu'une instruction ou une section spécifique du module HLO est à l'origine du plantage, vous pouvez utiliser l'outil isolate_hlo.

Cet outil extrait une seule instruction HLO (et son contexte nécessaire) dans un nouveau module HLO plus petit. Cela est extrêmement utile pour créer un reproducteur minimal au niveau du compilateur.

Documentation et source : l'outil isolate_hlo est disponible dans le dépôt OpenXLA. Consultez le répertoire xla/tools dans le code source XLA.
Utilisation : compilez l'outil à partir de l'arborescence source XLA. Il prend généralement un fichier de module HLO d'entrée (texte ou proto), le nom de l'instruction à extraire et le chemin d'accès au fichier de sortie.
```
# Example usage after building XLA:
# ./build/tools/isolate_hlo --input=module.hlo --instruction_name=fusion.123 \
#   --output=isolated_fusion.123.hlo --input_format=txt --output_format=long_txt
```
Pour connaître les options et les formats spécifiques, consultez le message d'aide de l'outil (--help).

Outils XLA Restez organisé à l'aide des collections Enregistrez et classez les contenus selon vos préférences.

[run_hlo_module] Exécuter des modules HLO

Exécuter plusieurs modules HLO

[multihost_hlo_runner] Exécuter des modules HLO avec prise en charge de SPMD

Exécuter plusieurs modules HLO avec prise en charge de SPMD

[hlo-opt] Compiler le module HLO