OpInterface definitions

CollectiveOpInterface (Sdy_CollectiveOpInterface)

Interface for all collective ops. Encapsulates common get/set for outSharding attribute.

Constraints:

  • Operand must have a sharding or allowMissingInputSharding() returns true.
  • out_sharding is valid w.r.t the corresponding type.
  • Operand and result sharding must have the same mesh if allowDifferentMeshes() returns false.
  • Same rank for the operand and result sharding.

Methods:

getOutSharding

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

Returns the output tensor sharding of the collective op.

setOutShardingAttr

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

Sets the output tensor sharding of the collective op.

getTensor

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

Get the tensor operand of the collective op.

getType

::mlir::Type getType();

Get the type of the collective op result.

allowDifferentMeshes

bool allowDifferentMeshes();

Indicated whether the collective op allows the input and output sharding to have different meshes.

allowMissingInputSharding

bool allowMissingInputSharding();

Indicated whether the collective op allows the input to have no sharding, i.e, implicitly fully replicated.

ShardableDataFlowOpInterface (Sdy_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::SmallVector<mlir::sdy::TensorShardingAttr> getBlockArgumentEdgeOwnerShardings();

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

setBlockArgumentEdgeOwnerShardings

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

Sets shardings of all block argument edge owners.

getOpResultEdgeOwnerShardings

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

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

setOpResultEdgeOwnerShardings

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

Sets shardings of all op result edge owners.

transformTargetSharding

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

Transforms the sharding of the target depending on transformType

See DataFlowShardingTransformType for more information.

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::OpOperand*> getEdgeSources(mlir::Value owner);

Gets the data flow edge sources given the edge owner.

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.

shouldKeepEdgeOwnerShardingsDivisible

bool shouldKeepEdgeOwnerShardingsDivisible();

Returns true if the sharding of the edge owners divides the dimension size, to avoid the need for padding.

ShardingRuleOpInterface (Sdy_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.

shouldKeepOutputShardingsDivisible

bool shouldKeepOutputShardingsDivisible();

Returns true if the output sharding divides the dimension size, to avoid the need for padding.