Le workflow de développement XLA est généralement centré sur l'IR HLO, qui représente le calcul fonctionnel isolé fourni au compilateur. XLA est fourni avec plusieurs outils de ligne de commande (décrits ci-dessous) qui consomment du 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 le HLO est à la fois visualisable et piratable. Modifier et exécuter de manière itérée 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 d'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.
Extraits HLO en cours d'exécution: run_hlo_module
L'outil run_hlo_module fonctionne sur le HLO pré-optimisé et regroupe par défaut la compilation, l'exécution et la comparaison avec l'implémentation de l'interprète 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 pour tous les outils, --help permet d'obtenir la liste complète des options.
Exécuter des extraits HLO avec la prise en charge de SPMD: multihost_hlo_runner
Le programme d'exécution HLO multi-hôte est un outil très similaire, à l'exception qu'il est compatible avec SPMD, y compris la communication entre les hôtes. Pour en savoir plus, consultez la section Multi-Host HLO Runner.
Replay multi-HLO
L'appel avec plusieurs modules est compatible avec run_hlo_module et hlo_runner_main, ce qui est souvent pratique pour rejouer tous les modules dans un répertoire de vidage:
$ hlo_runner_main /dump/*before_optimizations*
Exécution des passes/étapes de compilation HLO: hlo-opt
Lors du débogage ou de la compréhension du fonctionnement du compilateur, il est souvent utile d'obtenir l'expansion pour un matériel particulier à un point particulier du pipeline (qu'il s'agisse de HLO, de HLO optimisé, de TritonIR ou de LLVM), pour une entrée HLO (stable) 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 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
hlo
llvm
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 myinput.hlo --platform=CUDA --stage=llvm
qui imprime le dump dans stdout (ou dans un fichier donné si -o a été spécifié).
Utilisation sans appareil
L'accès à un GPU n'est pas nécessaire pour la plupart de la compilation. En spécifiant une spécification de GPU sur la ligne de commande, nous pouvons obtenir, par exemple, 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_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"
La compilation sans appareil peut rencontrer des problèmes si un réglage automatique est nécessaire. 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_pcie_80.txtpb --xla_gpu_load_autotune_results_from=results.textpb input.hlo
Le fichier de réglage automatique est la sérialisation textuelle de autotune_results.proto, comme illustré ci-dessous:
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'ajustement automatique peut être sérialisée à l'aide de XLA_FLAGS=--xla_gpu_dump_autotune_results_t=<myfile.pbtxt>.
Exécuter une seule étape de compilation
Les options de XLA_FLAGS sont également acceptées. L'outil peut donc être utilisé pour tester l'exécution d'une seule passe:
$ hlo-opt --platform=CUDA --stage=hlo --passes=algebraic_simplifer input.hlo