Рабочий процесс разработки XLA обычно сосредоточен вокруг HLO IR, представляющего собой изолированные функциональные вычисления, передаваемые компилятору. XLA поставляется с несколькими инструментами командной строки (описанными ниже), которые обрабатывают HLO и либо запускают его, либо предоставляют промежуточный этап компиляции. Использование таких инструментов бесценно для быстрого цикла итераций compile->modify->run , поскольку HLO можно визуализировать и модифицировать, а итеративное изменение и запуск часто является самым быстрым способом понять и исправить проблемы с производительностью или поведением XLA.
Простейший способ получить HLO для программы, компилируемой с помощью XLA, обычно заключается в использовании переменной среды XLA_FLAGS :
$ XLA_FLAGS=--xla_dump_to=/tmp/myfolder ./myprogram-entry-point
Этот модуль сохраняет все файлы HLO, созданные до оптимизации, в указанной папке, а также множество других полезных артефактов.
[ run_hlo_module ] Запуск модулей HLO
bazel run //xla/tools:run_hlo_module -- [flags] <filename>
Инструмент run_hlo_module работает с HLO до оптимизации и по умолчанию включает компиляцию, запуск и сравнение с эталонной реализацией интерпретатора. Например, обычный вызов для запуска входного файла computation.hlo на графическом процессоре NVIDIA и проверки его корректности выглядит следующим образом:
run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo
Запуск нескольких модулей HLO
Для run_hlo_module поддерживается запуск нескольких модулей HLO одновременно. Чтобы запустить все модули HLO из каталога:
bazel run //xla/tools:run_hlo_module -- [flags] /dump/*before_optimizations*
[ multihost_hlo_runner ] Запуск модулей HLO с поддержкой SPMD
# Note: Binary name is `hlo_runner_main`.
bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- [flags] <filename>
Multihost HLO Runner — это очень похожий инструмент, с той оговоркой, что он поддерживает SPMD, включая межхостовое взаимодействие. Подробнее см. Multi-Host HLO Runner .
Запуск нескольких модулей HLO с поддержкой SPMD.
Подобно run_hlo_module , multihost_hlo_runner также поддерживает вызов с несколькими модулями.
bazel run //xla/tools/multihost_hlo_runner:hlo_runner_main -- [flags] /dump/*before_optimizations*
[ hlo-opt ] Компилировать модуль HLO
bazel run //xla/tools:hlo-opt -- --platform=[gpu|cpu|...] [more flags] <filename>
При отладке или понимании работы компилятора часто бывает полезно получить информацию о расширении для конкретного оборудования в конкретной точке конвейера (будь то HLO, оптимизированный HLO, TritonIR или LLVM) для заданного входного сигнала HLO или StableHLO.
hlo-opt поддерживает несколько этапов вывода: PTX, HLO после оптимизации, LLVM IR до оптимизации или TritonIR. Точный набор поддерживаемых этапов зависит от платформы (например, PTX является специфичным для NVIDIA) и может быть просмотрен с помощью команды --list-stages:
hlo-opt --platform=CUDA --list-stages
buffer-assignment
hlo
hlo-backend
html
llvm
llvm-after-optimizations
llvm-before-optimizations
ptx
После выбора этапа пользователь может записать результат преобразования для данной платформы в указанный поток:
hlo-opt --platform=cpu --stage=hlo input.hlo
Это выведет дамп в стандартный поток вывода (или в указанный файл, если указан параметр -o ).
Беспроводная компиляция для графического процессора
Для компиляции без использования графического процессора доступ к нему не требуется. Компиляция без использования графического процессора позволяет указать характеристики графического процессора в командной строке ( --xla_gpu_target_config_filename ) на этапах, где требуется доступ к графическому процессору, что исключает необходимость в самом графическом процессоре.
Пример: вывод PTX без доступа к графическому процессору:
hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=/xla/tools/hlo_opt/gpu_specs/a100_pcie_80.txtpb input.hlo
Спецификации популярных графических процессоров поставляются вместе с компилятором, а предоставленный файл представляет собой строковую сериализацию файла 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"
Дополнительные характеристики графического процессора находятся по адресу /xla/tools/hlo_opt/gpu_specs
Автотюнинг
Иногда компиляция может включать автоматическую настройку на основе параметра ` --stage . Для корректной работы компиляции без устройств пользователю необходимо либо...
Отключите автонастройку с помощью --xla_gpu_autotune_level=0
или
Загрузите уже имеющиеся результаты автонастройки с помощью --xla_gpu_load_autotune_results_from=<filename> (полученные с помощью --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
Файл автонастройки представляет собой текстовую сериализацию файла autotune_results.proto , пример которого выглядит следующим образом:
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
}
}
}
Базу данных автонастройки можно сериализовать с помощью XLA_FLAGS=--xla_gpu_dump_autotune_results_to=<myfile.pbtxt>
[ hlo-opt ] Разработка и отладка HLO-проходов
# 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>
Инструмент hlo-opt позволяет выполнять отдельные проходы независимо от этапов компиляции конкретной платформы. Такая изоляция помогает быстро запускать проходы на входном модуле hlo и точно определять первопричину сбоев.
hlo-opt --passes=schedule-aware-collective-cse input.hlo
Инструмент hlo-opt также поддерживает DebugOptions XLA_FLAGS .
hlo-opt --passes=schedule-aware-collective-cse
--xla_gpu_experimental_collective_cse_distance_threshold=20 input.hlo
Используйте опцию --list-passes , чтобы получить строковое имя пароля.
hlo-opt --list-passes
Пользователи могут создать собственный конвейер обработки данных, указав несколько проходов в параметре --passes .
hlo-opt --passes=pass1,pass2,pass3 input.hlo
Помощь в разработке нового пропуска HLO.
- Сначала напишите свой пропуск.
Зарегистрируйте новый пропуск в реестре пропусков инструмента
hlo-opt.RegisterPass<FooPass>(FooPassInputOptions)В зависимости от типа пропуска выберите одно из следующих мест для регистрации:
opt_lib.ccАппаратно-независимые проходы.
cpu_opt.ccспецифические проходы для ЦП.
gpu_opt.ccспецифические проходы для графического процессора.
compiled_opt.ccПередает общие параметры процессору, видеокарте и графическому процессору.
Не забудьте добавить зависимость сборки.Включите регистрацию пропуска в свой запрос на слияние ( пример ), чтобы пропуск был доступен для всех пользователей
hlo-opt.Пересоберите инструмент
hlo-opt, проверьте успешность регистрации прохода с помощью параметра--list-passes, а затем используйте параметр--passesдля запуска прохода.$ hlo-opt --passes=foo-pass input.hloПишете модульные тесты для этого теста? Подробнее см. https://openxla.org/xla/test_hlo_passes .
Пройдено измерение времени выполнения
Для больших моделей полная компиляция может занимать несколько минут, что затрудняет обнаружение незначительных изменений производительности. В отличие от этого, отдельные проходы компиляции с использованием hlo-opt позволяют точно измерять производительность и легко обнаруживать даже небольшое увеличение времени выполнения, вызванное изменениями в коде.
time hlo-opt --passes=reduce-window-rewriter,scatter_simplifier
--xla_reduce_window_rewrite_base_length=128 input.hlo
[ hlo-opt ] Преобразование форматов модулей HLO
# Use the light weight version of the `hlo-opt` tool.
bazel run //xla/hlo/tools:hlo-opt -- [flags] <filename>
Преобразовать HLO Text в HLO Proto
hlo-opt --emit-proto input.hlo
Преобразование HLO Proto или HLO Proto Binary -> HLO Text
hlo-opt input.pbtxt or input.pb
[ ptx-opt ] Модуль LLVM компилятора преобразован в PTX
Инструмент запустит конвейер оптимизации LLVMIR, а затем вызовет CompileToPtx.
bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 <filename>
Этот инструмент также может выводить LLVMIR после каждого пути.
bazel run //xla/hlo/tools/ptx-opt -- --arch=9.0 --xla_dump_to=<path> --xla_gpu_dump_llvmir <filename>
[ isolate_hlo ] Изолировать проблемные инструкции HLO
Если у вас есть большой дамп HLO и вы подозреваете, что причиной сбоя является конкретная инструкция или раздел внутри модуля HLO, вы можете использовать инструмент isolate_hlo .
Этот инструмент извлекает отдельную инструкцию HLO (и необходимый для неё контекст) в новый, меньший по размеру модуль HLO. Это чрезвычайно полезно для создания минимального средства воспроизведения ошибки на уровне компилятора.
- Документация и исходный код: Инструмент
isolate_hloдоступен в репозитории OpenXLA. См. каталогxla/toolsв исходном коде XLA. Использование: Соберите инструмент из исходного кода XLA. Обычно он принимает на вход файл модуля HLO (текстовый или протокольный), имя инструкции для извлечения и путь к выходному файлу.
# Example usage after building XLA: # ./build/tools/isolate_hlo --input=module.hlo --instruction_name=fusion.123 \ # --output=isolated_fusion.123.hlo --input_format=txt --output_format=long_txtДля получения информации о конкретных флагах и параметрах форматирования обратитесь к справке инструмента (
--help).