Эта страница была переведа с помощью Cloud Translation API.
Switch to English

MLIR CodeGen для XLA

XLA работает с HloInstruction и выполняет много оптимизаций в этом представлении, разделяя многие из них между целевыми устройствами. В какой-то момент вычисляется линейное расписание, и буфер памяти назначается каждому значению статически. Специфичный для устройства кодогенератор работает, просматривая эту последовательность и вызывая «эмиттеры» для генерации представления, подходящего для устройства (например, единственная функция LLVM для вычисления XLA на ЦП или последовательность «переходов», инкапсулирующих операции графического процессора и, возможно, сгенерированные PTX, когда нацелен на GPU).

В качестве промежуточного шага мы сейчас находимся в процессе перехвата процесса сразу после того, как XLA завершает фазу назначения буфера и вместо этого генерирует модуль lhlo диалекте lhlo . Оттуда мы выполняем кодогенерацию с использованием компонентов MLIR (в основном Linalg, affine и диалект графического процессора) в зависимости от устройства.

Ниже приведен план записи по постепенной миграции XLA / GPU с использованием lhlo в качестве входных данных для кодогенератора.

Задачи

Хост Устройство
Формат ввода HloInstruction * (Задание 1) HloInstruction * (Задание 1)
Выходной формат xla :: Thunk (Задача 2) LLVM IR (Задача 3)
  • Задача 1 изменяет формат ввода хоста и устройства с HloInstruction * на LHLO.
  • Задача 2 изменяет формат вывода хоста с thunks на «некоторую посадочную площадку для хоста» (см. Ниже).
  • Задача 3 переносит вывод устройства из LLVM IR в некоторую форму MLIR. Это необязательно для этого проекта, подробности см. В разделе «Миграция устройства LLVM IR».

В этом проекте приоритет отдается сквозным работающим моделям с максимально активными излучателями LHLO. Это означает, что следующий список задач упорядочен по приоритету:

  • Сделайте XLA / GPU совместимым с эмиттерами LHLO, при этом существующие преобразователи и эмиттеры не будут изменены.
  • Удалите ссылки на HloInstruction * в LHLO в каждом конкретном случае:
    • Переключите унаследованный эмиттер на эмиттер на основе MLIR (например, Linalg) или
    • Механически преобразовать существующий эмиттер в представление MLIR (перейти на стандартный с помощью диалекта графического процессора).

Миграция Thunks (Задача 2)

xla :: gpu :: Thunk - это структура данных, которая:

  • Может быть вызван с хоста (xla :: gpu :: Thunk :: ExecuteOnStream ()).
  • В своих подклассах переносит различные данные.
  • Взаимодействует с BufferAllocation :: Slice и StreamExecutor.
  • Запускает ядра
  • Вызывает все библиотеки времени выполнения.

В стоимость входит:

  • Представление данных конфигурации для конкретной операции (например, конфигурации свертки).
  • Перенос формы операции и форм операндов.
  • Представление дерева преобразователей (while, condition и т. Д.).

Работа по миграции не зависит от миграции LHLO / эмиттера. При ограниченных ресурсах приоритет отдается миграции LHLO / эмиттера.

У нас есть несколько вариантов, как опустить часть на стороне хоста из LHLO:

  • TFRT
    • (Pro) отличные обертки CUDA и HIP для использования.
    • (Pro) легко реализовать библиотечные вызовы (cuDNN, cuBLAS, cuFFT и т. Д.), Поскольку операции TFRT интерпретируются кодом C ++.
    • (Минусы) сторона хоста находится в стадии разработки и не тестируется.
  • Джитовый код ЦП
    • (Pro) отличная низкая способность. Создайте несколько циклов и условий, и готово.
    • (Против) GPUDialect еще не моделирует цепочки / потоки / асинхронность / распределение устройств.
    • (Против) Поддержка среды выполнения CUDA / HIP минимальна (путь к инструментарию, версия, динамическая загрузка и т. Д.).
  • Существующая (интерпретирующая) среда выполнения XLA

Решение: использовать TFRT, но также поддерживать jiting код процессора в TFRT.

Миграция устройства LLVM IR (Задача 3)

Элементный эмиттер генерирует целевую операцию, заполняя ее элемент за элементом. Каждый выходной элемент зависит от набора элементов из операндов. Все элементы описываются путем объединения буфера с динамическими индексами. Достаточно описать почти все «математические» операции, но по соображениям производительности только большая часть «математических» операций реализована непосредственно в (Cpu | Gpu) ElementalIrEmitter.

ElementalIrEmitter уникален тем, что:

  • Большая часть кода распределяется между XLA / GPU и CPU.
  • Он представляет собой большую часть операций, видимых в моделях, включая все операции по элементам.
  • Большинство слияний зависят исключительно от ElementalIrEmitter.
  • Он структурно прост, поскольку описывает DAG зависимости данных между элементами op и элементами операндов.
  • В основном это портативный и высокоуровневый (например, в отличие от GPU kReduce и GPU kCopy).
  • Поддержка динамических форм проста, по крайней мере, для поэлементных операций.

Теперь, для всех операций, с элементарно-генерируемыми или нет, есть несколько разновидностей конечного состояния каждой операции XLA:

  1. Код устройства остается LLVM IR.
  2. Реорганизуйте старый эмиттер, чтобы он был похож на LHLO -> MLIR LLVM Dialect:
    • (Стоимость) Будет одноразовой работой, если мы захотим в конечном итоге перейти на Стандарт.
    • (Преимущество) Это просто и механически. Можно сделать в короткие сроки.
    • (Преимущество) Это не приносит большей пользы по сравнению с (1).
  3. Реорганизуйте старые эмиттеры, чтобы они были похожи на LHLO -> MLIR GPU + Standard + Loops:
    • (Стоимость) Перевод существующих эмиттеров до уровня Standard создает некоторые проблемы. Указатели и GEP необходимо преобразовать в MemRefs и SubViews. Еще одна проблема - обеспечение полноты amdgpu.
    • (Стоимость) XLA / GPU сильно зависит от метаданных LLVM:
      • range для индексов блока / потока.
      • align , dereferenceable , invariant.load , alias.scope , noalias для загрузки / сохранения.
      • llvm.loop.unroll.disable , llvm.loop.unroll.full , llvm.loop.vectorize.enable для последовательных циклов.
    • (Преимущество) Может быть долгосрочным. Более портативный.
  4. Сделайте рефакторинг старых эмиттеров LHLO -> Linalg и напишите новые эмиттеры Linalg.
    • (Стоимость) Это зависит от конкретного случая. По сравнению с предыдущими вариантами, новая реализация, которая соответствует производительности XLA, должна пройти тест <-> оптимизировать рабочий процесс, что может быть значительными затратами для некоторых операций.
    • (Преимущество) единый стек; поддержка сообщества; портативность; больше возможностей оптимизации.

Выводы:

  • Не выбирайте (2). (1) или (3) лучше, чем (2). (2) стоит больше, чем (1), так как требует большого количества механического рефакторинга. С (1) мы все еще можем достичь цели, позволяющей XLA улавливать излучатели MLIR. Это делается с помощью LHLO -> LLVM IR -> запускать эмиттеры устаревших устройств.
  • Операции ElementalIrEmitter идут по (4), но не постепенно. Невозможно сделать это op за op, потому что все элементарно генерируемые операции связаны в один и тот же граф. Эта работа также может служить точкой объединения нескольких действующих сил (xla / service / mlir_gpu, генератор ядра, Linalg).
  • Все остальные операции выполняются для (1). В качестве расширенной цели их можно перенести в (3) или (4).

Приоритезация

Хотя все три упомянутые выше задачи можно распараллеливать, при ограниченных ресурсах их необходимо сериализовать. Приоритезация ориентирована на видимые результаты для выполнения каждой задачи.

Приоритеты следующие: Задача 1 (LHLO для унаследованных эмиттеров)> Задача 2 (Thunks)> Задача 3 (передатчики MLIR).

К концу Задачи 1 пользователи XLA могут сгенерировать LHLO (например, генератор ядра) и выполнить их. Формат компиляции не будет сериализуемым MLIR.

К концу задачи 2 LHLO понижается до правильного сериализуемого MLIR. Это обеспечивает автономную компиляцию.

К концу задачи 3 все излучатели XLA в своей реализации основаны на MLIR.

Детальный дизайн

Шаг 1: (Задача 1) Завершите LHLO и заставьте устаревшие эмиттеры принимать LHLO

Этот шаг заставляет все существующие излучатели XLA / GPU взаимодействовать с операциями MLIR. Этот шаг - чистый рефакторинг и NFC.

Этот шаг в основном механический, но стоит отметить следующие несоответствия между не вложенными HloComputation и LHLO:

  • Каждая HloInstruction имеет прямой доступ к своим операндам (DAG потока данных). Напротив, каждая операция LHLO имеет доступ только к своим буферам операндов (разделение между операциями и буферами). Операции LHLO должны пройти через цепочки use-def для доступа к своим операциям операндов.
  • Невложенные унаследованные эмиттеры эмпирически почти никогда не обращаются к своим операндам. Единственное исключение - kReduce.
  • Невложенные устаревшие эмиттеры обращаются к BufferAssignment только для получения срезов, а не для доступа к вспомогательным структурам данных, таким как dataflow_analysis () или alias_analysis (). llvm_ir строит свой собственный alias_analysis () на основе информации о срезах.

Вывод таков, что LHLO должен быть установлен без особых проблем.

Шаг 2: (Необязательно) Поддержка профилирования

Этот шаг необходим только в том случае, если мы начнем отбрасывать часть логики преобразователя XLA (см. Следующий шаг).

Прежде чем фактически включать какие-либо излучатели на основе MLIR, нам нужно профилировать излучатели на основе MLIR.

В настоящее время XLA выполняет собственное профилирование, вызывая таймер StreamExecutor. Таймер под капотом вставляет два события до и после запуска ядра и измеряет время синхронизации между этими двумя событиями.

Существует примерно три подхода к поддержке профилирования в MLIR:

  • Сквозной запуск профилировщика
  • Добавьте операцию профиля для каждой операции в LHLO, используя внедренный профилировщик.

«Сквозной» подход прозрачен для MLIR, но страдает той же проблемой, из-за которой XLA не использует его в первую очередь: вызовы библиотеки, собранные профилировщиком (nvprof / ...), не могут легко относиться к HLO опс. Например, cuDNN запускает несколько ядер для каждого HLO, и трудно сказать, какие ядра соответствуют какому HLO.

Подход "внедренного профилировщика" требует:

  • LHLO, чтобы использовать профилировщик в качестве параметра.
  • вставка profile.start / profile.end до и после каждой операции.
  • переход от этого понижающего профиля. {start, end} к реализации на C ++.

Точное профилирование не может быть легко выполнено для операций, созданных MLIR, поскольку:

  • MLIR не имеет таймера и не зависит от TFRT / StreamExecutor.
  • MLIR нелегко вызывает функции C со сложными параметрами.

Шаг 3: (Задача 2) Миграция Thunks

Отметим, что существует примерно три типа преобразователей:

  • KernelThunk, запускающий ядро.
  • Преобразователи потока управления, которые имеют логику потока управления хостом (условную, в то время как, для, последовательность) и ядра тела запуска.
  • Преобразователи библиотеки: cuDNN, cuBLAS, cuFFT, NCCL и др.

План такой:

  • Сделайте преобразователи (де) сериализуемыми.
  • Помогите улучшить TFRT до состояния, в котором он может поддерживать эту семантику.
  • По мере улучшения состояния переносите отдельные переходники постепенно.

Эти действия заказаны только частично. Фактический порядок выполнения / инженерный параллелизм следует оценивать по ходу выполнения.

Шаг 4: (Задача 3) Перенесенный ElementalIrEmitter

Когда профилирование будет готово, мы можем завершить и настроить все эмиттеры на основе ElementalIrEmitter в MLIR. Затем мы включаем их по умолчанию, предполагая, что все эти излучатели на основе MLIR используют один поток.

Обратите внимание, что миграция XLA / CPU ElementalIrEmitter также выгодна, поскольку они разделяют большую часть кода.

После завершения всех тестов и поиска производительности (TODO: определение паритета производительности) мы включаем новый элементарный эмиттер на основе MLIR и удаляем устаревший ElementalIrEmitter.

Этот шаг также обеспечивает простые переходы слияния (вложенные операции) для последующей миграции.

Шаг 5: поддержка или удаление многопоточности

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

Шаг 6: (Задача 3) Операции с перенесенным устройством

Этот шаг переносит все невложенные операции, после чего мы можем удалить все невложенные эмиттеры.

Это требует перезаписи / рефакторинга для kCopy и kReduce. Над kReduce уже много работы, так что фактический объем работы, которую необходимо сделать, еще неизвестно.