-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

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

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

-mpmd-merge-forward-with-backward

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

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

-mpmd-merge-inferred-fragments

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

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

Параметры

-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 = transfer(x0) : m0 -> m1 x2 = transfer(x1) : m1 -> m2 x3 = transfer(x2) : m2 -> m3 x0_1 = transfer(x3) : m3 -> m0 x1_1 = transfer(x0_1) : m0 -> m1

~~>

x1 = transfer(x0) : m0 -> m1 x2 = transfer(x1) : m1 -> m2 x3 = transfer(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

Проверяет, успешно ли выполнено слияние фрагментов, отнесенных к этапам.

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

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

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