Рабочий процесс разработки 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 до оптимизации в указанной папке вместе со многими другими полезными артефактами.
Запуск фрагментов HLO: run_hlo_module
Инструмент run_hlo_module работает с HLO предварительной оптимизации и по умолчанию объединяет компиляцию, запуск и сравнение с реализацией эталонного интерпретатора. Например, обычный вызов для запуска входного файла computation.hlo на графическом процессоре NVIDIA и проверки его правильности:
$ run_hlo_module --platform=CUDA --reference_platform=Interpreter computation.hlo
Как и все другие инструменты, --help можно использовать для получения полного списка опций.
Запуск фрагментов HLO с поддержкой SPMD: multihost_hlo_runner
Multihost HLO runner — очень похожий инструмент, с той лишь оговоркой, что он поддерживает SPMD, включая межхостовую связь. Подробности см. в разделе Multi-Host HLO Runner .
Повтор мульти-HLO
Вызов с несколькими модулями поддерживается как для run_hlo_module , так и hlo_runner_main , что часто удобно для воспроизведения всех модулей в каталоге дампа:
$ hlo_runner_main /dump/*before_optimizations*
Запуск проходов/этапов компиляции HLO: hlo-opt
При отладке или понимании работы компилятора часто бывает полезно получить расширение для конкретного оборудования в определенной точке конвейера (будь то HLO, оптимизированный HLO, TritonIR или LLVM) для заданного (стабильного) входа HLO. .
hlo-opt поддерживает несколько этапов вывода: PTX, HLO после оптимизации, LLVM IR до оптимизации или TritonIR. Точный набор поддерживаемых этапов зависит от платформы (например, PTX зависит от NVIDIA) и его можно увидеть с помощью команды --list-stages:
$ hlo-opt --platform=CUDA --list-stages
hlo
llvm
ptx
После выбора этапа пользователь может записать результат конвертации для данной платформы в заданный поток:
$ hlo-opt myinput.hlo --platform=CUDA --stage=llvm
который будет печатать дамп на стандартный вывод (или в заданный файл, если был указан -o ).
Использование без устройств
Доступ к графическому процессору не требуется для большей части компиляции, и, указав спецификацию графического процессора в командной строке, мы можем получить, например, вывод PTX без доступа к ускорителю:
$ hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=(pwd)/tools/data/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"
При компиляции без устройств могут возникнуть проблемы, если требуется автонастройка. К счастью, мы также можем предоставить их в командной строке:
$ 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_t=<myfile.pbtxt>
Запуск одного прохода компилятора
Флаги из XLA_FLAGS также поддерживаются, поэтому этот инструмент можно использовать для тестирования выполнения одного прохода:
$ hlo-opt --platform=CUDA --stage=hlo --xla-hlo-enable-passes- input.hlo