Le workflow de développement XLA est généralement centré sur l'IR HLO, qui représente le calcul fonctionnel isolé transmis au compilateur. XLA est fourni avec plusieurs outils de ligne de commande (décrits ci-dessous) qui utilisent le HLO et l'exécutent ou fournissent une étape de compilation intermédiaire. L'utilisation de ces outils est inestimable pour un cycle d'itération compile->modify->run rapide, car le HLO est à la fois visualisable et hackable, et 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'une XLA.
Le moyen le plus simple d'obtenir le HLO pour un programme en cours de compilation 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 d'optimisation antérieure dans le dossier spécifié, ainsi que de nombreux autres artefacts utiles.
Extraits HLO exécutés: run_hlo_module
L'outil run_hlo_module fonctionne sur l'opération HLO pré-optimisation 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'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
Comme avec tous les outils, vous pouvez utiliser --help pour obtenir la liste complète des options.
Exécuter des extraits HLO compatibles avec SPMD: multihost_hlo_runner
L'exécuteur d'HLO à hôtes multiples est un outil très similaire, avec une mise en garde : il est compatible avec le protocole SPMD, y compris la communication entre hôtes. En général, un appel se présente comme suit:
hlo_runner_main --device_type=gpu --use_spmd_partitioning=true --num_partitions=4 --num_replicas=1 --hlo_file=computation.hlo
Exécutions/étapes de la compilation HLO: hlo-opt
Lors du débogage ou de la compréhension du fonctionnement du compilateur, il est souvent utile d'obtenir l'extension pour un matériel particulier à un point précis du pipeline (HLO, HLO optimisé, TritonIR ou LLVM) pour une entrée HLO (stable) donnée.
hlo-opt est compatible avec plusieurs étapes de sortie: PTX, HLO après les optimisations, IR LLVM avant optimisations ou TritonIR. L'ensemble exact d'étapes compatibles dépend de la plate-forme (par exemple, PTX est spécifique à NVIDIA). Vous pouvez le vérifier à l'aide de la commande --list-stages:
$ hlo-opt --platform=CUDA --list-stages
hlo
llvm
ptx
Après avoir sélectionné une étape, l'utilisateur peut écrire le résultat de la conversion d'une plate-forme donnée dans un flux donné:
$ hlo-opt myinput.hlo --platform=CUDA --stage=llvm
qui imprimerait le vidage sur stdout (ou dans un fichier donné si -o était spécifié).
Utilisation sans appareil
L'accès à un GPU n'est pas nécessaire pour la majeure partie de la compilation. En spécifiant une spécification de GPU sur la ligne de commande, nous pouvons, par exemple, obtenir une sortie PTX sans accès à un accélérateur:
$ hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=(pwd)/tools/data/gpu_specs/a100_80.txtpb input.hlo
Les spécifications des GPU les plus courants sont fournies avec le compilateur. Le fichier fourni est la 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"
La compilation sans appareil peut rencontrer des problèmes si un réglage automatique est requis. Heureusement, nous pouvons également les fournir sur la ligne de commande:
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
Le fichier autotune est une sérialisation textuelle de autotune_results.proto. Voici un exemple:
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
}
}
}
La base de données d'ajustement automatique peut être sérialisée à l'aide de XLA_FLAGS=--xla_gpu_dump_autotune_results_t=<myfile.pbtxt>.
Exécuter une carte de compilateur unique
Les indicateurs de XLA_FLAGS sont également acceptés. L'outil peut donc être utilisé pour tester l'exécution d'une seule carte:
hlo-opt --platform=CUDA --stage=hlo --xla-hlo-enable-passes-only=algebraic_simplifer input.hlo