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

Вызов с несколькими модулями HLO поддерживается для run_hlo_module . Чтобы запустить все модули 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 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

  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
,

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

Вызов с несколькими модулями HLO поддерживается для run_hlo_module . Чтобы запустить все модули 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 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

  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
,

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

Вызов с несколькими модулями HLO поддерживается для run_hlo_module . Чтобы запустить все модули 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 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

  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
,

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

Вызов с несколькими модулями HLO поддерживается для run_hlo_module . Чтобы запустить все модули 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 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

  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