Dialecto "mpmd"

Operaciones

mpmd.assign (mpmd::AssignOp)

Operación de asignación

Sintaxis:

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

Asigna un tensor local a una malla como completamente replicado dentro de esa malla.

Esta es una operación temporal que se introduce cuando se reducen las operaciones de JAX para pasar de tipos locales a tipos de malla. Estas operaciones se eliminarán durante la importación, cuando las entradas y los resultados de la operación func se conviertan en tensores de malla.

El nombre de la malla del tipo de resultado debe corresponder a una malla en la topología, y su tipo global debe ser idéntico al tipo de operando.

El origen de la operación de asignación es el origen de la malla, p.ej., named_computation, inferencia de malla, etcétera.

Rasgos: AlwaysSpeculatableImplTrait y HasParent<::mlir::func::FuncOp, ForOp>

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Efectos: MemoryEffects::Effect{}

Atributos:

Operandos:

Resultados:

mpmd.broadcast (mpmd::BroadcastOp)

Operación de transmisión

Sintaxis:

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

Permite que un tensor se transfiera (o replique) en cualquier malla en la que se use. Cada vez que se transfiere, el origen de la transferencia es la ubicación actual del operando.

Rasgos: AlwaysSpeculatableImplTrait y SameOperandsAndResultType

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Efectos: MemoryEffects::Effect{}

Operandos:

Resultados:

mpmd.call (mpmd::CallOp)

Función de llamada específica de MPMD

Sintaxis:

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

Operación de llamada a función. Es útil para envolver el cuerpo de los bucles en declaraciones de funciones para reducir el tamaño del código, por ejemplo.

Interfaces: ArgAndResultAttrsOpInterface, CallOpInterface, SymbolUserOpInterface

Atributos:

Operandos:

Resultados:

mpmd.for (mpmd::ForOp)

Para el operador

Devuelve el resultado de ejecutar una función de cuerpo para una cantidad fija de iteraciones, con el índice de iteración disponible en el cuerpo.

Se puede especificar un factor de desenrollado opcional, que debe dividir la cantidad de iteraciones, para desenrollar el cuerpo de la operación según ese factor, es decir, para un factor de desenrollado N, el cuerpo se replica para crear N copias y la cantidad de iteraciones se reduce en un factor de 1/N. Cada copia, excepto la primera, usa los resultados de la copia anterior en lugar de los argumentos del bloque, y el índice de iteración se multiplica por el factor de desenrollado y se incrementa después de cada copia.

Un operador for puede aceptar y devolver cualquier tipo, pero el TypeID de estos debe ser el mismo, p.ej., todos los tipos de tensores o todos los tipos de malla de MPMD, etc. Esto nos permite usar el operador en varios niveles, compartir la implementación y las transformaciones.

Rasgos: HLO_PairwiseSameOperandAndResultType, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Interfaces: ConditionallySpeculatable, LoopLikeOpInterface, OpAsmOpInterface y ShardableDataFlowOpInterface

Atributos:

Operandos:

Resultados:

mpmd.fragment (mpmd::FragmentOp)

Operación de fragmento

Asigna un cálculo, es decir, un bloque de operaciones, a una malla específica en una topología de MPMD, que se ejecutará como un fragmento de programa SPMD individual.

El fragmento solo toma y devuelve tensores de malla que se asignan a la misma malla que el fragmento.

El nombre de la malla del fragmento debe corresponder a una malla en la topología.

El fragmento incluye una lista de orígenes, es decir, metadatos con información sobre las named_computations originales que formaron este fragmento, y un staged_id definido iff si es un fragmento definido por el usuario, es decir, si tiene una lista de orígenes no vacía. El argumento opcional in_shardings especifica la fragmentación de los argumentos de bloque de un fragmento, que corresponden a los operandos. El argumento opcional out_shardings especifica la fragmentación de los resultados.

La región del fragmento no debe tener variables libres, y el tipo de cada argumento de bloque y valor devuelto en la región es el tipo de tensor global del tensor de malla correspondiente.

Rasgos: HasParent<::mlir::func::FuncOp, ForOp>, IsolatedFromAbove, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Interfaces: ConditionallySpeculatable, ShardableDataFlowOpInterface

Atributos:

Operandos:

Resultados:

mpmd.fragment_call (mpmd::FragmentCallOp)

Operación de llamada de fragmento

Representa una llamada a una función que contiene un fragmento del cuerpo de MPMD, es decir, un cálculo asignado a una malla específica en una topología de MPMD, que se pretende ejecutar como un fragmento de programa SPMD individual.

El nombre de la malla del fragmento debe corresponder a una malla en la topología de la función envolvente, y esa forma de malla debe coincidir con la de la función llamada.

El origen especifica los cálculos con nombre del usuario que contribuyeron a esta llamada de fragmento, p.ej., a través de la combinación.

Los tipos de entrada y resultado de la función del destinatario deben ser los tipos de tensor locales de los tensores de malla correspondientes de los operandos y resultados de esta op, respectivamente.

Ejemplo:

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

Rasgos: HasParent<::mlir::func::FuncOp> y MemRefsNormalizable

Interfaces: ArgAndResultAttrsOpInterface, CallOpInterface, SymbolUserOpInterface

Atributos:

Operandos:

Resultados:

mpmd.named_computation (mpmd::NamedComputationOp)

Operación de alcance con nombre

Agrupa un cálculo, es decir, un bloque de operaciones, y le asigna un nombre y un recuento de la transposición a través del atributo UserOrigin. Este NamedComputation se puede usar para asignar una malla al cálculo en MPMD o para optimizaciones.

El recuento de transposición (predeterminado=0) indica si el cálculo con nombre se produjo a partir de una cierta cantidad de transformaciones de transposición de AD de JAX.

La región de la op no debe tener variables libres, y el tipo de cada argumento de bloque y los valores devueltos en la región deben ser los mismos que el tipo de las entradas y el tipo de datos que se devuelve de la op.

Rasgos: IsolatedFromAbove, RecursiveMemoryEffects, RecursivelySpeculatableImplTrait, SingleBlockImplicitTerminator<ReturnOp>, SingleBlock

Interfaces: ConditionallySpeculatable

Atributos:

Operandos:

Resultados:

mpmd.named_tensor (mpmd::NamedTensorOp)

Cómo asignar un tensor a una malla

Sintaxis:

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

Es una operación de identidad que asocia el resultado del tensor con un nombre determinado. Este NamedTensor se puede usar para asignar una malla al tensor en MPMD.

Rasgos: AlwaysSpeculatableImplTrait y SameOperandsAndResultType

Interfaces: ConditionallySpeculatable, InferTypeOpInterface, NoMemoryEffect (MemoryEffectOpInterface)

Efectos: MemoryEffects::Effect{}

Atributos:

Operandos:

Resultados:

mpmd.reduce (mpmd::ReduceOp)

Operación de reducción entre mallas

Sintaxis:

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

Permite reducir un tensor en diferentes mallas y, luego, transmitirlo a donde sea necesario usarlo.

Rasgos: AlwaysSpeculatableImplTrait y SameOperandsAndResultType

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Efectos: MemoryEffects::Effect{}

Atributos:

Operandos:

Resultados:

mpmd.return (mpmd::ReturnOp)

La operación mpmd.return finaliza las regiones adjuntas a las operaciones basadas en regiones de mpmd. Es variádica: toma como argumentos una lista de valores cuyos tipos pueden ser cualquiera (pero del mismo tipo, p.ej., AnyTensor) y, por lo tanto, se puede reutilizar en varios niveles de la pila de RI de MPMD.

Sintaxis:

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

Rasgos: AlwaysSpeculatableImplTrait y Terminator

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Efectos: MemoryEffects::Effect{}

Operandos:

mpmd.transfer (mpmd::TransferOp)

Operación de transferencia

Sintaxis:

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

Transfiere un tensor distribuido de una malla a otra.

Los nombres de malla de los tipos de operando y resultado deben corresponder a mallas en la topología, y sus tipos globales deben ser idénticos.

Rasgos: AlwaysSpeculatableImplTrait y HasParent<::mlir::func::FuncOp>

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), ShardingRuleOpInterface

Efectos: MemoryEffects::Effect{}

Operandos:

Resultados:

mpmd.unassign (mpmd::UnassignOp)

Operación de anulación de asignación

Sintaxis:

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

Anula la asignación de un tensor replicado por completo de una malla.

Esta es una operación temporal que se introduce cuando se reducen las operaciones de JAX para pasar de tipos locales a tipos de malla. Estas operaciones se eliminarán durante la importación, cuando las entradas y los resultados de la operación func se conviertan en tensores de malla.

El nombre de la malla del tipo de operando debe corresponder a una malla en la topología, y su tipo global debe ser idéntico al tipo de resultado.

Rasgos: AlwaysSpeculatableImplTrait, HasParent<::mlir::func::FuncOp, ForOp>, InferTensorType

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

Efectos: MemoryEffects::Effect{}

Atributos:

Operandos:

Resultados:

Atributos

MeshWithOriginsAttr

Se adapta a sus orígenes.

Sintaxis:

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

Parámetros:

MeshesWithOriginsAttr

Una lista de mallas con sus orígenes.

Sintaxis:

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

Parámetros:

NamedMeshAttr

Es un par con un nombre y una malla.

Sintaxis:

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

Parámetros:

OriginAttr

Origen de la asignación de malla.

Sintaxis:

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

Es el origen de una asignación de malla.

origin_label es una etiqueta legible para el origen. Está diseñado para usarse con fines de depuración.

Parámetros:

ReductionAttr

Denota una reducción.

Sintaxis:

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

Parámetros:

TopologyAttr

Topología de mallas con nombre.

Sintaxis:

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

Parámetros:

UserOriginAttr

Origen del cálculo especificado por el usuario.

Sintaxis:

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

Parámetros:

Tipos

MeshTensorType

Tipo de tensor de malla

Asigna un RankedTensorType a una malla SPMD específica en la topología MPMD de mallas del programa. El tipo contiene un sharding opcional que especifica cómo se fragmenta el tensor con respecto a la malla de SPMD. Si no hay fragmentación, el tensor se replica por completo.

Parámetros:

Enums

ReductionType

Atributo de tipo de reducción

Casos:

EdgeNodeType

Enumeración del tipo de nodo perimetral

Casos:

PropagationDirection

Enum de dirección de propagación

Casos:

ReductionOp

Enum de operación de reducción

Casos: