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

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

Операции

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)

Край потока данных op.

Синтаксис:

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

Граница потока данных некоторой операции X определяет мост между набором источников (каждый является либо операндом 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

Атрибуты:

Операнды:

Результаты:

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.

Признаки: IsolatedFromAbove , RecursiveMemoryEffects , SingleBlockImplicitTerminator<ReturnOp> , SingleBlock

Атрибуты:

Операнды:

Результаты:

sdy.mesh (sdy::MeshOp)

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

Синтаксис:

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

Определяет новую именованную сетку. Все сетки в модуле должны иметь одинаковое количество устройств (за исключением сеток с одним идентификатором устройства). Сетка — это операция 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 , ShardableDataFlowOpInterface .

Атрибуты:

Операнды:

Результаты:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

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

Синтаксис:

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

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

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

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

Признаки: AlwaysSpeculatableImplTrait , Elementwise , SameOperandsAndResultType

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

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

Атрибуты:

Операнды:

Результаты:

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). Добавьте шаблон канонизации для удаления избыточных операций // reshard.

Признаки: AlwaysSpeculatableImplTrait , Elementwise , SameOperandsAndResultType

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

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

Атрибуты:

Операнды:

Результаты:

sdy.return (sdy::ReturnOp)

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

Синтаксис:

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), чтобы указать, что именно так должен быть сегментирован этот тензор или подмножество его применений.

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

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

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

Черты: Elementwise , SameOperandsAndResultType

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

Атрибуты:

Операнды:

Результаты:

sdy.sharding_group (sdy::ShardingGroupOp)

Операция группы шардинга

Синтаксис:

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

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

Атрибуты:

Операнды:

Атрибуты

ОсьRefAttr

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

Синтаксис:

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

Параметры:

Диммаппингаттр

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

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

Параметры:

Размер ШардингАттр

Шардинг измерений

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

Параметры:

МануалАксесАттр

Синтаксис:

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

Параметры:

Мешаттр

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

Синтаксис:

#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]>

Параметры:

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
  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, даже если они не могут быть сегментированы, это сделано главным образом для полноты, поскольку многие операции, такие как точечные операции, имеют измерения размера один, которые соответствуют операндам и результатам.

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

Параметры:

Субаксисинфоаттр

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

Синтаксис:

#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.

Параметры:

TensorMappingAttr

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

Синтаксис:

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

Параметры:

TensorShardingAttr

Тензорное шардинг

Синтаксис:

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

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

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

Параметры:

TensorShardingPerValueAttr

Тензорное сегментирование для каждого операнда/результата операции

Синтаксис:

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

Параметры:

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

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

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

Случаи: