-mpmd-convert-sdy-constants

Преобразует операции sdy.constant в stablehlo.constant .

Преобразует любую операцию sdy.constant , не подлежащую свёртыванию, в операцию stablehlo.constant . Нет причин предотвращать свёртывание констант, поскольку мы удаляем шардинги из констант в mpmd-convert-sdy-shardings-to-mpmd-types .

-mpmd-convert-sdy-shardings-to-mpmd-types

Перемещает шардинги из op attrs в типы !mpmd.mesh_tensor .

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

Этот проход также удаляет любое шардинг из операций, которые не имеют MeshTensorType , т. е. операций внутри операций mpmd.fragment .

-mpmd-enforce-user-shardings

Обеспечивает заданное пользователем разделение входов и выходов.

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

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

Предварительное условие:

• Шардинги пользователей устанавливаются на основе аргументов функции и результатов в виде атрибутов.
• Шардинги фрагментов задаются с помощью атрибутов in_shardings и out_shardings .

-mpmd-extract-reshards-from-inter-mesh-transfers

Перемещает перераспределение SPMD вокруг межсеточного трансфера внутрь фрагмента.

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

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

Этот проход применяется только к функциям MPMD в глобальном представлении и с однородной топологией.

Условие: все шардинги указаны как атрибуты операции, а не в типах.

-mpmd-simplify-program

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

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