В этом документе описывается путь модуля высокоуровневого оптимизатора (HLO) XLA от начального состояния до конечного исполняемого файла. Иногда мы будем опускать слово «модуль» и называть его просто «HLO».
Предварительная оптимизация HLO
Начнём с модуля предварительной оптимизации HLO. Предварительно оптимизированный HLO не содержит операций ( ops ), которые считаются внутренними для XLA, таких как fusion или bitcast . На этом этапе у операций нет макета, а если он есть, он будет проигнорирован. Предварительно оптимизированный HLO обычно создаётся фреймворками более высокого уровня, такими как TensorFlow и JAX. При использовании флага XLA -xla_dump_to предварительная оптимизация HLO сохраняется в файле с суффиксом «before_optimizations.txt».
Оптимизировать модуль HLO
Конвейер XLA:GPU преобразует предоптимизированный HLO в оптимизированный HLO, выполняя последовательность проходов. Проходы можно семантически сгруппировать и выполнять в следующем порядке:
Шардинг, связанный с проходами
Сюда входят такие проходы, как Shardy Partitioner или проходы для шардинга SPMD.
Оптимизация проходит
Это может включать как пропуска по легализации, так и пропуска по упрощению.
Коллективные оптимизационные проходы
Аналогично проходам оптимизации , но фокусируется на коллективных операциях.
Пропуска по макету
Каждой операции HLO назначается макет, являющийся частью формы инструкции. Макет определяет физическое расположение тензора в памяти.
Пример формы с макетом:
f32[10,20,30]{2,0,1}
После типа элемента следуют логические размеры фигуры, за которыми следует перестановка макета в порядке от меньшего к большему. В этом примере самый малый размер — 30, второй по величине малый размер — 10, а самый большой размер — 20.
Целью назначения макета является минимизация количества требуемых физических транспозиций с помощью жадной стратегии. Она начинается с определённых ограничений макета (например, библиотеки cuDNN/cuBLAS ожидают последовательных измерений) и распространяется «вниз», а затем «вверх» по графу HLO. В конце процесса распространения макета некоторые инструкции могут иметь конфликтующие макеты: один из них распространяется от операнда, другой — от пользователя. Для разрешения этого конфликта вставляется инструкция copy HLO, которая меняет макет с макета операнда на макет инструкции.
Проходы нормализации макета
Учитывая, что довольно сложно определить физическую форму, нормализация макета пытается переписать форму так, чтобы она использовала макет по умолчанию {rank-1, rank-2, …, 0} . В приведенном выше примере нормализованная форма будет f32[20,10,30]{2,1,0} . Операции копирования, которые изменяют макеты, переписываются как комбинация transpose и bitcast . Учитывая, что в настоящее время мы не можем нормализовать все операции, все еще есть некоторые операции, которые могут иметь нестандартные макеты, в первую очередь gather и dot . На границах между нормализованными операциями и ненормализованными операциями будут операции bitcast , которые представляют транспозицию, т. е. транспозицию с назначенной компоновкой, которая делает ее физически пустой операцией.
Нормализация макета также делает некоторые неявные транспозиции явными, что важно, поскольку кодогенерация может обрабатывать явные транспозиции с помощью выделенного генератора. Например, при перестройке технически допускается различная физическая компоновка операнда и результата (например, из-за разного ранга). Проход ReshapeDecomposer , выполняемый в рамках проходов нормализации макета, превращает перестройку в последовательность transpose , reshape bitcast и transpose .
Проходы оптимизации после назначения макета
Наиболее важными проходами здесь являются слияния Triton (слияния GEMM + слияния Softmax/Layernorm) или переписывание библиотечных вызовов. На этом этапе также выполняется автонастройка, где XLA выбирает между различными эмиттерами, выбирает оптимальный алгоритм для свёрток или точек, находит наилучшее тайловое покрытие для слияний, обрабатываемых эмиттером Triton, и т. д.
Fusion проходит
Двумя основными проходами являются PriorityFusion и Multi-Output fusion.
В PriorityFusion мы формируем слияния, руководствуясь моделью стоимости. При слиянии мы допускаем дублирование операций с несколькими пользователями, если операция может быть слита со всеми пользователями. Мы также допускаем расширение существующих слияний Triton Softmax, если это возможно.
Объединение Multi-Output — это отдельный проход, позволяющий объединять операции/объединения, использующие один и тот же операнд. Оно также позволяет объединять операнды/объединения операндов в пользователей без дублирования, добавляя дополнительные выходы, чтобы другие пользователи объединяемой операции могли быть перенаправлены на эти выходы. Этот проход следует выполнять осторожно, чтобы не создавать циклов в графе HLO.
После слияния нескольких выходов запускается процедура исключения общих подвыражений (проход HloCSE ), потенциально объединяющая ранее дублированные операции, если они попали в одно слияние.
Несколько проходов после слияния
Несколько проходов связаны с коллективами (например, превращение их в асинхронные или обеспечение определенного относительного порядка коллективов).
Наконец, мы запускаем CopyInsertion где копии добавляются для того, чтобы гарантировать, что операции на месте не перезапишут данные, которые все еще нужны в другом месте.
По завершении оптимизации оптимизированный HLO выгружается в файл с суффиксом «after_optimizations.txt» с использованием флага -xla_dump_to . Если вы хотите выгрузить HLO после промежуточных проходов, которые фактически изменяют HloModule, можно использовать флаг -xla_dump_hlo_pass_re=.* (или специальное регулярное выражение, чтобы ограничить выгрузку определёнными проходами).
Планирование
Модуль HLO без расписания всё ещё имеет некоторую степень свободы в порядке обработки операций. Любая топологическая сортировка, учитывающая связи операндов и результатов, а также зависимости управления, допустима. Планирование определяет конкретный порядок выполнения. На этом этапе основное внимание уделяется максимальному потреблению памяти, которое зависит от времени жизни тензоров. На начальном этапе мы пробуем различные алгоритмы планировщика и выбираем расписание, которое должно минимизировать пиковое потребление памяти. Обратите внимание, что на этом этапе мы пока не работаем с физическими буферами (это произойдёт в разделе «Назначение буферов») и моделируем использование памяти.
Затем запускается проход LatencyHidingScheduler , который пытается максимально увеличить перекрытие вычислений и коммуникаций. Но это может снова увеличить потребление памяти.
Наконец, если пиковое потребление памяти превышает объём доступной памяти, мы запускаем HloRematerialization . Этот проход пытается сократить использование памяти за счёт производительности, например, некоторые слияния могут быть разделены, а некоторые операции могут быть дублированы для сокращения времени жизни буфера. Если происходит рематериализация, может быть полезно изучить способы снижения требований к памяти на стороне модели (например, использование пакетов меньшего размера).
Назначение буфера
Непосредственно перед переходом в LLVM IR мы выполняем проходы назначения буфера, которые назначают буферные срезы каждой инструкции в графе HLO. Назначение буфера выполняется в несколько этапов:
HloDataflowAnalysisприсваивает инструкциямHloValues(по сути, логические буферы). Для операций, выполняемых на месте,HloValueоперанда может быть использовано повторно. Операция может определять более одногоHloValue(например, с формой результата в виде кортежа).HloAliasAnalysisпытается объединить буферы для операций совмещения имен и вычисляет сопоставлениеHloValueсHloBuffer.BufferAssignmentвычисляет сопоставление буферовHloBuffersс буферными срезами внутри большого буфера таким образом, чтобы один и тот же буферный срез не использовался для разных буферовHloBuffersс перекрывающимся временем жизни. Для операций, которые могут быть псевдонимами, допустимо небольшое перекрытие (время окончания одногоHloBufferможет совпадать с временем начала другогоHloBuffer). При использовании флага-xla_dump_toнекоторая информация о назначении буфера выгружается в файл с суффиксом "after_optimizations-buffer-assignment.txt".
Танкс
После оптимизации и планирования графа HLO он преобразуется в линейную последовательность преобразователей для конкретного бэкэнда (ЦП или ГП).
В XLA Thunk — это абстракция самодостаточной единицы работы, выполняемой средой выполнения. Это может быть запуск скомпилированного ядра, определённая операция, вызов библиотеки, конструкция управления потоком, коллективное взаимодействие и т. д. Последовательность Thunk представляет собой весь исполняемый файл для конкретного бэкенда.
Эмиссия звука
Процесс преобразования запланированного вычисления HLO в последовательность thunk называется «выбросом thunk». За него отвечает специальный класс-излучатель в каждом бэкенде.
Для бэкэнда GPU этим занимается IrEmitterUnnested . EmitHloComputation перебирает запланированный список инструкций HLO в вычислении и передает управление специализированному методу Emit... (например, EmitFusion , EmitConvolutionThunk , EmitWhile ). Каждый из этих методов создаёт соответствующий объект(ы) Thunk и добавляет их в последовательность Thunk.
Для бэкенда ЦП эту роль выполняет ThunkEmitter , организованный аналогичным образом. Финальный ThunkSequence встроен в CpuExecutable .
Обратите внимание, что каждая инструкция в вычислении входа модуля HLO может соответствовать ни одному ( kTuple , kConstant , ..), одному или нескольким (например, инструкция сортировки) переходам в конечной последовательности переходов.
Буферы команд: оптимизация выполнения на графическом процессоре
Современное аппаратное обеспечение графических процессоров позволяет записывать последовательность операций графического процессора (запуск ядра, копирование памяти и т. д.) один раз, а затем воспроизводить её многократно с минимальной нагрузкой на процессор. Это критически важная оптимизация производительности, особенно для рабочих нагрузок с большим количеством небольших быстро запускающихся ядер. XLA использует буфер команд как абстракцию графов CUDA или графов HIP. Интерфейс ядра определён в GpuCommandBuffer .
Буфер команд представлен в последовательности переходов CommandBufferThunk .
Эмиттер не создаёт этот thunk напрямую из инструкций HLO. Вместо этого это делает CommandBufferConversionPass , работающий непосредственно над ThunkSequence.
Проход идентифицирует смежные подпоследовательности совместимых thunk-ов (например, последовательность KernelThunk и GemmThunk ). Затем он заменяет найденную подпоследовательность одним CommandBufferThunk . Новый thunk инкапсулирует логику исходных thunk-ов в виде списка легковесных объектов CommandBufferCmd. Когда CommandBufferThunk выполняется впервые в данном потоке GPU, он «записывает» свою последовательность команд в аппаратный буфер команд. При всех последующих выполнениях он просто выдаёт одну команду GPU для «воспроизведения» записанной последовательности. Это позволяет избежать накладных расходов на процессор, связанных с запуском каждого отдельного ядра.
Исполняемый файл
Конечным продуктом конвейера компиляции XLA является самостоятельный, платформенно-зависимый исполняемый объект Executable . Этот объект инкапсулирует всю информацию, необходимую для запуска скомпилированной программы на целевом устройстве, таком как центральный процессор или графический процессор. Он служит мостом между компилятором и средой выполнения. Современные среды выполнения, такие как PJRT, используют несколько более высокоуровневые абстракции (см. PjRtExecutable ), но в конечном итоге они представляют собой обёртку для исполняемого файла, специфичного для бэкенда.
Executable содержит несколько ключевых фрагментов информации, генерируемых во время компиляции. Хотя точное содержимое различается в зависимости от бэкенда, обычно оно включает в себя:
Скомпилированный код: это низкоуровневый машинный код, который будет выполняться на устройстве. Для центральных процессоров это обычно один или несколько объектных файлов. Для графических процессоров это скомпилированный код устройства в формате PTX или HSACO, который загружается в графический процессор во время выполнения.
План выполнения (ThunkSequence): ядро логики среды выполнения. Это линейная последовательность объектов Thunk. Каждый Thunk представляет собой единицу работы, например, запуск ядра, вызов библиотечной функции (например, cuBLAS) или управление потоком управления. Среда выполнения выполняет программу, итеративно проходя по этой последовательности.
Структура памяти (BufferAssignment): этот критически важный фрагмент метаданных, создаваемый BufferAssigner, описывает полную структуру памяти для вычислений. Он определяет размер каждого буфера и порядок выделения и повторного использования памяти для параметров, выходных данных и временных значений. Среда выполнения использует эти данные для выделения памяти устройства и передачи корректных указателей каждому преобразователю.
(необязательно) Модуль HLO: для отладки и профилирования исполняемый файл часто сохраняет ссылку на окончательный оптимизированный модуль HloModule, из которого он был скомпилирован.
Созданием финального исполняемого файла управляет компилятор для каждого конкретного бэкенда. Метод RunBackend реализации компилятора является заключительным этапом процесса компиляции, который упаковывает все скомпилированные артефакты в объект Executable. GpuCompiler и CpuCompiler ориентированы на графический процессор (GPU) и центральный процессор (CPU) соответственно.
Когда пользователь вызывает Execute... для исполняемого файла, среда выполнения использует BufferAssignment для выделения памяти, а затем вызывает ThunkSequence для запуска операций на устройстве с использованием скомпилированного кода.