El flujo de trabajo de desarrollo de XLA suele centrarse en el IR de HLO, que representa el cálculo 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 lo ejecutan o proporcionan una etapa de compilación intermedia. El uso de estas herramientas es invaluable para un ciclo de iteración compile->modify->run rápido, ya que HLO se puede visualizar y modificar, y cambiarlo y ejecutarlo 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 sencilla de obtener el 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 HLO previos a la optimización en la carpeta especificada, junto con muchos otros artefactos útiles.
[run_hlo_module] Ejecuta módulos de HLO
bazel run //xla/tools:run_hlo_module -- [flags] <filename>
La herramienta run_hlo_module opera en HLO previo a la optimizació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 su corrección es la siguiente:
run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo
Ejecuta varios módulos de HLO
La invocación con varios módulos de HLO es compatible con run_hlo_module. Para ejecutar todos los módulos de HLO de un directorio, haz lo siguiente:
bazel run //xla/tools:run_hlo_module -- [flags] /dump/*before_optimizations*
[multihost_hlo_runner] Ejecuta módulos de HLO con compatibilidad con SPMD
# Note: Binary name is `hlo_runner_main`.
bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- [flags] <filename>
El ejecutor de HLO para varios hosts es una herramienta muy similar, con la advertencia de que admite SPMD, incluida la comunicación entre hosts. Consulta Multi-Host HLO Runner para obtener más detalles.
Ejecuta varios módulos de HLO con compatibilidad con SPMD
Al igual que run_hlo_module, multihost_hlo_runner también admite la invocación con varios módulos.
bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- [flags] /dump/*before_optimizations*
[hlo-opt] Compila el módulo de HLO
bazel run //xla/tools:hlo-opt -- --platform=[gpu|cpu|...] [more flags] <filename>
Cuando se depura o se intenta comprender 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 de HLO o StableHLO determinada.
hlo-opt admite varias etapas de salida: PTX, HLO después de las optimizaciones, LLVM IR antes de las optimizaciones o TritonIR. El conjunto exacto de etapas admitidas 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
buffer-assignment
hlo
hlo-backend
html
llvm
llvm-after-optimizations
llvm-before-optimizations
ptx
Después de seleccionar una etapa, el usuario puede escribir el resultado de la conversión para una plataforma determinada en un flujo determinado:
hlo-opt --platform=cpu --stage=hlo input.hlo
que imprimiría el volcado en stdout (o en un archivo determinado si se especificó -o).
Compilación sin dispositivo para GPU
La compilación sin dispositivo no necesita acceso a una GPU. La compilación sin dispositivo proporciona una forma de especificar las especificaciones de la GPU en la línea de comandos (--xla_gpu_target_config_filename) para las etapas en las que se requiere acceso a la GPU, lo que elimina la necesidad de un dispositivo de GPU.
Ejemplo: Salida de PTX sin acceso a un dispositivo de GPU:
hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=/xla/tools/hlo_opt/gpu_specs/a100_pcie_80.txtpb input.hlo
Las especificaciones de las GPUs populares se incluyen 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"
Puedes encontrar más especificaciones de la GPU en /xla/tools/hlo_opt/gpu_specs.
Ajuste automático
A veces, la compilación puede implicar el ajuste automático basado en un --stage de compilación.
Para que la compilación sin dispositivo funcione, el usuario debe
inhabilitar el ajuste automático con --xla_gpu_autotune_level=0
o
cargar resultados de ajuste automático preexistentes con
--xla_gpu_load_autotune_results_from=<filename> (obtenido con
--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
El archivo de ajuste automático es una serialización de texto de autotune_results.proto, con un ejemplo como el siguiente:
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_to=<myfile.pbtxt>.
[hlo-opt] HLO Pass Development And Debugging
# 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
La herramienta hlo-opt permite la ejecución de un pase individual, independientemente de las etapas de compilación de la plataforma determinada. Este aislamiento ayuda a ejecutar rápidamente pases en el módulo de HLO de entrada y a identificar la causa raíz de las fallas.
hlo-opt --passes=schedule-aware-collective-cse input.hlo
La herramienta hlo-opt también admite DebugOptions XLA_FLAGS.
hlo-opt --passes=schedule-aware-collective-cse
--xla_gpu_experimental_collective_cse_distance_threshold=20 input.hlo
Usa la opción--list-passes para obtener la cadena del nombre del pase.
hlo-opt --list-passes
Los usuarios pueden crear su propia canalización personalizada especificando más de un pase para la opción --passes.
hlo-opt --passes=pass1,pass2,pass3 input.hlo
Asiste el desarrollo de nuevos pases de HLO
- Primero, escribe tu pase.
Registra el pase nuevo en el registro de pases de la herramienta
hlo-opt.RegisterPass<FooPass>(FooPassInputOptions)Según el tipo de pase, elige una de las siguientes ubicaciones para el registro:
opt_lib.ccPases independientes del hardware.
cpu_opt.ccSon pases específicos de la CPU.
gpu_opt.ccPases específicos de la GPU.
compiled_opt.ccPases comunes a CPU, GPU y XPU.
No olvides agregar la dependencia de compilación.Incluye el registro del pase como parte de tu RP(ejemplo) para que el pase esté disponible para todos los usuarios de
hlo-opt.Vuelve a compilar la herramienta
hlo-opt, valida el registro correcto del pase con la opción--list-passesy, luego, usa la opción--passespara ejecutar el pase.$ hlo-opt --passes=foo-pass input.hlo¿Escribir pruebas de unidades para el pase? Consulta https://openxla.org/xla/test_hlo_passes para obtener más detalles.
Medición del tiempo de ejecución del pase
En el caso de los modelos grandes, las ejecuciones de compilación completas pueden tardar hasta unos minutos, lo que dificulta la detección de regresiones de rendimiento sutiles. En cambio, las ejecuciones de pases individuales con hlo-opt permiten una medición precisa del rendimiento y la detección sencilla incluso de pequeños aumentos en el tiempo de ejecución causados por nuevos cambios en el código.
time hlo-opt --passes=reduce-window-rewriter,scatter_simplifier
--xla_reduce_window_rewrite_base_length=128 input.hlo
[hlo-opt] Convierte formatos de módulos de 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 o HLO Proto Binary -> HLO Text
hlo-opt input.pbtxt or input.pb
[ptx-opt] Módulo del compilador LLVM a PTX
La herramienta ejecutará la canalización de optimización de LLVMIR y, luego, llamará a CompileToPtx.
bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 <filename>
La herramienta también puede volcar LLVMIR después de cada ruta.
bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 --xla_dump_to=<path> --xla_gpu_dump_llvmir <filename>