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 ).

Компиляция без устройств для GPU

Компиляция без устройств не требует доступа к графическому процессору. Компиляция без устройств позволяет указать спецификацию графического процессора в командной строке ( --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 Pass 

# 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 Pass

  1. Сначала напишите свой пропуск.

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

    RegisterPass<FooPass>(FooPassInputOptions)

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

    Включите регистрацию пропуска в свой PR ( пример ), чтобы пропуск был доступен для использования всеми пользователями 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>