Диалект «mpmd»

Операции

mpmd.assign (mpmd::AssignOp)

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

Синтаксис:

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

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

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

Название сетки результирующего типа должно соответствовать сетке в топологии, а её глобальный тип должен быть идентичен типу операнда.

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

Черты: 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.

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

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

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

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

Traits: 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 — это удобочитаемая метка для обозначения точки отправления. Она предназначена для использования в целях отладки.

Параметры:

ReductionAttr

Обозначает уменьшение.

Синтаксис:

#mpmd.reduction<
  `none` | `add` | `max` | `min` | `mul` | `or` | `and`   # reduction_type
>

Параметры:

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

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

Синтаксис:

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

Параметры:

UserOriginAttr

Источник вычислений, заданных пользователем.

Синтаксис:

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

Параметры:

Типы

MeshTensorType

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

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

Параметры:

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

ReductionType

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

Случаи:

EdgeNodeType

перечисление типов узлов ребер

Случаи:

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

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

Случаи:

ReductionOp

Сокращение открытого списка

Случаи: