Definiciones de OpInterface

ShardableDataFlowOpInterface (ShardableDataFlowOpInterface)

Una interfaz operativa que permite que los fragmentos se propaguen a través de los bordes del 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
                  })

Esta 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::ArrayRef<mlir::sdy::TensorShardingAttr> getBlockArgumentEdgeOwnerShardings();

Muestra las fragmentaciones de todos los propietarios de perímetro de flujo de datos de argumentos de bloques.

setBlockArgumentEdgeOwnerSharding

void setBlockArgumentEdgeOwnerSharding(unsigned index, mlir::sdy::TensorShardingAttr sharding);

Establece el sharding del propietario del borde del argumento de bloque con el index determinado.

setBlockArgumentEdgeOwnerShardings

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

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

getOpResultEdgeOwnerShardings

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

Muestra las fragmentaciones de todos los propietarios del perímetro de flujo de datos de los resultados de la operación.

setOpResultEdgeOwnerSharding

void setOpResultEdgeOwnerSharding(unsigned index, mlir::sdy::TensorShardingAttr sharding);

Establece el objeto sharding del propietario del perímetro del resultado de la operación con el index determinado.

setOpResultEdgeOwnerShardings

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

Configura los shardings de todos los propietarios del perímetro de los resultados de la operació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::Value> getEdgeSources(mlir::Value target);

Obtiene las fuentes de aristas del flujo de datos según un valor target.

getEdgeOwnerFromTarget

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

Obtiene el propietario target de un perímetro de flujo de datos según una 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 propietarios de un borde de flujo de datos dado el borde owner.

ShardingRuleOpInterface (ShardingRuleOpInterface)

Una interfaz operativa que permite que la operación defina su propia regla de fragmentación. Una regla de fragmentación especifica cómo se puede particionar una operación según varias propiedades de 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.