-mpmd-absorb-inferred-fragments

Фрагменты корня поглощают выведенные фрагменты.

Заставляет корневые фрагменты поглощать выведенные фрагменты, т. е. путем слияния выведенных фрагментов-производителей/потребителей в корневые фрагменты, где корневой фрагмент — это любой фрагмент, который:

• фрагмент пользователя, или
• не используется никаким другим фрагментом (например, фрагментом, используемым только оператором возврата или передачи), или
• не пользователь значения, созданного каким-либо другим фрагментом (например, пользователь аргументов блока или передач).

Для этого проход применяет следующие шаблоны, пока не достигнет фиксированной точки:

(1) Если задан корневой фрагмент rf и существует выведенный фрагмент ipf , такой что ipf является производителем rf , а rf — ближайшим потребителем ipf , то ipf объединяется с rf .

И вдвойне:

(2) При наличии корневого фрагмента rf , если существует выведенный потребитель icf такой, что icf является потребителем rf , а rf является ближайшим производителем icf , то icf объединяется с rf .

Это означает, что мы сохраняем структуру/форму программы, определенную пользователем, посредством именованных вычислений и назначения этапа/сетки.

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

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

Параметры

-absorb-on-entry-point-function : Whether to absorb inferred fragments into user-defined fragments on entry-point functions, in addition to targets of mpmd.calls.

-mpmd-call-inline

Встраивает все операции mpmd.call .

Встраивает операции mpmd.call , копируя их атрибуты во все встроенные операции.

-mpmd-copy-constants

Копирует константы, созданные в одном фрагменте, их потребителям.

Копирует константы из фрагментов-производителей в фрагменты-потребители, возможно, посредством трансферов.

Пример:

%f = fragment () () {
  return constant
}
%t = transfer %f
fragment (%t) (%arg) {
  op(... %arg ...)
  ...
}

 ~~>

%f = fragment () () {
  return constant
}
%t = transfer %f
fragment (%t) (%arg) {
  %c = constant
  op(... %c ...)
  ...
}

Это может положительно сказаться на производительности выполнения: позволяет проводить потенциальные оптимизации, объединяя константу с её пользователями, и избегает передачи констант. Кроме того, это улучшает использование памяти: мы сокращаем объём, необходимый для параметров вычислений.

-mpmd-erase-unused-callee-block-arguments

Удаляет любой аргумент вызываемого блока mpmd, который не используется вычислением (hlo).

Удаляет неиспользуемые аргументы блока из функций, вызываемых mpmd.calls. Аргумент блока считается неиспользуемым, если он не используется или используется только терминатором функции, то есть не используется ни одним вычислением hlo.

-mpmd-fragment-dce

Устраняет неиспользуемые аргументы/результаты фрагментов и упрощает области фрагментов.

Удаляет неиспользуемые аргументы и результаты фрагментов, одновременно упрощая области фрагментов, по сути устраняя мертвый код MPMD.

-mpmd-fragment-dedup

Удаляет все дублирующиеся операнды и приводит к фрагментам.

Удаляет все дублирующиеся используемые аргументы и результаты во фрагментах. При этом дублирующиеся аргументы и результаты останутся неиспользованными. Для удаления неиспользуемых аргументов и результатов следует выполнить другие проходы.

-mpmd-from-unroll-to-call-counter

Преобразует счетчик разворотов attr вызова op в счетчик вызовов attr.

Всякий раз, когда у операции вызова есть атрибут unroll_counter , этот проход заменяет его атрибутом call_counter . Это необходимо в случаях, когда последовательность вызовов получается в результате развёртывания цикла в MLIR (например, через -mpmd-unroll-for-loops ), а не развёртывания цикла на уровне Python.

-mpmd-merge-forward-with-backward

Объединить прямые фрагменты с обратными фрагментами.

Объединить прямой фрагмент производителя с обратным фрагментом потребителя, если первый расположен непосредственно перед вторым. Это справедливо только для последнего этапа в расписании 1F1B, поэтому фрагменты на предыдущих этапах не будут объединены, что и предполагается.

-mpmd-merge-inferred-fragments

Объединяет выведенные фрагменты с фрагментами, определенными пользователем.

Объединяет выведенные фрагменты с пользовательскими или другими выведенными фрагментами. Этот проход полезен для очистки/упрощения модуля и может быть полезен после других проходов компилятора, которые вводят выведенные фрагменты, в то время как указанный выше -mpmd-transfer-aware-merge более инвазивный и должен использоваться только в целях оптимизации.

Если clone_inferred_fragments=true , то этот проход слияния позволяет клонировать определённые фрагменты. В частности, если мы встречаем пару фрагментов f1 и f2, таких что:

• f2 использует f1, а
• f1 выведен, чист и достаточно прост (единственная операция без возврата и единственный результат), затем мы объединяем клон f1 с f2, т.е. сам f1 (и другие пользователи) остаются независимыми от f2. Объединять выведенные фрагменты-производители без клонирования может быть нежелательно, поскольку это может создать ненужные зависимости между фрагментами. Например,

%inferred = frag m1 { return stablehlo.const … }
%frag1 = frag m1 (%inferred, …)
%frag2 = frag m1 (%inferred, …)

~>

%inferred_frag1 = frag m1 (…) { … return const_m1, … }
%frag2 = frag m2 (inferred_frag1, …)

Итак, frag2 теперь зависит от inferred_frag1, и мы создаем зависимость.

Однако иногда нам требуется выполнить слияние на месте, например, когда выведенный фрагмент содержит внутри себя коллективы.

Параметры

-clone-inferred-fragments : Whether to clone inferred fragments. Chains of clonable fragments are merged one-by-one into their consumers and recursively.
-merge-any-consumer       : Whether to merge with any consumer or only the closest consumer.
-merge-sideways           : Whether to merge with the next fragment in the same mesh (neighbor), even if not a consumer.

-mpmd-merge-transfers

Объединяет наборы переводов, которые имеют одни и те же фрагменты производителя и потребителя.

Объединяет наборы передач с одним и тем же типом полезной нагрузки, использующие одни и те же фрагменты производителя и потребителя. Значения полезной нагрузки этих передач содержат меньше элементов, чем заданное пороговое значение, не сегментируются и не хранятся в pinned_host.

Объединение набора передач означает: объединение переданных значений на стороне производителя и разделение их на стороне потребителя.

-mpmd-merge-user-fragments-into-scheduling-units

Объединение пользовательских фрагментов до проходов планирования конвейера.

Объединяет пары пользовательских фрагментов для совместного использования с проходами планирования конвейера.

-mpmd-move-transfers-to-producer

Перемещает трансферы рядом со своими производителями.

Перемещает трансферы рядом с их производителями: если операнд является аргументом блока, перемещает трансфер в начало блока, в противном случае перемещает его после определяющей операции.

-mpmd-remove-transfer-cycles

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

Удаляет циклы переноса.

Например, в символах:

x1 = перевод(x0) : m0 -> m1 x2 = перевод(x1) : m1 -> m2 x3 = перевод(x2) : m2 -> m3 x0_1 = перевод(x3) : m3 -> m0 x1_1 = перевод(x0_1) : m0 -> m1

~~>

x1 = перевод(x0) : m0 -> m1 x2 = перевод(x1) : m1 -> m2 x3 = перевод(x2) : m2 -> m3 x0_1 = x0 x1_1 = x1

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

Обратите внимание, что это может привести к увеличению нагрузки на память, поскольку передача данных с устройства на устройство и обратно означает, что существует период, когда данные отсутствуют на устройстве. Поэтому мы делаем это только в том случае, если цикл содержит только передачи данных с устройства на устройство, например, поскольку цикл device -> host -> device может быть связан с памятью.

Здесь не используется канонизатор MLIR, поскольку он не гарантирует, что все будет канонизировано, а также его применение обходится дороже.

-mpmd-rule-based-merge

Объединяет фрагменты на основе правил, определенных пользователем.

Объединяет фрагменты на основе указанного списка правил, каждое из которых определяет список исходных фрагментов для объединения (по информации об их фрагментах) и целевую информацию для маркировки объединенного фрагмента.

Параметры

-rules                       : A list of fragment merge rules, each with a list of source fragment infos and a target fragment info.
-remove-control-dependencies : Whether to remove control dependencies at the end of the pass.

-mpmd-scheduling-units-verifier

Проверяет, содержит ли программа требуемые единицы планирования.

-mpmd-sink-negligible-ops-into-call-op

Тратит незначительные операции на операции вызова (т. е. вызываемую функцию).

Переводит (незначительное) количество операций в вызываемые функции: если существует операция с нулевыми операндами и единственным результатом, который используется как конкретный операнд всех операций вызова той же функции, то мы переводим его в эти операции вызова, то есть клонируем его в вызываемую функцию и заменяем все использования соответствующего аргумента клоном. Утраченные операции удаляются из вызывающей функции, а неиспользуемые аргументы вызываемой функции (и операнды соответствующих операций вызова) удаляются. Обратите внимание, что это может потенциально дублировать вычисления между многими микропакетами при использовании операций вызова для микропакетирования. Однако, скорее всего, это вычисление пренебрежимо мало, поскольку не принимает операндов.

-mpmd-split-and-prioritize-transfer-independent-computations

Разделяет обратные фрагменты на основе переданных результатов.

Разделяет фрагмент на два, чтобы можно было начать вычисления заранее. То есть, мы разделяем фрагмент на два фрагмента A -> B, где A не зависит от результата переноса и имеет максимальный размер, а B зависит от результата переноса.

-mpmd-split-bwd-fragments

Разделяет обратные фрагменты на основе переданных результатов.

Разделяет фрагменты в обратном направлении, так что любые вычисления, не попадающие в перенесённые результаты, становятся самостоятельными фрагментами. Исходный фрагмент возвращает некоторые остаточные значения, которые передаются в качестве дополнительных операндов в разделённые фрагменты.

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

-mpmd-uniquify-function-inputs-outputs

Делает уникальным любое значение, возвращаемое несколько раз, или любой аргумент блока, непосредственно возвращаемый функцией.

Если функция возвращает одно и то же значение несколько раз, она создаёт несколько версий этого значения, присваивая фрагмент ячейке этого значения, который возвращает значение несколько раз. После этого прохода каждый возвращаемый операнд уникален. Это важно для того, чтобы гарантировать размещение соответствующих результатов в разных буферах, как в следующем примере jax.jit :

def f(x):
  y = x + x
  return y, y

z1, z2 = f(5)
z1 += 1
print(z1) ~~> 6
print(z2) ~~> 5

Аналогично, если функция возвращает аргумент блока, этот проход создает фрагмент идентификации для этого аргумента блока, гарантируя, что значения передаются в функцию по значению, а не по ссылке.

Параметры

-use-transfer-instead-of-fragment : Whether to use mpmd.transfer or mpmd.fragment for uniquification

-mpmd-unroll-for-loops

Полностью разворачивает циклы mpmd.for .

Создает проход, который полностью разворачивает операции mpmd.for , прикрепляя атрибут unroll_counter к каждой развернутой операции.

Требуется: фактор развертывания должен быть равен количеству итераций.

-mpmd-verify-stage-merging

Проверяет, что объединение фрагментов, назначенных на этапы, прошло успешно.

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

Два фрагмента эквивалентны с точки зрения назначения и счетчиков, если а) они назначены на одну и ту же сетку, б) они назначены на один и тот же этап, в) они имеют одинаковое количество транспонирований, и г) либо у обоих одинаковый счетчик вызовов, либо у одного из них не определен счетчик вызовов (т. е. неопределенный счетчик вызовов совпадает с любым счетчиком вызовов).

Это необходимо для того, чтобы гарантировать пользователю, что любые вычисления, назначенные на один и тот же этап, выполняются непрерывно.