Dialeto 'mpmd'

Operações

mpmd.assign (mpmd::AssignOp)

Operação de atribuição

Sintaxe:

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

Atribui um tensor local a uma malha como totalmente replicado nela.

Essa é uma operação temporária introduzida ao reduzir as operações do JAX para mover de tipos locais para tipos de malha. Essas operações serão eliminadas durante a importação, quando as entradas e os resultados da operação de função se tornarem tensores de malha.

O nome da malha do tipo de resultado precisa corresponder a uma malha na topologia, e o tipo global precisa ser idêntico ao tipo de operando.

A origem da operação de atribuição é a origem da malha, por exemplo, named_computation, inferência de malha etc.

Características: AlwaysSpeculatableImplTrait, HasParent<::mlir::func::FuncOp, ForOp>

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Efeitos: MemoryEffects::Effect{}

Atributos:

Operandos:

Resultados:

mpmd.broadcast (mpmd::BroadcastOp)

Operação de transmissão

Sintaxe:

operation ::= `mpmd.broadcast` attr-dict $tensor `:` type($tensor)

Permite que um tensor seja transferido (ou replicado) em qualquer malha em que ele seja usado. Sempre que transferida, a origem da transferência é o local atual do operando.

Características: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Efeitos: MemoryEffects::Effect{}

Operandos:

Resultados:

mpmd.call (mpmd::CallOp)

Função de chamada específica do MPMD

Sintaxe:

operation ::= `mpmd.call` $callee `(` $tensors `)` attr-dict `:` functional-type(operands, results)

Uma operação de chamada de função. Útil para encapsular o corpo de loops em declarações de função para reduzir o tamanho do código, por exemplo.

Interfaces: ArgAndResultAttrsOpInterface, CallOpInterface, SymbolUserOpInterface

Atributos:

Operandos:

Resultados:

mpmd.for (mpmd::ForOp)

Para operador

Retorna o resultado da execução de uma função de corpo para um número fixo de iterações, com o índice de iteração disponível no corpo.

Um fator de abertura opcional, que precisa dividir o número de iterações, pode ser especificado para abrir o corpo da operação por esse fator. Por exemplo, para o fator de abertura N, o corpo é replicado para criar N cópias, e o número de iterações é reduzido por um fator de 1/N. Cada cópia, exceto a primeira, usa os resultados da cópia anterior em vez dos argumentos de bloco, e o índice de iteração é multiplicado pelo fator de desenrolamento e incrementado após cada cópia.

Um operador "for" pode aceitar e retornar qualquer tipo, mas o TypeID deles precisa ser o mesmo. Por exemplo, todos os tipos de tensor ou todos os tipos de malha MPMD etc. Isso permite usar a operação em vários níveis, compartilhando implementação e transformações.

Traços: HLO_PairwiseSameOperandAndResultType, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Interfaces: ConditionallySpeculatable, LoopLikeOpInterface, OpAsmOpInterface, ShardableDataFlowOpInterface

Atributos:

Operandos:

Resultados:

mpmd.fragment (mpmd::FragmentOp)

Operação de fragmento

Atribui uma computação, ou seja, um bloco de operações, a uma malha específica em uma topologia MPMD, que deve ser executada como um fragmento de programa SPMD individual.

O fragmento usa e retorna apenas tensores de malha atribuídos à mesma malha que o fragmento.

O nome da malha do fragmento precisa corresponder a uma malha na topologia.

O fragmento inclui uma lista de origens, ou seja, metadados com informações sobre as named_computations originais que formaram esse fragmento, e um staged_id definido iff se for um fragmento definido pelo usuário, ou seja, se tiver uma lista não vazia de origens. O in_shardings opcional especifica o sharding dos argumentos de bloco de um fragmento, que correspondem aos operandos. O out_shardings opcional especifica os fragmentos dos resultados.

A região do fragmento não pode ter variáveis livres, e o tipo de cada argumento de bloco e valor retornado na região é o tipo de tensor global do tensor de malha correspondente.

Características: HasParent<::mlir::func::FuncOp, ForOp>, IsolatedFromAbove, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Interfaces: ConditionallySpeculatable, ShardableDataFlowOpInterface

Atributos:

Operandos:

Resultados:

mpmd.fragment_call (mpmd::FragmentCallOp)

Operação de chamada de fragmento

Representa uma chamada para uma função que contém um corpo de fragmento MPMD, ou seja, um cálculo atribuído a uma malha específica em uma topologia MPMD, que deve ser executado como um fragmento de programa SPMD individual.

O nome da malha do fragmento precisa corresponder a uma malha na topologia da função de inclusão, e o formato da malha precisa corresponder ao do destinatário da chamada.

A origem especifica as computações nomeadas pelo usuário que contribuíram para essa chamada de fragmento, por exemplo, por mesclagem.

Os tipos de entrada e resultado da função do recebedor da chamada precisam ser os tipos de tensor local dos tensores de malha correspondentes dos operandos e resultados dessa op, respectivamente.

Exemplo:

%2 = mpmd.fragment_call<mesh="m1",> origin=[] @my_fragment(%0, %1) :
  (m<esh>_tensor..., m<esh>_te>nsor...) - m<esh>_tensor...

Características: HasParent<::mlir::func::FuncOp>, MemRefsNormalizable

Interfaces: ArgAndResultAttrsOpInterface, CallOpInterface, SymbolUserOpInterface

Atributos:

Operandos:

Resultados:

mpmd.named_computation (mpmd::NamedComputationOp)

Operação de escopo nomeado

Agrupa um cálculo, ou seja, um bloco de operações, e atribui um nome e uma contagem de transposição usando o atributo UserOrigin. Essa NamedComputation pode ser usada para atribuir uma malha à computação em MPMD ou para otimizações.

A contagem de transposição (padrão=0) indica se a computação nomeada foi produzida por um determinado número de transformações de transposição JAX AD.

A região da operação não pode ter variáveis livres, e o tipo de cada argumento de bloco e valor retornado na região precisa ser o mesmo que o tipo das entradas e o tipo de retorno da operação.

Traços: IsolatedFromAbove, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Interfaces: ConditionallySpeculatable

Atributos:

Operandos:

Resultados:

mpmd.named_tensor (mpmd::NamedTensorOp)

Atribuir um tensor a uma malha

Sintaxe:

operation ::= `mpmd.named_tensor` $tensor `name````=```$name attr-dict `:` type($result)

Uma operação de identidade que associa o resultado do tensor a um determinado nome. Esse NamedTensor pode ser usado para atribuir uma malha ao tensor em MPMD.

Características: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Efeitos: MemoryEffects::Effect{}

Atributos:

Operandos:

Resultados:

mpmd.reduce (mpmd::ReduceOp)

Operação de redução entre malhas

Sintaxe:

operation ::= `mpmd.reduce` `` $reduction attr-dict $tensors `:` functional-type(operands, results)

Permite que um tensor seja reduzido em diferentes malhas e transmitido para onde precisar ser usado.

Características: AlwaysSpeculatableImplTrait, SameOperandsAndResultType

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Efeitos: MemoryEffects::Effect{}

Atributos:

Operandos:

Resultados:

mpmd.return (mpmd::ReturnOp)

A operação mpmd.return encerra as regiões anexadas a operações baseadas em regiões do mpmd. Ele é variádico: recebe como argumentos uma lista de valores cujos tipos podem ser quaisquer (mas do mesmo tipo, por exemplo, AnyTensor) e, portanto, podem ser reutilizados em vários níveis da pilha de IR do MPMD.

Sintaxe:

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

Características: AlwaysSpeculatableImplTrait, Terminator

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Efeitos: MemoryEffects::Effect{}

Operandos:

mpmd.transfer (mpmd::TransferOp)

Operação de transferência

Sintaxe:

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

Transfere um tensor distribuído de uma malha para outra.

Os nomes das malhas dos tipos de operando e resultado precisam corresponder às malhas na topologia, e os tipos globais precisam ser idênticos.

Características: AlwaysSpeculatableImplTrait, HasParent<::mlir::func::FuncOp>

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), ShardingRuleOpInterface

Efeitos: MemoryEffects::Effect{}

Operandos:

Resultados:

mpmd.unassign (mpmd::UnassignOp)

Operação de remoção da atribuição

Sintaxe:

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

Cancela a atribuição de um tensor totalmente replicado de uma malha.

Essa é uma operação temporária introduzida ao reduzir as operações do JAX para mover de tipos locais para tipos de malha. Essas operações serão eliminadas durante a importação, quando as entradas e os resultados da operação de função se tornarem tensores de malha.

O nome da malha do tipo de operando precisa corresponder a uma malha na topologia, e o tipo global precisa ser idêntico ao tipo de resultado.

Características: AlwaysSpeculatableImplTrait, HasParent<::mlir::func::FuncOp, ForOp>, InferTensorType

Interfaces: ConditionallySpeculatable, InferShapedTypeOpInterface, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Efeitos: MemoryEffects::Effect{}

Atributos:

Operandos:

Resultados:

Atributos

MeshWithOriginsAttr

Mesh com suas origens.

Sintaxe:

#mpmd.mesh_with_origins<
  ::llvm::StringRef,   # mesh_name
  ::llvm::ArrayRef<OriginAttr>   # origins
>

Parâmetros:

MeshesWithOriginsAttr

Uma lista de malhas com as origens delas.

Sintaxe:

#mpmd.meshes_with_origins<
  ::llvm::ArrayRef<MeshWithOriginsAttr>   # value
>

Parâmetros:

NamedMeshAttr

Um par com um nome e uma malha.

Sintaxe:

#mpmd.named_mesh<
  ::llvm::StringRef,   # name
  sdy::MeshAttr   # mesh
>

Parâmetros:

OriginAttr

Origem da atribuição de malha.

Sintaxe:

#mpmd.origin<
  ::llvm::StringRef   # origin_label
>

A origem de uma atribuição de malha.

origin_label é um rótulo legível para a origem. Ele é usado para fins de depuração.

Parâmetros:

ReductionAttr

Indica uma redução.

Sintaxe:

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

Parâmetros:

TopologyAttr

Topologia de malhas nomeadas.

Sintaxe:

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

Parâmetros:

UserOriginAttr

Origem do cálculo especificado pelo usuário.

Sintaxe:

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

Parâmetros:

Tipos

MeshTensorType

Tipo de tensor de malha

Atribui um RankedTensorType a uma malha SPMD específica na topologia MPMD de malhas do programa. O tipo contém um sharding opcional que especifica como o tensor é fragmentado em relação à malha SPMD. Se o sharding não estiver presente, o tensor será totalmente replicado.

Parâmetros:

Tipos enumerados

ReductionType

Atributo de tipo de redução

Casos:

EdgeNodeType

Enumeração do tipo de nó de borda

Casos:

PropagationDirection

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

Casos:

ReductionOp

Enumeração da operação de redução

Casos: