Использование инструментов 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 до оптимизации в указанной папке вместе со многими другими полезными артефактами.

Запуск фрагментов 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, включая межхостовую связь. Типичный вызов выглядит так:

hlo_runner_main  --device_type=gpu --use_spmd_partitioning=true --num_partitions=4 --num_replicas=1 --hlo_file=computation.hlo

Запуск проходов/этапов компиляции 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_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_80.txtpb --xla_gpu_load_autotune_results_from=results.textpb input.hlo

Файл автонастройки представляет собой текстовую сериализацию autotune_results.proto , пример которого выглядит так:

version: 2
results {
  device: "sm_8.0 with 42331013120B RAM, 108 cores, 1410000KHz clock, 1215000KHz mem clock, 41943040B L2$"
  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-only=algebraic_simplifer input.hlo