Диалект «mpmd»

Операции

mpmd.assign (mpmd::AssignOp)

Назначить операцию

Синтаксис:

operation ::= `mpmd.assign` attr-dict $tensor `:` functional-type(operands, results)

Назначает локальный тензор сетке, полностью реплицированный внутри этой сетки.

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

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

Источником операции назначения является источник сетки, например named_computation, вывод сетки и т. д.

Черты: AlwaysSpeculatableImplTrait , HasParent<::mlir::func::FuncOp, ForOp>

Интерфейсы: ConditionallySpeculatable , NoMemoryEffect (MemoryEffectOpInterface)

Эффекты: MemoryEffects::Effect{}

Атрибуты:

Операнды:

Результаты:

mpmd.broadcast (mpmd::BroadcastOp)

Трансляционная операция

Синтаксис:

operation ::= `mpmd.broadcast` attr-dict $tensor `:` type($tensor)

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

Черты: AlwaysSpeculatableImplTrait , SameOperandsAndResultType

Интерфейсы: ConditionallySpeculatable , InferTypeOpInterface , NoMemoryEffect (MemoryEffectOpInterface)

Эффекты: MemoryEffects::Effect{}

Операнды:

Результаты:

mpmd.call (mpmd::CallOp)

Специфическая функция вызова MPMD

Синтаксис:

operation ::= `mpmd.call` $callee `(` $tensors `)` attr-dict `:` functional-type(operands, results)

Операция вызова функции. Полезно, например, для обёртывания тела циклов в объявления функций с целью уменьшения размера кода.

Интерфейсы: ArgAndResultAttrsOpInterface , CallOpInterface , SymbolUserOpInterface

Атрибуты:

Операнды:

Результаты:

mpmd.for (mpmd::ForOp)

Для оператора

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

Необязательный коэффициент развёртки, который должен делить количество итераций, может быть указан для развёртывания тела операции на этот коэффициент. Например, при коэффициенте развёртки N тело реплицируется для создания N копий, а количество итераций уменьшается в 1/N раз. Каждая копия, кроме первой, использует результаты предыдущей копии вместо аргументов блока, а индекс итерации умножается на коэффициент развёртки и увеличивается после каждой копии.

Оператор for может принимать и возвращать любые типы, но их TypeID должен быть одинаковым — например, все типы тензоров или все типы сеток MPMD и т. д. Это позволяет нам использовать оператор на разных уровнях, разделяя реализацию и преобразования.

Черты: HLO_PairwiseSameOperandAndResultType , RecursiveMemoryEffects , RecursivelySpeculatableImplTrait , SingleBlockImplicitTerminator<ReturnOp> , SingleBlock

Интерфейсы: ConditionallySpeculatable , LoopLikeOpInterface , OpAsmOpInterface , ShardableDataFlowOpInterface

Атрибуты:

Операнды:

Результаты:

mpmd.fragment (mpmd::FragmentOp)

Операция «Фрагмент»

Назначает вычисление, т. е. блок операций, определенной сетке в топологии MPMD, которая должна выполняться как отдельный фрагмент программы SPMD.

Фрагмент принимает и возвращает только те тензоры сетки, которые назначены той же сетке, что и фрагмент.

Имя сетки фрагмента должно соответствовать сетке в топологии.

Фрагмент включает список источников, т.е. метаданные с информацией об исходных named_computations, которые сформировали этот фрагмент, и staged_id, определяемый только тогда, когда это пользовательский фрагмент, т.е. он имеет непустой список источников. Необязательный параметр in_shardings задаёт шардинг аргументов блока фрагмента, соответствующих операндам. Необязательный параметр out_shardings задаёт шардинг результатов.

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

Черты: HasParent<::mlir::func::FuncOp, ForOp> , IsolatedFromAbove , RecursiveMemoryEffects , RecursivelySpeculatableImplTrait , SingleBlockImplicitTerminator<ReturnOp> , SingleBlock

Интерфейсы: ConditionallySpeculatable , ShardableDataFlowOpInterface

Атрибуты:

Операнды:

Результаты:

mpmd.fragment_call (mpmd::FragmentCallOp)

Операция вызова фрагмента

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

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

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

Типы входных данных и результата функции вызываемой функции должны быть локальными тензорными типами соответствующих тензоров сетки операндов и результатов этой операции соответственно.

Пример:

%2 = mpmd.fragment_call<mesh="m1", origin=[]> @my_fragment(%0, %1) :
  (mesh_tensor<...>, mesh_tensor<...>) -> mesh_tensor<...>

Признаки: HasParent<::mlir::func::FuncOp> , MemRefsNormalizable

Интерфейсы: ArgAndResultAttrsOpInterface , CallOpInterface , SymbolUserOpInterface

Атрибуты:

Операнды:

Результаты:

mpmd.named_computation (mpmd::NamedComputationOp)

Операция с именованной областью действия

Группирует вычисление, т.е. блок операций, и присваивает ему имя и количество транспонирований через атрибут UserOrigin. Этот NamedComputation можно использовать для назначения сетки вычислению в MPMD или для оптимизации.

Счетчик транспонирования (по умолчанию = 0) указывает, было ли указанное вычисление произведено с помощью определенного количества преобразований транспонирования JAX AD.

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

Черты: IsolatedFromAbove , RecursiveMemoryEffects , RecursivelySpeculatableImplTrait , SingleBlockImplicitTerminator<ReturnOp> , SingleBlock

Интерфейсы: ConditionallySpeculatable

Атрибуты:

Операнды:

Результаты:

mpmd.named_tensor (mpmd::NamedTensorOp)

Назначить тензор сетке

Синтаксис:

operation ::= `mpmd.named_tensor` $tensor `name````=```$name attr-dict `:` type($result)

Идентификационная операция, которая связывает результат тензора с заданным именем. Этот NamedTensor можно использовать для назначения сетки тензору в MPMD.

Черты: AlwaysSpeculatableImplTrait , SameOperandsAndResultType

Интерфейсы: ConditionallySpeculatable , InferTypeOpInterface , NoMemoryEffect (MemoryEffectOpInterface)

Эффекты: MemoryEffects::Effect{}

Атрибуты:

Операнды:

Результаты:

mpmd.reduce (mpmd::ReduceOp)

Операция по уменьшению перекрестной сетки

Синтаксис:

operation ::= `mpmd.reduce` `` $reduction attr-dict $tensors `:` functional-type(operands, results)

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

Черты: AlwaysSpeculatableImplTrait , SameOperandsAndResultType

Интерфейсы: ConditionallySpeculatable , NoMemoryEffect (MemoryEffectOpInterface)

Эффекты: MemoryEffects::Effect{}

Атрибуты:

Операнды:

Результаты:

mpmd.return (mpmd::ReturnOp)

Операция mpmd.return завершает работу с регионами, прикреплёнными к операциям mpmd, основанным на регионах. Она вариативна: в качестве аргументов она принимает список значений, типы которых могут быть любыми (но одного вида, например, AnyTensor ), и поэтому может использоваться повторно на различных уровнях стека IR MPMD.

Синтаксис:

operation ::= `mpmd.return` attr-dict $results (`:` type($results)^)?

Черты: AlwaysSpeculatableImplTrait , Terminator

Интерфейсы: ConditionallySpeculatable , NoMemoryEffect (MemoryEffectOpInterface)

Эффекты: MemoryEffects::Effect{}

Операнды:

mpmd.transfer (mpmd::TransferOp)

Операция переноса

Синтаксис:

operation ::= `mpmd.transfer` attr-dict $tensor `:` functional-type(operands, results)

Переносит распределенный тензор из одной сетки в другую.

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

Черты: AlwaysSpeculatableImplTrait , HasParent<::mlir::func::FuncOp>

Интерфейсы: ConditionallySpeculatable , NoMemoryEffect (MemoryEffectOpInterface) , ShardingRuleOpInterface

Эффекты: MemoryEffects::Effect{}

Операнды:

Результаты:

mpmd.unassign (mpmd::UnassignOp)

Операция отмены назначения

Синтаксис:

operation ::= `mpmd.unassign` attr-dict $tensor `:` functional-type(operands, results)

Отменяет назначение полностью реплицированного тензора из сетки.

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

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

Признаки: AlwaysSpeculatableImplTrait , HasParent<::mlir::func::FuncOp, ForOp> , InferTensorType

Интерфейсы: ConditionallySpeculatable , InferShapedTypeOpInterface , InferTypeOpInterface , NoMemoryEffect (MemoryEffectOpInterface)

Эффекты: MemoryEffects::Effect{}

Атрибуты:

Операнды:

Результаты:

Атрибуты

MeshWithOriginsAttr

Связь со своим происхождением.

Синтаксис:

#mpmd.mesh_with_origins<
  ::llvm::StringRef,   # mesh_name
  ::llvm::ArrayRef<OriginAttr>   # origins
>

Параметры:

MeshesWithOriginsAttr

Список сеток с указанием их происхождения.

Синтаксис:

#mpmd.meshes_with_origins<
  ::llvm::ArrayRef<MeshWithOriginsAttr>   # value
>

Параметры:

NamedMeshAttr

Пара с именем и сеткой.

Синтаксис:

#mpmd.named_mesh<
  ::llvm::StringRef,   # name
  sdy::MeshAttr   # mesh
>

Параметры:

OriginAttr

Происхождение назначения сетки.

Синтаксис:

#mpmd.origin<
  ::llvm::StringRef   # origin_label
>

Происхождение назначения сетки.

origin_label — это понятная человеку метка источника. Она предназначена для использования в целях отладки.

Параметры:

СокращениеАттр

Обозначает сокращение.

Синтаксис:

#mpmd.reduction<
  ::mlir::mpmd::ReductionType   # reduction_type
>

Параметры:

ТопологияАттр

Топология именованных сеток.

Синтаксис:

#mpmd.topology<
  ::llvm::ArrayRef<NamedMeshAttr>   # meshes
>

Параметры:

UserOriginAttr

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

Синтаксис:

#mpmd.user_origin<
  ::mlir::StringAttr,   # userName
  int64_t   # transposeCount
>

Параметры:

Типы

MeshTensorType

Тип тензора сетки

Назначает тип RankedTensorType конкретной SPMD-сетке в топологии MPMD программы. Этот тип содержит необязательный шардинг, определяющий, как шардинг тензора осуществляется относительно SPMD-сетки. Если шардинг отсутствует, тензор полностью реплицируется.

Параметры:

Перечисления

Тип сокращения

Атрибут типа сокращения

Случаи:

EdgeNodeType

Перечисление типов граничных узлов

Случаи:

НаправлениеРаспространения

Направление распространения перечисления

Случаи: