Por lo general, el flujo de trabajo de desarrollo de XLA se centra en el IR de HLO, que representa el procesamiento funcional aislado que se le proporciona al compilador. XLA incluye varias herramientas de línea de comandos (que se describen a continuación) que consumen HLO y la ejecutan o proporcionan una etapa de compilación intermedia. El uso de estas herramientas es inestimable para un ciclo de iteración rápido de compile->modify->run, ya que la HLO se puede visualizar y hackear, y cambiarla e implementarla de forma iterativa suele ser la forma más rápida de comprender y corregir el rendimiento o el comportamiento de XLA.
La forma más fácil de obtener la HLO para un programa que se compila con XLA suele ser usar la variable de entorno XLA_FLAGS:
$ XLA_FLAGS=--xla_dump_to=/tmp/myfolder ./myprogram-entry-point
que almacena todos los archivos de HLO anteriores a la optimización en la carpeta especificada, junto con muchos otros artefactos útiles.
Fragmentos de HLO en ejecución: run_hlo_module
La herramienta run_hlo_module opera en HLO de preoptimización y, de forma predeterminada, agrupa la compilación, la ejecución y la comparación con la implementación del intérprete de referencia. Por ejemplo, la invocación habitual para ejecutar un archivo de entrada computation.hlo en una GPU de NVIDIA y verificar que sea correcta es la siguiente:
$ run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo
Al igual que con todas las herramientas, --help se puede usar para obtener la lista completa de opciones.
Ejecuta fragmentos de HLO con compatibilidad con SPMD: multihost_hlo_runner
El ejecutor de HLO multihost es una herramienta muy similar, con la salvedad de que admite SPMD, incluida la comunicación entre hosts. Consulta Multi-Host HLO Runner para obtener más detalles.
Repetición de varios HLO
La invocación con varios módulos es compatible con run_hlo_module y hlo_runner_main, lo que suele ser conveniente para volver a reproducir todos los módulos en un directorio de volcado:
$ hlo_runner_main /dump/*before_optimizations*
Ejecución de pases o etapas de compilación de HLO: hlo-opt
Cuando se depura o se comprende el funcionamiento del compilador, a menudo es útil obtener la expansión para un hardware en particular en un punto determinado de la canalización (ya sea HLO, HLO optimizado, TritonIR o LLVM), para una entrada de HLO (estable) determinada.
hlo-opt admite varias etapas de salida, ya sea PTX, HLO después de las optimizaciones, IR de LLVM antes de las optimizaciones o TritonIR. El conjunto exacto de etapas compatibles depende de la plataforma (por ejemplo, PTX es específico de NVIDIA) y se puede ver con el comando --list-stages:
$ hlo-opt --platform=CUDA --list-stages
hlo
llvm
ptx
Después de seleccionar una etapa, el usuario puede escribir el resultado de la conversión de una plataforma determinada en un flujo determinado:
$ hlo-opt myinput.hlo --platform=CUDA --stage=llvm
que imprimiría el volcado en stdout (o en un archivo determinado si se especificó -o).
Uso sin dispositivos
No se necesita acceso a una GPU para la mayor parte de la compilación. Si especificas una especificación de GPU en la línea de comandos, podemos obtener, por ejemplo, una salida de PTX sin acceso a un acelerador:
$ hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=(pwd)/tools/data/gpu_specs/a100_pcie_80.txtpb input.hlo
Las especificaciones de las GPUs populares se envían con el compilador, y el archivo proporcionado es la serialización de cadenas 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 compilación sin dispositivos puede tener problemas si se requiere el ajuste automático. Por suerte, también podemos proporcionarlos en la línea de comandos:
$ 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
El archivo de ajuste automático es la serialización de texto de autotune_results.proto, con un ejemplo que se ve de la siguiente manera:
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 datos de ajuste automático se puede serializar con XLA_FLAGS=--xla_gpu_dump_autotune_results_t=<myfile.pbtxt>.
Cómo ejecutar un solo pase del compilador
También se admiten las marcas de XLA_FLAGS, por lo que la herramienta se puede usar para probar la ejecución de un solo pase:
$ hlo-opt --platform=CUDA --stage=hlo --passes=algebraic_simplifer input.hlo