Outils XLA

Le workflow de développement XLA est généralement centré sur l'HLO, 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 tels outils est inestimable pour un cycle d'itération compile->modify->run rapide, 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 les modules HLO

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

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

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

Exécuter plusieurs modules HLO

L'invocation avec plusieurs modules HLO est acceptée pour run_hlo_module. Pour exécuter tous les modules HLO d'un répertoire :

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

[multihost_hlo_runner] Exécuter des modules HLO avec la prise en charge 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 prend en charge SPMD, y compris la communication entre hôtes. Pour en savoir plus, consultez Multi-Host HLO Runner.

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

Comme run_hlo_module, multihost_hlo_runner est également compatible avec l'invocation 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>

Lorsque vous déboguez ou essayez de comprendre le fonctionnement du compilateur, il est souvent utile d'obtenir l'expansion pour un matériel spécifique à 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 optimisations, LLVM IR avant optimisations ou TritonIR. L'ensemble exact des étapes prises en charge 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 affichera le dump sur la sortie standard (ou dans un fichier donné si -o a été spécifié).

Compilation sans appareil pour GPU

La compilation sans appareil n'a pas besoin d'accéder à un GPU. La compilation sans appareil permet de spécifier les caractéristiques 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 de sortie PTX sans accès à un 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 populaires 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"

Pour en savoir plus sur les spécifications des GPU, consultez /xla/tools/hlo_opt/gpu_specs.

Réglage automatique

La compilation peut parfois impliquer un réglage automatique basé sur une compilation --stage. Pour que la compilation sans appareil fonctionne, l'utilisateur doit
désactiver l'optimisation automatique avec --xla_gpu_autotune_level=0
ou
charger des résultats d'optimisation 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. Voici un exemple :

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 d'optimisation 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 HLO Pass

# 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 d'exécuter des passes individuelles indépendamment des étapes de compilation de la plate-forme donnée. Cette isolation permet d'exécuter rapidement des passes 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 du nom de la carte.

hlo-opt --list-passes

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

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

Aider à développer un nouveau pass HLO

  1. Commencez par rédiger votre pass.

  2. Enregistrez le nouveau pass dans le registre des pass d'outils hlo-opt.

    RegisterPass<FooPass>(FooPassInputOptions)

    En fonction du type de carte, choisissez l'un des emplacements suivants pour l'enregistrement :
    opt_lib.cc Cartes indépendantes du matériel.
    cpu_opt.cc Passes spécifiques au processeur.
    gpu_opt.cc Passes spécifiques au GPU.
    compiled_opt.cc Passes communs au CPU, au GPU et à l'XPU.
    N'oubliez pas d'ajouter une dépendance de compilation.

    Incluez l'enregistrement du pass dans votre PR(exemple) afin que le pass soit disponible pour tous les utilisateurs hlo-opt.

  3. 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

  4. Pour savoir comment écrire des tests unitaires pour les passes, consultez https://openxla.org/xla/test_hlo_passes.

Mesure de l'exécution des passes

Pour les grands modèles, les compilations complètes peuvent prendre plusieurs minutes, ce qui rend difficile la détection des régressions de performances subtiles. En revanche, les exécutions de passes individuelles utilisant hlo-opt permettent de mesurer précisément les performances et de détecter facilement même les petites augmentations du temps d'exécution causées par les nouvelles modifications du 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 -> HLO Proto

hlo-opt --emit-proto input.hlo

Convertir HLO Proto ou HLO Proto Binary → HLO Text

hlo-opt input.pbtxt or input.pb

[ptx-opt] Compiler le module LLVM en PTX

L'outil exécutera le pipeline d'optimisation LLVMIR, puis appellera CompileToPtx.

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

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

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