OpInterface definitions

ShardableDataFlowOpInterface (ShardableDataFlowOpInterface)

An op interface that allows shardy to propagate shardings through data flow edges of ops that extend this interface.

A data flow edge of some op X defines a bridge between a set of sources (each is either an operand of X or an operand of X's block terminator) and a set of targets (each is either a result of X or a block argument of X), such that all sources and targets should be sharded in the same way. An op can have multiple data flow edges that are orthogonal to one another.

An owner is a user specified target of the data flow edge used by shardy's propagation. The user can choose it arbitrarily but it needs to be static.

For example:

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

This custom_op has two types for data flow edges, n edges each between return_value_i (sources) and y_i (targets) and n edges between x_i(sources) and body_arg_i(targets). In this case the edge owners are the same as the targets.

Here is an example of an op with multiple targets:

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

This while op has n data flow edges, the i-th data flow edges is between sources x_i, return_value_i and targets y_i, pred_arg_i, body_arg_i.

Methods:

getBlockArgumentEdgeOwnerShardings

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

Returns the shardings of all block argument data flow edge owners.

setBlockArgumentEdgeOwnerSharding

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

Sets the sharding of the block argument edge owner with the given index.

setBlockArgumentEdgeOwnerShardings

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

Sets shardings of all block argument edge owners.

getOpResultEdgeOwnerShardings

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

Returns the shardings of all op result data flow edge owners.

setOpResultEdgeOwnerSharding

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

Sets the sharding of the op result edge owner with the given index.

setOpResultEdgeOwnerShardings

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

Sets shardings of all op result edge owners.

getBlockArgumentEdgeOwners

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

Gets all block argument edge owners.

getOpResultEdgeOwners

mlir::ResultRange getOpResultEdgeOwners();

Gets all op result edge owners.

getEdgeSources

mlir::SmallVector<mlir::Value> getEdgeSources(mlir::Value target);

Gets the data flow edge sources given a target value.

getEdgeOwnerFromTarget

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

Gets the owner target of a data flow edge given a target that may or may not be the owner.

getEdgeOwnerFromSource

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

Gets the owner target of a data flow edge given a source.

getNonEdgeOwnerTargets

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

Gets the non-owner targets of a data flow edge given the edge owner.

ShardingRuleOpInterface (ShardingRuleOpInterface)

An op interface that allows the op to define its own sharding rule. A sharding rule specifies how an operation can be partitioned according to various properties on the op - any attributes, the shape of operands, the shape of the results, etc. See OpShardingRuleAttr for more details.

Methods:

getShardingRule

mlir::sdy::OpShardingRuleAttr getShardingRule();

Returns the sharding rule of the op.