В этом документе описывается путь модуля высокоуровневой оптимизации (HLO) XLA от его начального состояния до конечного исполняемого файла. Иногда мы будем опускать слово «модуль» и называть его просто «HLO».
Предварительная оптимизация HLO
Начнём с модуля HLO до оптимизации. HLO до оптимизации не содержит операций ( ops ), которые считаются внутренними для XLA, таких как fusion ) или bitcast . На этом этапе у операций нет структуры (layout), или, если она есть, она игнорируется. 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 , 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 он преобразуется в линейную последовательность thunk-запросов для конкретного бэкэнда (CPU или GPU).
В XLA Thunk — это абстракция самодостаточной единицы работы, выполняемой средой выполнения. Это может быть запуск скомпилированного ядра, конкретная операция, вызов библиотеки, конструкция управления потоком выполнения, коллективная коммуникация и так далее. Thunk Sequence представляет собой весь исполняемый файл для конкретного бэкенда.
Излучение грома
Процесс преобразования запланированного вычисления HLO в последовательность thunk-запросов называется «генерацией thunk-запросов». В каждом бэкэнде это обрабатывается специальным классом-генератором.
Для бэкенда GPU это обрабатывается функцией IrEmitterUnnested . Функция EmitHloComputation перебирает запланированный список инструкций HLO в вычислении и передает управление специализированному методу Emit... (например, EmitFusion , EmitConvolutionThunk , EmitWhile ). Каждый из этих методов создает соответствующий(ие) объект(ы) Thunk и добавляет их к последовательности Thunk.
Для бэкенда ЦП эту роль выполняет ThunkEmitter , который организован аналогичным образом. Финальная ThunkSequence встраивается в CpuExecutable .
Обратите внимание, что каждая инструкция в вычислении входа модуля HLO может соответствовать отсутствию ( kTuple , kConstant , ...), одной или нескольким (например, инструкции сортировки) вспомогательным функциям в итоговой последовательности вспомогательных функций.
Буферы команд: оптимизация выполнения на графическом процессоре
Современное аппаратное обеспечение GPU позволяет записывать последовательность операций GPU (запуски ядра, копирование памяти и т. д.) один раз, а затем воспроизводить эту последовательность многократно с минимальной нагрузкой на ЦП. Это критически важная оптимизация производительности, особенно для рабочих нагрузок с множеством небольших, быстро запускаемых ядер. XLA использует Command Buffer в качестве абстракции CUDA Graphs или HIP Graphs. Основной интерфейс определен в GpuCommandBuffer .
Буфер команд представляется в последовательности thunk объектом CommandBufferThunk .
Излучатель не генерирует этот thunk напрямую из инструкций HLO. Вместо этого это делается с помощью функции CommandBufferConversionPass , которая выполняется непосредственно над ThunkSequence.
На этом этапе определяются смежные подпоследовательности совместимых thunk-функций (например, серия KernelThunk и GemmThunk ). Затем найденная подпоследовательность заменяется одной CommandBufferThunk . Новая thunk-функция инкапсулирует логику исходных thunk-функций в виде списка легковесных объектов Command. Когда CommandBufferThunk выполняется впервые на заданном потоке GPU, она «записывает» свою последовательность команд в аппаратный буфер команд. При всех последующих выполнениях она просто отправляет одну команду GPU для «воспроизведения» записанной последовательности. Это позволяет избежать накладных расходов ЦП на запуск каждого отдельного ядра.
Исполняемый файл
Конечным продуктом конвейера компиляции XLA является самодостаточный, платформенно-специфичный исполняемый файл . Этот объект инкапсулирует всю информацию, необходимую для запуска скомпилированной программы на целевом устройстве, таком как ЦП или ГП. Он служит связующим звеном между компилятором и средой выполнения. Современные среды выполнения, такие как PJRT, используют несколько более высокоуровневые абстракции (см. PjRtExecutable ), но в конечном итоге они представляют собой оболочку для исполняемого файла, специфичного для бэкэнда.
Executable файл содержит несколько ключевых элементов информации, генерируемых в процессе компиляции. Хотя точное содержимое может различаться в зависимости от бэкенда, обычно оно включает в себя:
Скомпилированный код: это низкоуровневый машинный код, который будет выполняться на устройстве. Для центральных процессоров это обычно один или несколько объектных файлов. Для графических процессоров это скомпилированный код устройства в формате PTX или HSACO, который загружается на графический процессор во время выполнения.
План выполнения (ThunkSequence): ядро логики среды выполнения. Это линейная последовательность объектов Thunk. Каждый объект Thunk представляет собой отдельную единицу работы, например, запуск ядра, вызов библиотечной функции (например, cuBLAS) или обработка потока управления. Среда выполнения выполняет программу, итерируя по этой последовательности.
Структура памяти (BufferAssignment): Этот важный фрагмент метаданных, создаваемый BufferAssigner, описывает полную структуру памяти для вычислений. Он определяет размер каждого буфера, а также то, как память выделяется и повторно используется для параметров, выходных данных и временных значений. Среда выполнения использует эти данные для выделения памяти устройства и передачи правильных указателей каждому thunk-объекту.
(необязательный) модуль HLO: для отладки и профилирования исполняемый файл часто сохраняет ссылку на окончательный, оптимизированный модуль HloModule, из которого он был скомпилирован.
Создание конечного исполняемого файла организуется компилятором для каждого конкретного бэкенда. Метод RunBackend реализации компилятора является заключительным этапом процесса компиляции, который упаковывает все скомпилированные артефакты в объект Executable. GpuCompiler и CpuCompiler ориентированы на GPU и CPU соответственно.
Когда пользователь вызывает Execute... для исполняемого файла, среда выполнения использует BufferAssignment для выделения памяти, а затем вызывает ThunkSequence для запуска операций на устройстве с использованием скомпилированного кода.