Dialecto 'sdy'

O dialeto Shardy (SDY) define uma representação de fragmentação de tensor com base em eixo, além de outros componentes da API para anexar fragmentações aos tensores.

Operações

sdy.constant (sdy::ConstantOp)

Operação constante

Produz um tensor output de uma value constante.

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

Exemplo:

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

Características: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Efeitos: MemoryEffects::Effect{}

Atributos:

Resultados:

sdy.data_flow_edge (sdy::DataFlowEdgeOp)

Operação de borda do fluxo de dados.

Sintaxe:

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

Uma borda de fluxo de dados de alguma operação X define uma ponte entre um conjunto de origens (cada uma é um operando de X ou um operando do terminador de bloco de X) e um conjunto de destinos (cada um é um resultado de X ou um argumento de bloco de X), de modo que todas as origens e destinos precisam ser particionados da mesma maneira.

Uma operação pode ter várias arestas de fluxo de dados ortogonais entre si.

Exemplo:

  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
                  })

Essa operação while tem n bordas de fluxo de dados, e a i-ésima borda de fluxo de dados está entre as fontes x_i, return_value_i e os destinos y_i, pred_arg_i e body_arg_i.

Um sdy.data_flow_edge usa como entrada o destino raiz de uma aresta (pode ser qualquer um dos destinos, mas de preferência um resultado de operação em vez de um argumento de bloco), que não pode ter outros usos. Essa operação não é pura porque pode receber uma entrada que originalmente não tinha nenhum uso.

O sdy.data_flow_edge também mantém um fragmentação opcional para todos os destinos do edge, e essa fragmentação precisa ser atualizada em vez da fragmentação dos destinos (se puder ser anexada) durante a propagação. Isso é útil quando uma operação tem muitas arestas, porque é muito mais eficiente:

• se propagam separadamente em cada borda.
• Atualize o sharding de cada borda separadamente em vez de todas as metas de uma vez. Por exemplo, uma operação tem um único TensorShardingPerValueAttr imutável para shardings de resultados.
• Adicione cada borda à lista de trabalhos separadamente quando o sharding de uma origem tiver mudado.

A propagação vai propagar os particionamentos entre todas as origens e destinos de um sdy.data_flow_edge como se fosse uma operação regular com as origens como operandos e os destinos como resultados, além de uma identidade sdy.op_sharding_rule. Isso significa que a propagação para frente é de origens para destinos, e a propagação para trás é de destinos para origens.

Não permitimos que a entrada de um sdy.data_flow_edge seja definida por uma operação SdyDialect. Portanto, podemos presumir que ela é definida por uma operação com atributo sdy.sharding não registrado.

Características: SameOperandsAndResultType

Interfaces: InferTypeOpInterface

Atributos:

Operandos:

Resultados:

sdy.manual_computation (sdy::ManualComputationOp)

Operação de paralelismo em vários dispositivos com coletivos manuais

Sintaxe:

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)

Vá para uma região escrita em termos de código local por dispositivo com coletivos explícitos, em que as formas lógicas correspondem às formas de buffer físico por dispositivo, e os coletivos correspondem exatamente à comunicação física entre dispositivos.

O corpo é local wrt os manuais_eixos. A propagação vai ocorrer pelo corpo em qualquer eixo livre, ou seja, aqueles que não estão na lista manual_axes.

Características: IsolatedFromAbove, RecursiveMemoryEffects, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Atributos:

Operandos:

Resultados:

sdy.mesh (sdy::MeshOp)

Malha nomeada

Sintaxe:

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

Define uma nova malha com nome. Todas as malhas em um módulo precisam ter o mesmo número de dispositivos, exceto as malhas com um único device_id. A malha é uma operação Symbol que aparece na SymbolTable do módulo e pode ser referenciada pelo name.

Características: HasParent<ModuleOp>

Interfaces: Symbol

Atributos:

sdy.named_computation (sdy::NamedComputationOp)

Operação de cálculo nomeada

Sintaxe:

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)

Agrupa uma computação, ou seja, um bloco de operações, e dá um nome a ela. A propagação vai fluir para dentro/fora da região como se tudo estivesse inline.

Isso pode ser usado para processar a propagação por instruções de chamada para outras funções. Todos os usuários do Shardy precisam escrever um cartão de importação/exportação que converta as operações de chamada em operações sdy.named_computation, duplicando/copiando o corpo da função chamada para o corpo do named_computation.

O tipo de cada argumento de bloco e os valores retornados na região precisam ser idênticos ao tipo dos operandos e ao tipo de resultados da operação.

Exemplo:

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

Características: IsolatedFromAbove, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Interfaces: ConditionallySpeculatable, ShardableDataFlowOpInterface

Atributos:

Operandos:

Resultados:

sdy.propagation_barrier (sdy::PropagationBarrierOp)

Operação de barreira de propagação

Sintaxe:

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

Essa operação funciona como uma operação de identidade, gerando o mesmo valor que recebeu como entrada. Mas, em termos de propagação, isso só permite que a propagação flua por ele em uma determinada direção.

Isso impede que os fragmentos sejam propagados entre os usos do resultado da operação de barreira e do operando.

• FORWARD significa que os fragmentos só podem fluir do operando para o resultado.
• BACKWARD significa que as fragmentações só podem fluir do resultado para o operando.
• NONE significa que nenhum sharding pode ser propagado por essa operação.
• Não é possível especificar BOTH, porque essa operação seria redundante.

Características: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Efeitos: MemoryEffects::Effect{}

Atributos:

Operandos:

Resultados:

sdy.reshard (sdy::ReshardOp)

Refragmenta um tensor para uma fragmentação diferente

Sintaxe:

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

Refaz o sharding do tensor de entrada com o sharding especificado, que é diferente do sharding atual do tensor de entrada.

Tanto ShardingConstraintOp quanto ReshardOp anexam um particionamento a um tensor. O ciclo de vida deles é:

1. Antes da propagação de fragmentação, o ShardingConstraintOp é adicionado pelos usuários.
2. A propagação da fragmentação consome ShardingConstraintOp. Não há ShardingConstraintOp nos resultados da propagação de fragmentação. Em vez disso, ReshardOp pode ser adicionado, se necessário.
3. Um particionador converte uma ReshardOp em uma operação coletiva (ou uma operação de identidade). Não deve haver ReshardOp nos resultados do particionador.

// TODO(b/331680067). Adicione um padrão de canonização para remover operações // de refragmentação redundantes.

Características: AlwaysSpeculatableImplTrait, Elementwise, SameOperandsAndResultType

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Efeitos: MemoryEffects::Effect{}

Atributos:

Operandos:

Resultados:

sdy.return (sdy::ReturnOp)

A operação sdy.return encerra as regiões anexadas às operações baseadas em região sdy e a qualquer outra operação baseada em região Shardy. Ele é variado: usa como argumentos uma lista de valores com tipos que podem ser qualquer um (mas do mesmo tipo, por exemplo, AnyTensor) e, portanto, podem ser reutilizados em vários níveis da pilha de IR fragmentada.

Sintaxe:

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

Características: AlwaysSpeculatableImplTrait, Terminator

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Efeitos: MemoryEffects::Effect{}

Operandos:

sdy.sharding_constraint (sdy::ShardingConstraintOp)

Restrição de um tensor ao particionamento especificado

Sintaxe:

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

Anexa um sharding a um tensor intermediário (por exemplo, o resultado de uma matmul) para indicar como esse tensor ou um subconjunto de usos dele precisa ser fragmentado.

Se o particionamento tiver dimensões abertas e eixos sem restrição, isso significa que o tensor pode ser particionado de acordo com as dimensões abertas.

Essa operação pode:

• Não tem uso (solto), o que significa que o particionamento anexado é como o tensor de entrada precisa ser particionado.
• Ter usos: significa que o particionamento anexado é como os usos da operação de restrição de particionamento precisam ser particionados, enquanto outros usos do tensor de entrada podem ter um particionamento diferente. Se o tensor de entrada não tiver outros usos, o comportamento será o mesmo do caso sem usos.

Características: Elementwise, SameOperandsAndResultType

Interfaces: InferTypeOpInterface

Atributos:

Operandos:

Resultados:

sdy.sharding_group (sdy::ShardingGroupOp)

Operação de fragmentação de grupo

Sintaxe:

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

Essa operação fornece uma interface para atribuir tensores a grupos de fragmentação ( grupos de tensores que serão forçados a ter fragmentações idênticas). Durante a propagação, assim que um elemento de grupo é dividido, todos os outros membros são divididos da mesma forma. Essa operação usa o ID do grupo de argumentos e não retorna nenhum resultado, mas modifica a representação do grupo de fragmentação interna para adicionar o tensor de entrada ao grupo com o ID fornecido.

Atributos:

Operandos:

Atributos

AxisRefAttr

Referência a um eixo completo ou a um subeixo dividido

Sintaxe:

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

Parâmetros:

DimMappingAttr

Lista de índices de fatores para uma dimensão

Todos os índices de fatores precisam estar no intervalo [0, num_factors], e uma lista vazia indica que este é um mapeamento nulo (analisado/impresso com *), ou seja, a dimensão não é mapeada para nenhum fator.

Parâmetros:

DimensionShardingAttr

Fragmentação de dimensão

Lista de nomes de eixos para fragmentar uma dimensão de tensor do maior para o menor, um booleano que indica se a dimensão pode ser ainda mais fragmentada e um número inteiro opcional que indica a prioridade da fragmentação da dimensão, que será respeitada durante a propagação da fragmentação. As prioridades são originadas de anotações de fragmentação do usuário, e um valor menor indica uma prioridade mais alta. A prioridade mais alta é assumida quando está faltando a prioridade na anotação.

Parâmetros:

ManualAxesAttr

Sintaxe:

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

Parâmetros:

MeshAttr

Malha de eixos e uma lista de dispositivos

Sintaxe:

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

Uma malha é uma lista de eixos e uma lista opcional de IDs de dispositivos que especificam a ordem dos dispositivos.

Se a lista de eixos estiver vazia, a malha terá um eixo implícito sem nome de tamanho 1. Nesse caso, se uma lista de IDs de dispositivo não for fornecida, a lista implícita de IDs de dispositivos será [0]. Se uma lista de IDs de dispositivos for fornecida, ela precisará conter um único número inteiro de qualquer valor não negativo. Chamamos isso de caso de fragmentação máxima.

Para todos os casos de fragmentação não máxima, se uma lista de IDs de dispositivos for especificada, o produto dos tamanhos de eixo precisa corresponder ao número de dispositivos. Se uma lista de IDs de dispositivo não for especificada, a lista implícita de IDs de dispositivo será iota(product(axes)). Para simplificar, também não é permitido especificar uma lista de IDs de dispositivo igual a iota(product(axes)); nesse caso, uma lista de IDs de dispositivo não pode ser especificada.

Confira alguns exemplos de malhas:

• Uma malha vazia representa uma malha de marcador de posição que pode ser substituída durante a propagação: <[]>
• Uma malha com um eixo sem nome e um ID de dispositivo explícito, que normalmente é usado para representar o máximo de fragmentação: <[], device_ids=[3]>
• Uma malha com dois eixos e IDs de dispositivo implícitos iota(6): <["a"=2, "b"=3]>
• Uma malha com dois eixos e IDs de dispositivo explícitos que especificam a ordem do dispositivo: <["a"=3, "b"=2], device_ids=[0, 2, 4, 1, 3, 5]>

Parâmetros:

MeshAxisAttr

Eixo nomeado em uma malha

Sintaxe:

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

Parâmetros:

OpShardingRuleAttr

Especifica como uma operação pode ser particionada.

Sintaxe:

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

Uma regra de fragmentação especifica como uma operação pode ser particionada de acordo com várias propriedades na operação, como qualquer atributo, a forma dos operandos, a forma dos resultados etc. Por exemplo:

%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>

Permitimos fatores com tamanho 1, mesmo que não possam ser divididos, principalmente para completar, já que muitas operações, como operações pontuais, têm dimensões de tamanho um que correspondem a operandos e resultados.

is_custom_rule descreve se essa é uma regra definida por um usuário para uma operação stablehlo.custom_call. Como o particionador não sabe como particionar essas operações, o usuário precisa informar como fazer isso. Quando é uma regra personalizada, ela é sempre preservada/nunca removida. is_custom_rule só pode ser verdadeiro para operações stablehlo.custom_call.

Parâmetros:

SubAxisInfoAttr

Informações sobre como esse eixo secundário é derivado do eixo completo

Sintaxe:

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

Ao dividir um eixo completo em n subeixos, ele é remodelado em [k_1,...,k_n], e o i-ésimo subeixo pode ser expresso pelo produto de todos os tamanhos de eixo à esquerda de m=prod(k_1,...,k_(i-1)) (também conhecido como pré-tamanho) e tamanho k_i. Portanto, o atributo sub-axis-info contém esses dois números e é indicado da seguinte maneira: (m)k para o pré-tamanho m e o tamanho k.

Parâmetros:

TensorMappingAttr

Mapeamentos de fatores para cada dimensão de um tensor.

Sintaxe:

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

Parâmetros:

TensorShardingAttr

Fragmentação de tensor

Sintaxe:

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

Um fragmentação de tensor é vinculado a uma malha específica e só pode referenciar nomes de eixos dessa malha. Os shardings de dimensão informam, para cada dimensão do tensor, ao longo de quais eixos (ou subeixos) ele é dividido de maior para menor. Todos os outros eixos que não fragmentam uma dimensão são replicados de forma implícita ou explícita (se aparecerem na lista de eixos replicados).

A malha a que esse fragmentação está vinculada pode ser especificada por um nome de símbolo, fazendo referência a um símbolo MeshOp correspondente ou a um MeshAttr inline.

Parâmetros:

TensorShardingPerValueAttr

Fragmentação de tensor por operando/resultado de uma operação

Sintaxe:

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

Parâmetros:

Enums

PropagationDirection

Enumeração de direção de propagação

Casos: