'сди' Диалект

Шарди-диалект (SDY)

Диалект Shardy (SDY) определяет представление сегментации тензоров на основе осей, а также дополнительные компоненты API для привязки сегментов к тензорам.

Журнал версий: 0.0.1: Добавлены неуменьшенные оси в TensorShardingAttr.

Операции

sdy.all_gather (sdy::AllGatherOp)

Осуществляет всестороннюю связь по осям.

Синтаксис:

operation ::= `sdy.all_gather` $gathering_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

Собирает фрагменты тензора вдоль осей, указанных в gathering_axes .

Параметр gathering_axes представляет собой список списков осей. Внешний список охватывает измерения тензора. Каждый внутренний список указывает оси, вдоль которых следует выполнить отдельную операцию gather для соответствующего измерения. Она будет применена к сегментации операнда ( tensor ) для получения сегментации результата ( out_sharding ).

Обратите внимание, что out_sharding не используется для определения сегментации результата. Вместо этого сегментация результата определяется сегментацией операнда и параметра gathering_axes , и out_sharding должен соответствовать этой выведенной сегментации.

Пример:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b", "c"}, {}, {"d"}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.all_gather [{"b", "c"}, {}, {"d"}\] %1 out_sharding=<@mesh, [{"a"}, {}, {}\]> : tensor<8x8x8xf32>

Ограничения:

• Необходимо соблюдать ограничения, указанные в Sdy_CollectiveOpInterface .
• Элементы в gathering_axes должны удовлетворять ограничениям, перечисленным в AxisRefListAttr .
• Применение gathering_axes к операнду sharding приводит к получению out_sharding .

Характеристики: SameOperandsAndResultType

Интерфейсы: InferTypeOpInterface , Sdy_CollectiveOpInterface , SymbolUserOpInterface

Атрибуты:

Операнды:

Результаты:

sdy.all_reduce (sdy::AllReduceOp)

Выполните обмен данными по всем осям путем редукции.

Синтаксис:

operation ::= `sdy.all_reduce` $reduction_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

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

Ограничения:

• Необходимо соблюдать ограничения, указанные в Sdy_CollectiveOpInterface .
• reduction_axes должен удовлетворять ограничениям, перечисленным в AxisRefListAttr .
• Сортировка reduction_axes должна производиться относительно сетки.
• Оперы sharding и out_sharding должны иметь эквивалентные размерности сегментирования.
• reduction_axes не должна перекрываться с осью сегментирования и реплицированными осями операнда (она может перекрываться с неуменьшенными осями).
• reduction_axes не должны перекрываться с неуменьшенными осями out_sharding ). Другими словами, out_sharding должен быть продублирован вдоль reduction_axes ) (явно или неявно).

Характеристики: SameOperandsAndResultType

Интерфейсы: CollectiveOpInterface , InferTypeOpInterface , SymbolUserOpInterface

Атрибуты:

Операнды:

Результаты:

sdy.all_slice (sdy::AllSliceOp)

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

Синтаксис:

operation ::= `sdy.all_slice` $slicing_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

Функция sdy.all_slice разделяет тензор на фрагменты вдоль осей, указанных в slicing_axes . Между sdy.all_slice и sdy.all_gather существует алгебраическая двойственность.

slicing_axes — это список списков осей. Внешний список охватывает измерения тензора. Каждый внутренний список указывает оси, вдоль которых следует выполнить срез по соответствующему измерению. Он будет применен к сегментированию операнда ( tensor ) для получения сегментирования результата ( out_sharding ).

Обратите внимание, что out_sharding не используется для определения сегментации результата. Вместо этого сегментация результата определяется сегментацией операнда и slicing_axes , и out_sharding должен соответствовать этой выведенной сегментации.

Пример:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}, {}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.all_slice [{"b", "c"}, {}, {"d"}\] %1 out_sharding=<@mesh, [{"a", "b", "c"}, {}, {"d"}\]> : tensor<8x8x8xf32>

Ограничения:

• Необходимо соблюдать ограничения, указанные в Sdy_CollectiveOpInterface .
• Элементы в slicing_axes должны удовлетворять ограничениям, перечисленным в AxisRefListAttr .
• Применение slicing_axes к операнду sharding дает значение out_sharding .

Характеристики: SameOperandsAndResultType

Интерфейсы: CollectiveOpInterface , InferTypeOpInterface , SymbolUserOpInterface

Атрибуты:

Операнды:

Результаты:

sdy.all_to_all (sdy::AllToAllOp)

Осуществляет связь «все со всеми» по осям.

Синтаксис:

operation ::= `sdy.all_to_all` $params $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

Для каждого кортежа (axes, src_dim, tgt_dim) в списке параметров эта операция разбивает тензор на фрагменты вдоль размерности tgt_dim и осей, указанных в axes , распределяет эти фрагменты вдоль осей и объединяет их вдоль размерности src_dim .

Эта операция по сути представляет собой комбинацию операции сбора данных вдоль src_dim и axes , за которой следует операция разрезания данных вдоль tgt_dim и axes , то есть к размерности сегментирования осей src_dim входного тензора добавляется суффикс размерности сегментирования осей tgt_dim на выходном тензоре.

Метод all-to-all будет применен к сегментации операнда ( tensor ) для получения сегментации результата ( out_sharding ).

Обратите внимание, что out_sharding не используется для определения сегментации результата. Вместо этого сегментация результата определяется сегментацией операнда, src_dim , tgt_dim и axes , и out_sharding должен соответствовать этой предполагаемой сегментации.

Пример:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b"}, {"c"}, {}, {}\]>]>} : tensor<8x8x4x4x32>
%2 = sdy.all_to_all [{"b"}: 0->2, {"c"}: 1->3] %1 out_sharding=<@mesh, [{"a"}, {}, {"b"}, {"c"}\]> : tensor<8x8x4x4x32>

Ограничения:

• Необходимо соблюдать ограничения, указанные в Sdy_CollectiveOpInterface .
• Список параметров не должен быть пустым.
• Для каждого параметра в params :
  • Элементы в axes должны удовлетворять ограничениям, заданным в AxisRefAttr .
  • src_dim и tgt_dim должны быть допустимыми размерностями (неотрицательными и меньше ранга тензора).
  • Любое src_dim или tgt_dim должно быть уникальным для всех параметров.
  • src_dim должно быть отсортировано в порядке возрастания по всем параметрам.
• Перемещение axes из src_dim в tgt_dim в операнде sharding приводит к out_sharding .

Характеристики: SameOperandsAndResultType

Интерфейсы: InferTypeOpInterface , Sdy_CollectiveOpInterface , SymbolUserOpInterface

Атрибуты:

Операнды:

Результаты:

sdy.collective_permute (sdy::CollectivePermuteOp)

Выполняет коллективно-перестановочную коммуникацию для замены осей.

Синтаксис:

operation ::= `sdy.collective_permute` $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

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

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

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

В приведенном ниже примере размер сегментированного тензора равен tensor<1x4x2xf32> , и это сохраняется при коллективной перестановке.

Пример:

sdy.mesh @mesh = <["a"=2, "b"=2, "c"=4, "d"=2, "e"=2, "f"=2]>
%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "c"}, {"f"}, {"d", "e"}\]>]>} : tensor<8x8x8xf32>
%2 = sdy.collective_permute %1 out_sharding=<@mesh, [{"c":(1)2, "b", "f"}, {"a"}, {"e", "d"}\]> : tensor<8x8x8xf32>

Ограничения:

• Необходимо соблюдать ограничения, указанные в Sdy_CollectiveOpInterface .
• Если входные и выходные сегменты имеют разные сетки, то эти сетки должны иметь абсолютно одинаковые оси и разный порядок идентификаторов устройств.
• Для каждого измерения произведение размеров осей сегментирования в out_sharding должно совпадать с размером сегментирования соответствующего операнда измерения.

Характеристики: SameOperandsAndResultType

Интерфейсы: CollectiveOpInterface , InferTypeOpInterface , SymbolUserOpInterface

Атрибуты:

Операнды:

Результаты:

sdy.constant (sdy::ConstantOp)

Постоянная работа

Создает output тензор из постоянного value .

См.: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#constant

Пример:

%output = sdy.constant dense<[[0.0, 1.0], [2.0, 3.0]]> : tensor<2x2xf32>

Характеристики: AlwaysSpeculatableImplTrait

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

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

Атрибуты:

Результаты:

sdy.data_flow_edge (sdy::DataFlowEdgeOp)

Операция на границе потока данных.

Синтаксис:

operation ::= `sdy.data_flow_edge` $input (`sharding````=``` $sharding^)? attr-dict `:` type($result)

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

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

Например:

  y_0, ..., y_n = while (x_0, ..., x_n)
                  ((pred_arg_0,... , pred_arg_n) { ... })
                  ((body_arg_0,..., body_arg_n) {
                    ...
                    return return_value_0, ..., return_value_n
                  })

В то время как операция имеет n ребер потока данных, i-е ребро потока данных находится между источниками x_i , return_value_i и целями y_i , pred_arg_i , body_arg_i .

Функция sdy.data_flow_edge принимает на вход владельца ребра (это может быть любой из целевых объектов, но предпочтительнее результат операции, а не аргумент блока), который не должен иметь других применений. Эта операция не является чистой, поскольку может принимать на вход данные, которые изначально не имели никаких применений.

Объект sdy.data_flow_edge также содержит необязательный сегмент для всех целевых объектов ребра, и этот сегмент должен обновляться вместо сегмента целевых объектов (если его можно прикрепить) во время распространения. Это полезно, когда операция имеет много ребер, поскольку гораздо эффективнее:

• распространяться по каждому краю отдельно.
• Обновляйте сегментацию каждого ребра отдельно, а не всех целей одновременно (например, операция имеет один неизменяемый TensorShardingPerValueAttr для сегментации результатов).
• Добавляйте каждое ребро в список задач отдельно при изменении сегментации источника.

В результате операции шардирование будет происходить между всеми источниками и целями sdy.data_flow_edge так же, как если бы это была обычная операция, где источники являются операндами, а цели — результатами, и используется правило идентичности sdy.op_sharding_rule . Это означает, что прямое распространение происходит от источников к целям, а обратное — от целей к источникам.

Мы не разрешаем определять входные данные для sdy.data_flow_edge операцией SdyDialect , поэтому можем предположить, что они определяются операцией с незарегистрированным атрибутом sdy.sharding .

Характеристики: SameOperandsAndResultType

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

Атрибуты:

Операнды:

Результаты:

sdy.manual_computation (sdy::ManualComputationOp)

Параллельная работа нескольких устройств с ручным объединением операций.

Синтаксис:

operation ::= `sdy.manual_computation` `(`operands`)`
              `in_shardings````=```custom<StrippedTensorShardingPerValueAttr>($in_shardings)
              `out_shardings````=```custom<StrippedTensorShardingPerValueAttr>($out_shardings)
              `manual_axes````=```$manual_axes
              custom<SingleBlockRegionNoBlockId>($body)
              attr-dict
              `:`
              functional-type(operands, results)

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

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

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

Ограничения:

• Элементы во in_shardings и out_shardings должны удовлетворять ограничениям, перечисленным в TensorShardingAttr .
• Количество глобальных и локальных тензорных входов/выходов в области операции должно совпадать.
• В каждом сегменте данных оси, заданные вручную, должны располагаться перед любыми свободными осями.
• В ручных настройках осей нельзя добавлять отступы. А именно, размер габарита должен быть кратен соответствующему размеру, заданному в ручных настройках.
• Глобальная и локальная формы областей действия аргументов/результатов должны совпадать.

Характеристики: IsolatedFromAbove , RecursiveMemoryEffects , SingleBlockImplicitTerminator<ReturnOp> , SingleBlock

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

Атрибуты:

Операнды:

Результаты:

sdy.mesh (sdy::MeshOp)

Именованная сетка

Синтаксис:

operation ::= `sdy.mesh` $sym_name `=` $mesh attr-dict

Определяет новую именованную сетку. Все сетки в модуле должны иметь одинаковое количество устройств (за исключением сеток с одним device_id). Сетка представляет собой операцию Symbol , которая отображается в SymbolTable модуля и может быть указана по своему name .

Характеристики: HasParent<ModuleOp>

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

Атрибуты:

sdy.named_computation (sdy::NamedComputationOp)

Именованная операция вычисления

Синтаксис:

operation ::= `sdy.named_computation` `<`$name`>` `` `(` $operands `)`
              (`in_shardings````=```custom<StrippedTensorShardingPerValueAttr>($in_shardings)^)?
              (`out_shardings````=```custom<StrippedTensorShardingPerValueAttr>($out_shardings)^)?
              custom<SingleBlockRegionNoBlockId>($body)
              attr-dict
              `:` functional-type($operands, results)

Группирует вычисления, то есть блок операций, и присваивает ему имя. Распространение будет происходить внутри/вне этой области так, как если бы все операции были встроены.

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

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

Пример:

%1 = sdy.named_computation<"foo">(%0) (%arg1: tensor<16x32xf32>) {
  sdy.return %arg1 : tensor<16x32xf32>
} : (tensor<16x32xf32>) -> tensor<16x32xf32>

Характеристики: IsolatedFromAbove , RecursiveMemoryEffects , RecursivelySpeculatableImplTrait , SingleBlockImplicitTerminator<ReturnOp> , SingleBlock

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

Атрибуты:

Операнды:

Результаты:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

Работа барьера распространения

Синтаксис:

operation ::= `sdy.propagation_barrier` $input `allowed_direction````=```$allowed_direction attr-dict `:` type($input)

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

Это предотвращает распространение фрагментации данных между использованием результата операции барьера и её операнда.

• FORWARD означает, что фрагменты данных могут передаваться только от операнда к результату.
• BACKWARD означает, что фрагменты могут передаваться только от результата к операнду.
• Значение NONE означает, что в рамках этой операции шардинг невозможен.
• Невозможно указать BOTH , так как это будет излишним.

Характеристики: AlwaysSpeculatableImplTrait , SameOperandsAndResultType

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

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

Атрибуты:

Операнды:

Результаты:

sdy.reduce_scatter (sdy::ReduceScatterOp)

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

Синтаксис:

operation ::= `sdy.reduce_scatter` $reduce_scatter_axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

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

Ограничения:

• Необходимо соблюдать ограничения, указанные в Sdy_CollectiveOpInterface .
• Элементы в reduce_scatter_axes должны удовлетворять ограничениям, перечисленным в AxisRefListAttr .
• Применение функции reduce_scatter_axes к операнду sharding дает значение out_sharding .

Характеристики: SameOperandsAndResultType

Интерфейсы: CollectiveOpInterface , InferTypeOpInterface , SymbolUserOpInterface

Атрибуты:

Операнды:

Результаты:

sdy.replicated_to_unreduced (sdy::ReplicatedToUnreducedOp)

Переместите неявно или явно реплицированные оси на нередуцированные оси.

Синтаксис:

operation ::= `sdy.replicated_to_unreduced` $axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

axes должны быть неявно или явно продублированы в операнде. Эта операция делает их нередуцированными в результате. Мы имеем следующее соотношение:

all-reduce(replicated-to-unreduced(x, axes), axes) = x

Пример:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}, {}, {}\], replicated={"c", "d"}, unreduced={"e"}>]>} : tensor<8x8x8xf32>
%2 = sdy.replicated_to_unreduced {"a", "c", "f"} %1 out_sharding=<@mesh, [{"b"}, {}, {}\], replicated={"d"}, unreduced={"a", "c", "e", "f"}> : tensor<8x8x8xf32>

Ограничения:

• Необходимо соблюдать ограничения, указанные в Sdy_CollectiveOpInterface .
• axes должны удовлетворять ограничениям, перечисленным в AxisRefListAttr .
• axes должны быть отсортированы относительно сетки.
• axes не пусты.
• Входные и выходные сегменты должны иметь одинаковые размерности.
• axes должны быть неявно или явно продублированы в операнде сегментирования.
• inUnreducedAxes + axes = outUnreducedAxes.

Характеристики: SameOperandsAndResultType

Интерфейсы: InferTypeOpInterface , Sdy_CollectiveOpInterface , SymbolUserOpInterface

Атрибуты:

Операнды:

Результаты:

sdy.reshard (sdy::ReshardOp)

Перераспределяет тензор на другой сегмент.

Синтаксис:

operation ::= `sdy.reshard` $input $sharding attr-dict `:` type($result)

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

И ShardingConstraintOp, и ReshardOp добавляют сегментацию к тензору. Срок их действия:

1. Перед распространением шардинга пользователи добавляют параметр ShardingConstraintOp.
2. При распространении шардинга используется операция ShardingConstraintOp. В результатах распространения шардинга операция ShardingConstraintOp отсутствует. Вместо этого при необходимости может быть добавлена ​​операция ReshardOp.
3. Разделитель преобразует операцию ReshardOp в коллективную операцию (или операцию идентификации). В результатах работы разделителя не должно быть операций ReshardOp.

Характеристики: AlwaysSpeculatableImplTrait , SameOperandsAndResultType

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

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

Атрибуты:

Операнды:

Результаты:

sdy.return (sdy::ReturnOp)

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

Синтаксис:

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

Характеристики: AlwaysSpeculatableImplTrait , ReturnLike , Terminator

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

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

Операнды:

sdy.sharded_to_unreduced (sdy::ShardedToUnreducedOp)

Переместите некоторые фрагментированные оси операнда на нередуцированные оси результата.

Синтаксис:

operation ::= `sdy.sharded_to_unreduced` $axes $tensor `out_sharding````=```$out_sharding attr-dict `:` type($result)

axes следует использовать для разделения операнда на части. Эта операция делает их нередуцированными в результате. Мы имеем следующее соотношение:

all-gather(x, axes) = all-reduce(sharded-to-unreduced(x, axes), axes), где all-gather, sharded-to-unreduced и all-reduce применяются к одним и тем же осям.

Пример:

%1 = stablehlo.tanh(%0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a", "b", "c"}, {}, {"d"}\], unreduced={"e"}>]>} : tensor<8x8x8xf32>
%2 = sdy.sharded_to_unreduced [{"b", "c"}, {}, {"d"}\] %1 out_sharding=<@mesh, [{"a"}, {}, {}\], unreduced={"b", "c", "d", "e"}> : tensor<8x8x8xf32>

Ограничения:

• Необходимо соблюдать ограничения, указанные в Sdy_CollectiveOpInterface .
• Элементы axes должны удовлетворять ограничениям, перечисленным в AxisRefListAttr .
• Применение axes к операнду sharding дает значение out_sharding .

Характеристики: SameOperandsAndResultType

Интерфейсы: InferTypeOpInterface , Sdy_CollectiveOpInterface , SymbolUserOpInterface

Атрибуты:

Операнды:

Результаты:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

Ограничивает тензор заданным сегментированием.

Синтаксис:

operation ::= `sdy.sharding_constraint` $input $sharding attr-dict `:` type($result)

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

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

Эта операция может быть либо:

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

Характеристики: SameOperandsAndResultType

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

Атрибуты:

Операнды:

Результаты:

sdy.sharding_group (sdy::ShardingGroupOp)

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

Синтаксис:

operation ::= `sdy.sharding_group` $input `group_id````=```$group_id attr-dict `:` type($input)

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

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

Атрибуты:

Операнды:

Атрибуты

AllToAllParamAttr

Параметр «все ко всем»

Синтаксис:

#sdy.all_to_all_param<
  ::llvm::ArrayRef<AxisRefAttr>,   # axes
  int64_t,   # src_dim
  int64_t   # tgt_dim
>

Кортеж, содержащий оси и исходные/целевые измерения, по которым необходимо выполнить операцию "все ко всем".

Параметры:

AllToAllParamListAttr

Список параметров «все ко всем»

Синтаксис:

#sdy.all_to_all_param_list<
  ::llvm::ArrayRef<AllToAllParamAttr>   # value
>

Параметры:

AxisRefAttr

Ссылка либо на полную ось, либо на разделенную подось.

Синтаксис:

#sdy.axis_ref<
  ::llvm::StringRef,   # name
  SubAxisInfoAttr   # sub_axis_info
>

Ограничения:

• name должно присутствовать в связанном MeshAttr .
• Если sub_axis_info присутствует, он должен удовлетворять ограничениям параметра SubAxisInfoAttr .

Параметры:

AxisRefListAttr

Список ссылок на оси

Синтаксис:

#sdy.axis_ref_list<
  ::llvm::ArrayRef<AxisRefAttr>   # value
>

Ограничения:

• Элементы в value должны удовлетворять ограничениям AxisRefAttr .
• Отсутствуют дублирующиеся ссылки на оси или перекрывающиеся между собой подоси.
• Никакие две смежные оси не являются последовательными подосью одной и той же полной оси, то есть они могут быть объединены в одну подось или в полную ось.

Параметры:

AxisToPropagationDetailsAttr

Детали распространения потока по краям для конкретной оси и источника.

Синтаксис:

#sdy.axis_to_propagation_details<
  ::mlir::sdy::AxisRefAttr,   # axis_name
  ::mlir::sdy::EdgeValueRefAttr,   # source
  ::llvm::ArrayRef<EdgeValueRefAttr>   # targets
>

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

Параметры:

DimMappingAttr

Список факторных индексов для измерения

Пустой список указывает на то, что это пустое сопоставление (оно анализируется/выводится с помощью * ), то есть измерение не сопоставлено ни с одним фактором.

Ограничения:

• Существует как минимум один факторный индекс.
• Индексы факторов должны находиться в диапазоне [0, $factor_sizes ).
• Если факторов несколько, ни один из них не может иметь размер 1.
• Отсутствуют повторяющиеся индексы факторов.

Параметры:

DimensionShardingAttr

Разделение измерений

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

Ограничения:

• Элементы axes должны удовлетворять ограничениям, перечисленным в AxisRefListAttr .
• Если сегментирование измерений имеет приоритет:
  • Приоритет больше или равен 0.
  • Если размерность замкнута, то она имеет как минимум одну ось.

Параметры:

EdgeValueRefAttr

Ссылка на конкретный индекс ребра значения type .

Синтаксис:

#sdy.edge_value_ref<
  ::mlir::sdy::EdgeNodeType,   # type
  int64_t   # index
>

Параметры:

ListOfAxisRefListsAttr

Список ссылок на оси

Синтаксис:

#sdy.list_of_axis_ref_lists<
  ::llvm::ArrayRef<AxisRefListAttr>   # value
>

Параметры:

ManualAxesAttr

Список осей, по которым операция ManualComputationOp выполняется вручную.

Синтаксис:

#sdy.manual_axes<
  ::llvm::ArrayRef<StringAttr>   # value
>

Параметры:

MeshAttr

Сетка осей и список устройств

Синтаксис:

#sdy.mesh<
  ::llvm::ArrayRef<MeshAxisAttr>,   # axes
  ::llvm::ArrayRef<int64_t>   # device_ids
>

Сетка представляет собой список осей и необязательный список идентификаторов устройств, определяющих порядок расположения устройств.

Если список осей пуст

• Если device_ids не указан, это пустая сетка.
• Если параметр device_ids указан, он должен представлять собой одно неотрицательное целое число; мы называем такую ​​сеть максимально сегментированной сетью .

Если предоставлен список осей

• Если указан список идентификаторов устройств, произведение размеров осей должно соответствовать количеству устройств.
• Если список идентификаторов устройств не указан, используется неявный список идентификаторов устройств iota(product(axes)). Для простоты мы также запрещаем указывать список идентификаторов устройств, совпадающий с iota(product(axes)); в этом случае список идентификаторов устройств указывать не следует.
• Это не сетка с максимальным сегментированием, даже если общий размер осей равен 1.

Вот несколько примеров сеток:

• Пустая сетка представляет собой сетку-заполнитель, которую можно заменить во время распространения: <[]>
• Список сетки без осей и один неотрицательный идентификатор устройства, представляющий собой сетку с максимальным сегментированием: <[], device_ids=[3]>
• Сетка с двумя осями и неявными идентификаторами устройств iota(6): <["a"=2, "b"=3]>
• Сетка с двумя осями и явно указанными идентификаторами устройств, определяющими порядок устройств: <["a"=3, "b"=2], device_ids=[0, 2, 4, 1, 3, 5]>

Ограничения:

• Элементы в device_ids должны быть неотрицательными.
• Если axes пуста, размер device_ids может быть равен 0 (пустая сетка) или 1 (сетка с максимальным сегментированием).
• Если axes не пуста,
  • Элементы в axes не должны иметь одинаковых имен.
  • Если указан параметр device_ids , исходный device_ids не iota(product(axis_sizes)) , а отсортированный device_ids равен iota(product(axis_sizes)) .

Параметры:

MeshAxisAttr

Именованные оси в сетке

Синтаксис:

#sdy.mesh_axis<
  ::llvm::StringRef,   # name
  int64_t   # size
>

Параметры:

OpShardingRuleAttr

Указывает, как можно разделить операцию на этапы.

Синтаксис:

#sdy.op_sharding_rule<
  ::llvm::ArrayRef<int64_t>,   # factor_sizes
  ::llvm::ArrayRef<TensorMappingAttr>,   # operand_mappings
  ::llvm::ArrayRef<TensorMappingAttr>,   # result_mappings
  ::llvm::ArrayRef<int64_t>,   # reduction_factors
  ::llvm::ArrayRef<int64_t>,   # need_replication_factors
  ::llvm::ArrayRef<int64_t>,   # permutation_factors
  ::llvm::ArrayRef<int64_t>,   # blocked_propagation_factors
  bool   # is_custom_rule
>

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

%0 = stablehlo.add %arg0, %arg1 {
    sdy.sharding_rule = #sdy.op_sharding_rule<
        ([i, j],[i, j])->([i, j])
        {i=8, j=8}>
} : tensor<8x8xf32>

%1 = stablehlo.dot_general %arg2, %arg3, contracting_dims = [1] x [0] {
  sdy.sharding_rule = #sdy.op_sharding_rule<
      ([i, k],[k, j])->([i, j])
      {i=8, j=16, k=8}>
}: (tensor<8x8xf32>, tensor<8x16xf32>) -> tensor<8x16xf32>

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

Типы факторов:

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

blocked_propagation_factors содержит факторы, по которым распространение шардингов запрещено. Он ортогонален типам факторов. А именно, фактор, по которому распространение заблокировано, может быть любого из типов факторов.

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

Ограничения:

• Количество сопоставлений операндов и результатов должно соответствовать количеству операндов и результатов операции.
• Существует как минимум одно соответствие (нельзя иметь правило для операции без операндов/результатов).
• Ранг каждого TensorMappingAttr соответствует рангу соответствующего типа тензора.
• Для каждой группы факторов ( reduction_factors , need_replication_factors , permutation_factors ):
  • Элементы должны находиться в диапазоне [0, $factor_sizes ].
  • Внутри каждой группы и между группами отсутствуют дублирующиеся индексы факторов.

Параметры:

PropagationEdgesAttr

Метаданные границы распространения для всех этапов распространения.

Синтаксис:

#sdy.propagation_edges<
  ::llvm::ArrayRef<PropagationOneStepAttr>   # value
>

Список сведений о распространении значения по каждой оси, сгруппированный по индексу шага.

Параметры:

PropagationOneStepAttr

Метаданные распространения для каждого этапа.

Синтаксис:

#sdy.propagation_one_step<
  int64_t,   # step_index
  ::llvm::ArrayRef<AxisToPropagationDetailsAttr>   # axis_entries
>

Подробные сведения о распространении по всем осям для одного этапа распространения.

Параметры:

SubAxisInfoAttr

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

Синтаксис:

#sdy.sub_axis_info<
  int64_t,   # pre_size
  int64_t   # size
>

При разделении полной оси на n подосей ось преобразуется в [k_1,...,k_n], и i-я подось может быть выражена произведением размеров всех осей слева от нее m=prod(k_1,...,k_(i-1)) (также называемым предварительным размером) и размером k_i. Следовательно, атрибут sub-axis-info содержит эти два числа и обозначается следующим образом: (m)k для предварительного размера m и размера k.

Ограничения:

• pre-size составляет не менее 1.
• size больше 1.
• pre-size должен делить размер полной оси, то есть и pre-size , и size делят размер полной оси, а подось не выходит за пределы полной оси.
• Размер вспомогательной оси не равен размеру соответствующей полной оси, в этом случае следует использовать полную ось.

Параметры:

TensorMappingAttr

Факторные отображения для каждого измерения тензора.

Синтаксис:

#sdy.tensor_mapping<
  ::llvm::ArrayRef<DimMappingAttr>   # dim_mappings
>

Ограничения:

• Элементы в dim_mappings должны удовлетворять ограничениям, указанным в DimMappingAttr .
• Индексы факторов не дублируются по различным измерениям.

Параметры:

TensorShardingAttr

Разделение тензоров

Синтаксис:

#sdy.sharding<
  ::mlir::Attribute,   # mesh_or_ref
  ::llvm::ArrayRef<DimensionShardingAttr>,   # dim_shardings
  ::llvm::ArrayRef<AxisRefAttr>,   # replicated_axes
  ::llvm::ArrayRef<AxisRefAttr>   # unreduced_axes
>

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

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

Сетку, к которой привязан этот сегмент, можно указать либо с помощью имени символа, ссылающегося на соответствующий символ MeshOp , либо с помощью встроенного MeshAttr .

При сегментировании могут быть нередуцированные оси (указывающиеся параметром unreduced_axes ), то есть тензор не редуцирован вдоль этих осей. Например, если сжимающая размерность matmul сегментирована вдоль оси x как в левой, так и в правой части, результат не редуцирован вдоль x . Применение операции all-reduce к тензору вдоль нередуцированных осей приведет к дублированию тензора вдоль этих осей. Однако тензор с нередуцированными осями не обязательно должен быть немедленно подвергнут операции all-reduce; он может оставаться нередуцированным при передаче в линейные операции, такие как stablehlo.add (при условии, что как левая, так и правая части не редуцированы), и затем подвергаться операции all-reduce. Мы предполагаем, что тип редукции — sum, в будущем могут быть поддержаны и другие типы редукции.

Ограничения:

• Элементы в dim_shardings должны удовлетворять ограничениям, перечисленным в DimensionShardingAttr .
• Элементы в replicated_axes должны удовлетворять ограничениям, перечисленным в AxisRefListAttr .
• Элементы в unreduced_axes должны удовлетворять ограничениям, перечисленным в AxisRefListAttr .
• Если соответствующий тип тензора не является ShapedType , то сегментирование должно иметь ранг 0 и не должно содержать дублированных осей.
• Если это ShapedType , то:
  • Тензор должен иметь ранг.
  • Количество сегментов размерности равно рангу тензора.
  • Размеры 0 не сегментируются.
• В параметрах dim_shardings , replicated_axes и unreduced_axes отсутствуют дублирующиеся ссылки на оси или подоси, перекрывающиеся друг с другом.
• Элементы в replicated_axes и unreduced_axes упорядочены относительно mesh_or_ref (см. AxisRefAttr::getMeshComparator ).

Параметры:

TensorShardingPerValueAttr

Разделение тензора на операнды/результаты операций

Синтаксис:

#sdy.sharding_per_value<
  ::llvm::ArrayRef<TensorShardingAttr>   # shardings
>

Список атрибутов TensorShardingAttr , по одному для каждого операнда/результата операции.

Ограничения:

• Элементы в shardings должны удовлетворять ограничениям, заданным в TensorShardingAttr .

Параметры:

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

EdgeNodeType

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

Случаи:

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

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

Случаи: