-mpmd-copy-topology-from-main

Копирует топологию из основной функции в функции, на которые ссылается mpmd.call.

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

-mpmd-enforce-input-output-equisharding

Обеспечивает соблюдение ограничений эквисардинга для функций MPMD.

Обеспечивает соблюдение ограничений эквисардинга ввода-вывода для функций MPMD путем внедрения TransferOps при необходимости.

Параметры

-constraints : A list of constraint, each enforcing that an input and output should be assigned to the same mesh.

-mpmd-generate-sdy-meshes-from-topology

Генерирует сетки Шарди на основе топологии MPMD.

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

-mpmd-infer-mesh-assign-mesh-func-leaves

_Назначает сетку каждому неиспользуемому вычислению, выходу функции и входу функции с использованием анализа use_set и src set.

Этот проход назначает сетки листьям тела функции (т. е. результатам неиспользуемых вычислений, неиспользуемым аргументам функции и выходным данным функции) путем создания AssignOps или изменения типа с использованием информации use_set и src_set.

Мы также рассматриваем некоторые промежуточные значения как листья для удобства анализа. А именно: операнды mpmd.reduce и mpmd.broadcast рассматриваются как листья, и для них будет создана пара assign-unassign.

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

Этот проход завершится неудачей и выдаст ошибки, если use- и src-sets не будут правильно заполнены для листовых операций.

Предварительное условие: каждая операция имеет непустой src-set, иначе мы делаем вывод о трансферах.

Параметры

-infer-transfers : Whether to create transfers when needed, instead of erroring.
-error-limit     : The number of errors to emit. Set to -1 to emit all errors. Cannot be 0.

-mpmd-infer-mesh-assign-using-input-output-constraints

Назначает сетку входам и выходам в соответствии с ограничениями назначения входов-выходов.

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

Обратите внимание, что это гарантирует, что входные данные будут находиться в той же сетке, независимо от того, будут ли они впоследствии переданы в другие сетки. Но это означает, что нам не следует запускать populate-src-set после этого прохода.

Требует:

  • для любого входа i функции точки входа, которая может быть частью ограничения равноназначения: i имеет MeshTensorType, ИЛИ четко определенный набор использования.
  • для любого выхода o функции точки входа, которая может быть частью ограничения равноназначения: o имеет тип MeshTensorType, ИЛИ четко определенные исходные и используемые наборы.

Где чётко определённый набор использования значения включает все сетки, которым это значение (транзитивно) назначено посредством операций mpmd.assign, и никакие другие сетки. Чётко определённый набор источников включает все сетки, где тензору разрешено существовать, и никакие другие сетки.

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

Параметры

-verbose-logging : Whether to enable verbose logging
-constraints     : A list of constraint, each enforcing that an input and output should be assigned to the same mesh.

-mpmd-infer-mesh-convert-reduce-ops

Преобразует аннотированные операции reduce в операции mpmd.reduce и выравнивает цепочки операций reduce.

Преобразует аннотированные операции reduce в операции mpmd.reduce, а также выравнивает цепочки этих операций reduce.

В символах:

x = добавить(w0, w1) {mpmd.reduce = #mpmd.reduce } y = добавить(x, w2) {mpmd.reduce = #mpmd.reduce } ~~> r = mpmd.reduce (w0,w1,w2)

Параметры

-infer-cross-mesh-reductions : Whether to infer cross-mesh reductions. Will be enabled by default once stable.

-mpmd-infer-mesh-finalize

Применяет окончательную очистку после выведения сетки шаблонов.

Параметры

-infer-transfers : Whether to create transfers when needed, instead of erroring.

-mpmd-infer-mesh-populate-src-set

_Инициализирует src_set для UnassignOps и аргументов func и распространяет набор src.

Этот проход инициализирует src_set и распространяет его, заполняя граф информацией src_set.

Предварительное условие: для того, чтобы аргументы функции имели src_sets, use_set должен быть заполнен.

Инициализация: src_set объекта UnassignOp устанавливается на сетку, к которой он привязан. src_set объекта func arg устанавливается на его use_set.

Распространение: src_sets распространяет данные от операндов к самому оператору, используя пересечение операндов. Подробнее см. PropagateSrcSet .

-mpmd-infer-mesh-populate-use-set

_Инициализирует use_set для AssignOps и распространяет набор использования.

Этот проход инициализирует use_set и распространяет его в обратном направлении, заполняя граф информацией use_set.

Инициализация: use_set AssignOp устанавливается в соответствии с сеткой, которую он назначает.

Распространение: use_sets распространяются в обратном направлении от пользователей к самому оператору, используя объединение пользователей. use_set оператора по определению является объединением use_sets его пользователей, поскольку use_set — это множество транзитивных use.

-mpmd-infer-mesh-rewrite-using-analysis

_Переписывает операции в соответствии с набором использования.

Этот проход назначает операции без сетки, оборачивая их во фрагменты, используя анализы use_set и src_set.

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

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

TODO: jupvfranco — рассмотреть возможность переименования этого прохода, учитывая, что он больше не так сильно зависит от анализа.

Параметры

-max-clones : How many copies of a meshless operation we allow. Setting it to 1 means we never clone the op.

-mpmd-infer-mesh-validate-no-additional-transfers-needed

Проверяет, что для назначения сетки не требуется дополнительных передач.

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

Для операций без сетки, которые не являются операциями func, возникают ошибки, когда:

  1. use_set не содержится в src_set для данной операции, т.е. требуется передача.

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

Параметры

-error-limit : The number of errors to emit. Set to -1 to emit all errors. Cannot be 0.

-mpmd-infer-mesh-validate-src-set-not-empty

_Проверяет, что каждая операция без сетки имеет непустой набор src.

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

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

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

Предварительное условие: перед запуском этого прохода сокращения перекрестной сетки должны быть преобразованы в сокращения операций.

Параметры

-error-limit : The number of errors to emit. Set to -1 to emit all errors. Cannot be 0.

-mpmd-inline-nested-user-exposed-ops

_Встраивает любую предоставленную пользователем операцию mpmd, вложенную в именованное вычисление.

Встраивает любые операции named_computation, named_tensor, broadcast и reduce, вложенные в named_computation, проверяя, что их назначение сетки (если определено) совпадает с назначением родительской сетки.

Параметры

-assignment : Mapping between names (of computations and tensors) and mesh names, and optionally stage ids. E.g., 'n0@m0,n1@m1' defines that names n0 and n1 will be assigned to meshes m0 and m1, respectively. Alternatively 'n0@m0/0,n1@m1/1' means that these names are also assigned to the stages 0 and 1.

-mpmd-insert-nameless-clone-of-negligible-ops

Клонирует незначительные операции за пределами именованных вычислений.

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

Этот проход вообще НЕ изменяет названное вычисление.

-mpmd-introduce-transfers

Создает передачи данных на основе назначений пользовательских сеток.

Создаёт проход, который вводит операции передачи на основе назначений пользовательских сеток. Это включает в себя:

  1. Добавьте UnassignOp в вызовы mpmd, если результат UnassignOp позже будет назначен вызываемой стороне.
  2. Заменяет AssignOp UnassignOp на TransferOp.
  3. Назначить добавление к потребляющей сетке и ввести перенос, если между фрагментами имеется безсеточное добавление.

-mpmd-map-input-output-to-mesh

Назначает сетки входам и выходам функций.

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

Для входных аргументов этот проход:

  1. Преобразует входные тензоры, которые следует поместить в сетку, в тензор сетки.
  2. Обновляет сигнатуру функции.
  3. Добавляет mpmd.unassign перед использованием тензора.

Для выходных аргументов этот проход добавляет mpmd.assign перед возвратом тензора и обновляет сигнатуру функции.

Требуется: каждый индекс ввода/вывода является допустимым, и каждая сопоставленная сетка является допустимой сеткой в ​​топологии.

Параметры

-input-assignment  : Mapping between function input indices and assigned mesh names.E.g., '0@m0,1@m1' defines that input with index 0 will be assigned to mesh m0 and input with index 1 will be assigned to mesh m1.
-output-assignment : Mapping between function output indices and assigned mesh names.E.g., '0@m0,1@m1' defines that output with index 0 will be assigned to mesh m0 and output with index 1 will be assigned to mesh m1.

-mpmd-map-named-ops-to-mpmd-ops

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

Создаёт проход, опционально присваивая mpmd.named_tensor функции Assign(Unassign(%v)) (зависит от наличия записи в assignment ) и сопоставляя каждое named_computation с сеткой, используя определяемое пользователем сопоставление между named_computations и именами сеток. Это означает замену каждого named_computation фрагментом и создание AssignOps для операндов и UnassignOps для результатов этих фрагментов. Представленный шаблон Assign(Unassign(%v)) переписывается в Transfer(%v). После этого прохода операции named_computation/named_tensor больше не будут существовать.

Требуется: все named_computations и named_tensors должны находиться на верхнем уровне функции.

Параметры

-assignment : Mapping between names (of computations and tensors) and mesh names, and optionally stage ids. E.g., 'n0@m0,n1@m1' defines that names n0 and n1 will be assigned to meshes m0 and m1, respectively. Alternatively 'n0@m0/0,n1@m1/1' means that these names are also assigned to the stages 0 and 1.

-mpmd-simplify-named-computations

Упрощает входы и выходы именованных вычислительных операций.

Упрощает каждое именованное вычисление независимо. В частности, он:

  • дедуплицирует результаты и соответствующие им возвращаемые значения;
  • дедуплицирует операнды и соответствующие им аргументы блока;
  • удаляет результаты, соответствующий возвращаемый операнд которых является блочным аргументом оператора;
  • удаляет операнды, соответствующий аргумент блока которых больше не имеет применения (или не имел его изначально); и
  • удаляет неиспользуемые результаты.
  • заменяет шаблон arg -> stablehlo.optimization_barrier -> return в именованном вычислении на шаблон arg -> return , что позволяет дополнительно упростить задачу.

-mpmd-validate-named-ops-in-mpmd-func

Проверяет, что именованные операции вложены только в функции mpmd.

Проверяет, что NamedComputationOp и NamedTensorOp вложены только в функции mpmd, т. е. функции с атрибутом топологии.