XLA Инструменты

Рабочий процесс разработки 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.

  1. Сначала напишите свой пропуск.
  2. Зарегистрируйте новый пропуск в реестре пропусков инструмента hlo-opt .

    RegisterPass<FooPass>(FooPassInputOptions)
    

    В зависимости от типа пропуска выберите одно из следующих мест для регистрации:
    opt_lib.cc Аппаратно-независимые проходы.
    cpu_opt.cc специфические проходы для ЦП.
    gpu_opt.cc специфические проходы для графического процессора.
    compiled_opt.cc Передает общие параметры процессору, видеокарте и графическому процессору.
    Не забудьте добавить зависимость сборки.

    Включите регистрацию пропуска в свой запрос на слияние ( пример ), чтобы пропуск был доступен для всех пользователей hlo-opt .

  3. Пересоберите инструмент hlo-opt , проверьте успешность регистрации прохода с помощью параметра --list-passes , а затем используйте параметр --passes для запуска прохода.

    $ hlo-opt --passes=foo-pass input.hlo
    
  4. Пишете модульные тесты для этого теста? Подробнее см. 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 ).