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

Диалект Шарди (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 представляет собой список списков осей. Внешний список охватывает измерения тензора. Каждый внутренний список определяет оси, вдоль которых должна быть выполнена отдельная сборка для соответствующего измерения. Он будет применен к шардингу операнда ( 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 .
• Шардинг операнда и 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)

Разрезает фрагменты тензора вдоль осей, указанных в 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 выходного тензора.

К шардингу операнда ( 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
                  })

При этом у op имеется 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.reshard (sdy::ReshardOp)

Переносит тензор в другой шардинг

Синтаксис:

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

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

ShardingConstraintOp и ReshardOp прикрепляют шардинг к тензору. Срок их жизни составляет:

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

// TODO(b/331680067). Добавьте шаблон канонизации для удаления избыточных // операций перераспределения.

Черты: AlwaysSpeculatableImplTrait , SameOperandsAndResultType

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

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

Атрибуты:

Операнды:

Результаты:

sdy.return (sdy::ReturnOp)

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

Синтаксис:

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

Черты: AlwaysSpeculatableImplTrait , Terminator

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

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

Операнды:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

Ограничивает тензор указанным сегментированием

Синтаксис:

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

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

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

Эта операция может:

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

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

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

Если список осей пуст, сетка имеет неявную безымянную ось размера 1. В этом случае, если список идентификаторов устройств не указан, неявный список идентификаторов устройств равен [0]; если список идентификаторов устройств указан, он должен содержать одно целое число любого неотрицательного значения. Мы называем этот случай максимальным шардингом.

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

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

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

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

• Элементы 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, даже если их нельзя сегментировать. Это сделано главным образом для полноты, поскольку многие операции, например, точечные операции, имеют измерения размером 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 ].
  • Отсутствие дублирующихся факторных индексов внутри каждой группы и между группами.

Параметры:

РаспространениеКраяАттр

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

Синтаксис:

#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 к тензору вдоль нередуцированных осей сделает тензор реплицированным вдоль этих осей. Однако тензор с нередуцированными осями не обязательно должен быть полностью редуцирован немедленно, он может оставаться не редуцированным при передаче в линейные операции, такие как stablehlo.add (при условии, что обе левые и правые части не редуцированы) и полностью редуцированным впоследствии. Мы предполагаем, что тип редукции — сумма, другие редукции могут поддерживаться в будущем.

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

• Элементы в 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

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

Случаи:

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

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

Случаи: