Definiciones de OpInterface

CollectiveOpInterface (Sdy_CollectiveOpInterface)

Es la interfaz para todas las operaciones colectivas. Encapsula los métodos get/set comunes para el atributo outSharding.

Restricciones:

  • El operando debe tener un fragmento o allowMissingInputSharding() muestra verdadero.
  • out_sharding es válido en relación con el tipo correspondiente.
  • El fragmentación del operando y del resultado debe tener la misma malla si allowDifferentMeshes() muestra un valor falso.
  • Mismo rango para el operador y el fragmentación de resultados.

Métodos:

getOutSharding

::mlir::sdy::TensorShardingAttr getOutSharding();

Devuelve la fragmentación del tensor de salida de la operación colectiva.

setOutShardingAttr

void setOutShardingAttr(::mlir::sdy::TensorShardingAttr sharding);

Establece la fragmentación del tensor de salida de la operación colectiva.

getTensor

::mlir::TypedValue<::mlir::TensorType> getTensor();

Obtén el operando de tensor de la operación colectiva.

getType

::mlir::Type getType();

Obtén el tipo del resultado de la operación colectiva.

allowDifferentMeshes

bool allowDifferentMeshes();

Indica si la operación colectiva permite que el fragmentación de entrada y salida tenga mallas diferentes.

allowMissingInputSharding

bool allowMissingInputSharding();

Indica si la operación colectiva permite que la entrada no tenga fragmentación, es decir, que se replique de forma implícita y completa.

ShardableDataFlowOpInterface (Sdy_ShardableDataFlowOpInterface)

Una interfaz de operación que permite que shardy propague particiones a través de los bordes de flujo de datos de las operaciones que extienden esta interfaz.

Un borde de flujo de datos de alguna operación X define un puente entre un conjunto de fuentes (cada una es un operando de X o un operando del terminador de bloque de X) y un conjunto de destinos (cada uno es un resultado de X o un argumento de bloque de X), de modo que todas las fuentes y los destinos se deben dividir de la misma manera. Una operación puede tener varios bordes de flujo de datos que son ortogonales entre sí.

Un propietario es un destino especificado por el usuario del borde de flujo de datos que usa la propagación de fragmentos. El usuario puede elegirlo de forma arbitraria, pero debe ser estático.

Por ejemplo:

  y_1, ..., y_n = custom_op (x_1, ..., x_n)
                  ((body_arg_1,..., body_arg_n) {
                    ...
                    return return_value_1, ..., return_value_n
                  })

Este custom_op tiene dos tipos de aristas de flujo de datos, n aristas entre return_value_i (fuentes) y y_i (destinos) y n aristas entre x_i(fuentes) y body_arg_i(destinos). En este caso, los propietarios de los bordes son los mismos que los destinos.

Este es un ejemplo de una operación con varios objetivos:

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

Esta operación while tiene n aristas de flujo de datos, y la arista de flujo de datos en el paso i está entre las fuentes x_i, return_value_i y los destinos y_i, pred_arg_i y body_arg_i.

Métodos:

getBlockArgumentEdgeOwnerShardings

mlir::SmallVector<mlir::sdy::TensorShardingAttr> getBlockArgumentEdgeOwnerShardings();

Muestra los particionamientos de todos los propietarios de aristas de flujo de datos de argumentos de bloque.

setBlockArgumentEdgeOwnerShardings

void setBlockArgumentEdgeOwnerShardings(mlir::ArrayRef<mlir::sdy::TensorShardingAttr> shardings);

Establece shardings de todos los propietarios de aristas de argumentos de bloque.

getOpResultEdgeOwnerShardings

mlir::SmallVector<mlir::sdy::TensorShardingAttr> getOpResultEdgeOwnerShardings();

Muestra los particionamientos de todos los propietarios de aristas del flujo de datos de resultados de la operación.

setOpResultEdgeOwnerShardings

void setOpResultEdgeOwnerShardings(mlir::ArrayRef<mlir::sdy::TensorShardingAttr> shardings);

Establece shardings de todos los propietarios de aristas de resultados de la operación.

transformTargetSharding

mlir::sdy::TensorShardingAttr transformTargetSharding(mlir::Value target, mlir::sdy::TensorShardingAttr sharding, mlir::sdy::DataFlowShardingTransformType transformType);

Transforma el sharding del objetivo según transformType.

Consulta DataFlowShardingTransformType para obtener más información.

getBlockArgumentEdgeOwners

mlir::ArrayRef<mlir::BlockArgument> getBlockArgumentEdgeOwners();

Obtiene todos los propietarios de aristas de argumentos de bloque.

getOpResultEdgeOwners

mlir::ResultRange getOpResultEdgeOwners();

Obtiene todos los propietarios de los bordes de resultados de la operación.

getEdgeSources

mlir::SmallVector<mlir::OpOperand*> getEdgeSources(mlir::Value owner);

Obtiene las fuentes de aristas de flujo de datos según el borde owner.

getEdgeOwnerFromTarget

mlir::Value getEdgeOwnerFromTarget(mlir::Value target);

Obtiene el target propietario de un borde de flujo de datos dado un target que puede ser o no el propietario.

getEdgeOwnerFromSource

mlir::Value getEdgeOwnerFromSource(mlir::OpOperand&source);

Obtiene el objetivo del propietario de un borde de flujo de datos dado un source.

getNonEdgeOwnerTargets

mlir::SmallVector<mlir::Value> getNonEdgeOwnerTargets(mlir::Value owner);

Obtiene los destinos que no son del propietario de un borde de flujo de datos dado el borde owner.

ShardingRuleOpInterface (Sdy_ShardingRuleOpInterface)

Es una interfaz de operación que le permite definir su propia regla de fragmentación. Una regla de fragmentación especifica cómo se puede particionar una operación según variosas propiedades en la operación: cualquier atributo, la forma de los operandos, la forma de los resultados, etcétera. Consulta OpShardingRuleAttr para obtener más detalles.

Métodos:

getShardingRule

mlir::sdy::OpShardingRuleAttr getShardingRule();

Muestra la regla de fragmentación de la operación.