-mpmd-delay-inferred-fragments
Задерживает выполнение выведенных фрагментов как можно позже.
Перемещает выведенные фрагменты так, чтобы они выполнялись как можно позже, т. е. прямо перед их первым потребителем.
-mpmd-delay-transfers-from-cpu
Задерживает передачу данных от ЦП, чтобы она выполнялась как можно позже.
Переносит операции передачи данных с процессора на устройство непосредственно перед их первым потребителем. Это означает отсрочку выделения памяти для передаваемых буферов, что может быть полезно для использования HBM.
-mpmd-lower-to-fragment-calls
Понижает фрагменты MPMD до фрагментных вызовов и функций.
Заменяет все фрагменты вызовами фрагментов.
Этот проход создает функцию для каждой группы фрагментов, имеющих одинаковое тело и форму сетки, причем в качестве имени символа используется имя первого обнаруженного фрагмента в группе, и добавляет его в таблицу символов.
Тело каждой функции извлекается из тела первого обнаруженного фрагмента в соответствующей группе, а форма сетки определяется топологией по имени сетки этого фрагмента. Если у фрагмента есть атрибуты аргументов, касающиеся алиасинга входов и выходов, они будут назначены атрибутам аргументов функции, находящейся ниже.
Поскольку функции должны иметь уникальные имена, этот проход добавляет индекс к имени всех функций, кроме первой, с тем же исходным именем, т. е. i-я функция с именем «some_name» для i > 0 будет иметь имя «some_name_i».
Параметры
-group-across-meshes : Whether to do more aggressive fragment grouping, across meshes. This may not be desirable for heterogeneous systems.
-verbose-logging : Whether to enable verbose logging
-mpmd-mark-aliasing-and-donation
Отмечает каждый фрагмент псевдонимом или информацией о пожертвовании.
Устанавливает атрибут arg_attrs для операций фрагментации, когда любой из их входов может быть связан с выходом или передан. Каждый вход, для которого может быть связан с выходом, будет иметь атрибут tf.aliasing_output . В противном случае атрибут jax.buffer_donor = true . Например, {arg_attrs =\[{tf.aliasing_output = 0 : i32}, {jax.buffer_donor = true}, {}\]} показывает, что первый вход может быть сопоставлен с выходом 0, второй вход может быть передан, чтобы XLA нашла преобразованный выход, а третий вход не может быть сопоставлен или передан.
-mpmd-mark-fragment-reserved-memory
Отметьте каждый фрагмент объемом памяти, который необходимо зарезервировать для компиляции.
Каждому фрагменту назначается атрибут xla_tpu_user_reserved_hbm_bytes , который сообщает XLA, сколько байтов следует сохранять при компиляции каждого фрагмента. Отслеживая текущие тензоры в сетке, XLA будет знать о фактическом минимальном использовании памяти во время выполнения, и мы можем предотвратить применение оптимизаций в исполняемом файле, которые могли бы увеличить использование памяти сверх допустимого объёма устройства. ПРИМЕЧАНИЕ: этот проход предполагает, что фрагменты выполняются в порядке выполнения программы.
-mpmd-mark-input-output-with-layouts
Распространяет макеты из func args/results во fragment args/results.
Распространяет атрибуты mhlo.layout_mode из входных данных программы на фрагменты, являющиеся потребителями входных данных программы, и распространяет атрибуты mhlo.layout_mode из выходных данных программы на фрагменты, являющиеся источниками выходных данных. Если возвращается аргумент программы, атрибут mhlo.layout_mode распространяется на результаты программы и из них. Входные, выходные данные и аргументы/результаты фрагментов, связанные с операцией передачи, всегда перезаписываются макетом DEFAULT, поскольку передачи поддерживают только макеты по умолчанию.
Если аргумент программы установлен на макет AUTO и используется в нескольких фрагментах, то мы устанавливаем его на макет DEFAULT, чтобы настроить единообразный макет для всех фрагментов.
Если выходные данные программы, а также результат фрагмента, установлены на макет AUTO и используются в других фрагментах в качестве входных данных, мы также устанавливаем его на макет DEFAULT, чтобы настроить единообразный макет для всех фрагментов.
-mpmd-mark-offloaded-input-output
Помечает выгруженные входные и выходные значения, чтобы компилятор знал, что они находятся в памяти хоста.
Помечает аргументы и результаты фрагментов атрибутами, чтобы определить, какие значения находятся в памяти хоста, чтобы эта информация могла использоваться XLA. Также помечает аргументы и результаты функций точки входа, чтобы Pathways мог использовать эту информацию.