El flujo de trabajo de desarrollo de XLA suele centrarse en la IR HLO, que representa el procesamiento funcional aislado que se le proporciona al compilador. XLA incluye varias herramientas de línea de comandos (descritas a continuación) que consumen HLO y lo ejecutan, o proporcionan una etapa de compilación intermedia. El uso de esas herramientas es invaluable para un ciclo de iteración compile->modify->run rápido, ya que el HLO se puede visualizar y se puede hackear, y cambiar y ejecutarlo de manera iterativa suele ser la forma más rápida de entender y corregir el rendimiento o el comportamiento de XLA.
Por lo general, la forma más fácil de obtener el HLO para un programa que se compila con XLA es usar la variable de entorno XLA_FLAGS:
XLA_FLAGS=--xla_dump_to=/tmp/myfolder ./myprogram-entry-point
que almacena todos los archivos HLO previos a la optimización en la carpeta especificada, junto con muchos otros artefactos útiles.
Ejecuta fragmentos de HLO: run_hlo_module
La herramienta run_hlo_module funciona en HLO de optimización previa y, de forma predeterminada, compila, ejecuta y compara con la implementación del intérprete de referencia de forma predeterminada. Por ejemplo, la invocación habitual para ejecutar un archivo de entrada computation.hlo en una GPU NVIDIA y comprobar su precisión 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.
Ejecución de fragmentos de HLO con compatibilidad con SPMD: multihost_hlo_runner
El ejecutor de HLO de varios hosts es una herramienta muy similar, pero con la salvedad de que es compatible con SPMD, incluida la comunicación entre hosts. Una invocación típica se ve de la siguiente manera:
hlo_runner_main --device_type=gpu --use_spmd_partitioning=true --num_partitions=4 --num_replicas=1 --hlo_file=computation.hlo
Ejecución de pases o etapas de compilación de HLO: hlo-opt
Cuando depuras o comprendes el funcionamiento del compilador, suele ser útil obtener la expansión para un hardware en particular en un punto específico de la canalización (ya sea HLO, HLO optimizado, TritonIR o LLVM) para una entrada HLO (estable) determinada.
hlo-opt admite varias etapas de salida: PTX, HLO después de las optimizaciones, IR de LLVM antes de las optimizaciones o TritonIR. El conjunto exacto de etapas admitidas depende de la plataforma (p.ej., 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 una transmisión determinada:
$ 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
El acceso a una GPU no es necesario para la mayor parte de la compilación y, si se especifica una especificación de GPU en la línea de comandos, podemos obtener, p.ej., resultados de PTX sin acceso a un acelerador:
$ hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=(pwd)/tools/data/gpu_specs/a100_80.txtpb input.hlo
Las especificaciones para las GPU 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. Afortunadamente, también podemos proporcionarlos en la línea de comandos:
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
El archivo de ajuste automático es la serialización de texto de autotune_results.proto, cuyo ejemplo se ve de la siguiente manera:
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 datos de ajuste automático se puede serializar con XLA_FLAGS=--xla_gpu_dump_autotune_results_t=<myfile.pbtxt>.
Cómo ejecutar un pase de compilador único
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 --xla-hlo-enable-passes-only=algebraic_simplifer input.hlo