ShardingPropagation.cpp File Reference

#include "mlir/Dialect/Mesh/Transforms/Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Mesh/IR/MeshDialect.h"
#include "mlir/Dialect/Mesh/IR/MeshOps.h"
#include "mlir/Dialect/Mesh/Interfaces/ShardingInterface.h"
#include "mlir/IR/Verifier.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Pass/Pass.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/iterator_range.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
#include <algorithm>
#include <vector>
#include "mlir/Dialect/Mesh/Transforms/Passes.h.inc"

Go to the source code of this file.

Classes
struct	ShardingPropagation

Namespaces
	mlir
	Include the generated interface declarations.
	mlir::mesh

Macros
#define	GEN_PASS_DEF_SHARDINGPROPAGATION
#define	DEBUG_TYPE   "sharding-propagation"
#define	DBGS()   (llvm::dbgs() << "[" DEBUG_TYPE << "]: ")

Enumerations
enum class	ReshardingRquirementKind { NO_RESHARDING = 0 , NO_RESHARDING_FOR_EXPLICIT_ANNOTATIONS , RESHARDING_FOR_EXPLICIT_ANNOTATIONS }

Functions
static SmallVector< std::vector< MeshSharding > >	getOrderedPossibleShardingAttrs (ArrayRef< MeshSharding > mustShardings, ArrayRef< MeshSharding > optionalShardings)
ReshardingRquirementKind	getReshardingRquirementKind (Operation *op, const std::vector< MeshSharding > &operandAndResultShardings)
static FailureOr< ShardingOption >	selectShardingOption (ShardingInterface shardingOp, ArrayRef< std::vector< MeshSharding >> possibleOperandShardingAttrs, ArrayRef< std::vector< MeshSharding >> possibleResultShardingAttrs)
static LogicalResult	visitOp (Operation *op, OpBuilder &builder)

Macro Definition Documentation

DBGS

#define DBGS

(

)

(llvm::dbgs() << "[" DEBUG_TYPE << "]: ")

Definition at line 34 of file ShardingPropagation.cpp.

DEBUG_TYPE

#define DEBUG_TYPE   "sharding-propagation"

Definition at line 33 of file ShardingPropagation.cpp.

GEN_PASS_DEF_SHARDINGPROPAGATION

#define GEN_PASS_DEF_SHARDINGPROPAGATION

Definition at line 28 of file ShardingPropagation.cpp.

Enumeration Type Documentation

ReshardingRquirementKind

enum ReshardingRquirementKind

strong

Enumerator
NO_RESHARDING
NO_RESHARDING_FOR_EXPLICIT_ANNOTATIONS
RESHARDING_FOR_EXPLICIT_ANNOTATIONS

Definition at line 39 of file ShardingPropagation.cpp.

Function Documentation

getOrderedPossibleShardingAttrs()

static SmallVector<std::vector<MeshSharding> > getOrderedPossibleShardingAttrs ( ArrayRef< MeshSharding >  mustShardings, ArrayRef< MeshSharding >  optionalShardings  )

static

Definition at line 112 of file ShardingPropagation.cpp.

Referenced by visitOp().

getReshardingRquirementKind()

ReshardingRquirementKind getReshardingRquirementKind	(	Operation *	op,
		const std::vector< MeshSharding > &	operandAndResultShardings
	)

Definition at line 159 of file ShardingPropagation.cpp.

References mlir::Operation::getOperands(), mlir::Operation::getResults(), NO_RESHARDING, NO_RESHARDING_FOR_EXPLICIT_ANNOTATIONS, and RESHARDING_FOR_EXPLICIT_ANNOTATIONS.

Referenced by selectShardingOption().

selectShardingOption()

static FailureOr<ShardingOption> selectShardingOption ( ShardingInterface  shardingOp, ArrayRef< std::vector< MeshSharding >>  possibleOperandShardingAttrs, ArrayRef< std::vector< MeshSharding >>  possibleResultShardingAttrs  )

static

Definition at line 217 of file ShardingPropagation.cpp.

References DBGS, getReshardingRquirementKind(), mlir::mesh::ShardingOption::makeEmpty(), and NO_RESHARDING.

Referenced by visitOp().

visitOp()

static LogicalResult visitOp ( Operation *  op, OpBuilder &  builder  )

static

Definition at line 284 of file ShardingPropagation.cpp.

References DBGS, mlir::Operation::emitOpError(), mlir::mesh::getMeshSharding(), mlir::Operation::getNumOperands(), mlir::Operation::getNumResults(), mlir::Operation::getOpOperands(), getOrderedPossibleShardingAttrs(), mlir::Operation::getResults(), mlir::Operation::hasTrait(), and selectShardingOption().

Referenced by ShardingPropagation::runOnOperation().