От HLO до Thunks

Предварительная оптимизация HLO

Начнём с предоптимизационного HLO. Предоптимизационный HLO не содержит операций, которые считаются внутренними для 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) или переписывание библиотечных вызовов. Кроме того, на этом этапе запускается автонастройка, где мы выбираем наилучший алгоритм для свёрток или точек, наилучшее разбиение на тайлы для точек, обрабатываемых устаревшим эмиттером Triton GEMM, или решаем, следует ли использовать Triton или Cublas для определённого слияния точек.

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=.* (или специальное регулярное выражение, чтобы ограничить выгрузку определёнными проходами).

Планирование

Модуль HloModule без планирования всё ещё имеет некоторую степень свободы в выборе порядка выполнения операций. В принципе, любая топологическая сортировка в соответствии с отношением операнд/результат и зависимостями управления подходит. Планирование задаёт определённый порядок. Это влияет на объём необходимой памяти, поскольку мы не можем повторно использовать буфер, пока не будут обработаны все его читатели. На начальном этапе мы пробуем различные алгоритмы планировщика и выбираем расписание, минимизирующее пиковое потребление памяти.

В качестве последующего действия мы запускаем проход LatencyHidingScheduler , который пытается максимизировать перекрытие вычислений и коммуникаций, но может снова увеличить использование памяти.

После планирования мы запускаем HloRematerialization , которая пытается сократить использование памяти в случае, если пиковое потребление превышает объём доступной памяти. Это происходит за счёт производительности, например, некоторые слияния могут быть разделены, а некоторые операции могут быть дублированы для сокращения времени жизни буфера. Если происходит рематериализация, потенциально имеет смысл поискать способы на стороне модели уменьшить объём требуемой памяти (например, уменьшить размер партии).

Thunks и CommandBuffers

Будет определено

Назначение буфера

Непосредственно перед переходом в LLVM IR мы выполняем проходы назначения буфера, которые назначают буферные срезы каждой инструкции в графе HLO. Назначение буфера происходит в несколько этапов.

  1. HloDataflowAnalysis присваивает инструкциям HloValues (по сути, логические буферы). Для операций, выполняемых на месте, HloValue операнда может быть использовано повторно. Операция может определять более одного HloValue (например, с формой результата в виде кортежа).

  2. HloAliasAnalysis пытается объединить буферы для операций совмещения имен и вычисляет сопоставление HloValue с HloBuffer .

  3. BufferAssignment вычисляет сопоставление буферов HloBuffers с буферными срезами внутри большого буфера таким образом, чтобы один и тот же буферный срез не использовался для разных HloBuffers с перекрывающимся временем жизни. Для операций, которые могут быть псевдонимами, допустимо небольшое перекрытие (время окончания одного HloBuffer может совпадать с временем начала другого HloBuffer ). При использовании флага -xla_dump_to некоторая информация о назначении буфера выгружается в файл с суффиксом "after_optimizations-buffer-assignment.txt".