ShardingInterface.cpp File Reference

#include "mlir/Dialect/Shard/Interfaces/ShardingInterface.h"
#include "mlir/Dialect/Shard/Interfaces/ShardingInterfaceImpl.h"
#include "mlir/Dialect/Shard/IR/ShardOps.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/Support/Debug.h"
#include <utility>
#include "mlir/Dialect/Shard/Interfaces/ShardingInterface.cpp.inc"

Go to the source code of this file.

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

Functions
static LogicalResult	checkOperandAffineExprRecursively (AffineExpr expr, SmallVectorImpl< bool > &seenIds)
static FailureOr< llvm::SmallSet< unsigned, 2 > >	checkOperandAffineExpr (AffineExpr expr, unsigned numDims)
template<typename T >
SmallVector< GridAxesAttr >	fromArrayOfVector (MLIRContext *ctxt, const SmallVector< SmallVector< T >> &vec)
static Sharding	getSharding (OpResult result, const ShardingOption &shardingOption, AffineMap map, ArrayRef< utils::IteratorType > loopTypes)
static FailureOr< Sharding >	getSharding (OpOperand &opOperand, const ShardingOption &shardingOption, AffineMap map)
static LogicalResult	addShardOp (OpBuilder &b, OpResult result, const ShardingOption &shardingOption, AffineMap map, ArrayRef< utils::IteratorType > loopTypes)
static LogicalResult	addShardOp (OpBuilder &b, OpOperand &opOperand, const ShardingOption &shardingOption, AffineMap map)
static bool	isValueCompatibleWithFullReplicationSharding (Value value, const Sharding &sharding)
template<typename ValueRange , typename ShardingRage >
static bool	areValuesCompatibleWithFullReplicationShardings (ValueRange &&values, ShardingRage &&shardings)
static void	updateGridAxisAssignmentForLoopIterators (ArrayRef< GridAxis > gridAxesAssignmentForTensorAxis, AffineExpr indexingExpr, SmallVector< std::optional< SmallVector< GridAxis >>> &gridAxesAssignmentForLoopIterators)

Macro Definition Documentation

DBGS

#define DBGS

(

)

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

Definition at line 24 of file ShardingInterface.cpp.

DEBUG_TYPE

#define DEBUG_TYPE   "sharding-interface"

Definition at line 23 of file ShardingInterface.cpp.

Function Documentation

addShardOp() [1/2]

static LogicalResult addShardOp ( OpBuilder &  b, OpOperand &  opOperand, const ShardingOption &  shardingOption, AffineMap  map  )

static

Definition at line 473 of file ShardingInterface.cpp.

References mlir::remark::failed(), mlir::shard::getSharding(), and mlir::shard::maybeInsertSourceShardingAnnotation().

addShardOp() [2/2]

static LogicalResult addShardOp ( OpBuilder &  b, OpResult  result, const ShardingOption &  shardingOption, AffineMap  map, ArrayRef< utils::IteratorType >  loopTypes  )

static

Definition at line 461 of file ShardingInterface.cpp.

References mlir::shard::getSharding(), and mlir::shard::maybeInsertTargetShardingAnnotation().

Referenced by mlir::shard::detail::defaultAddShardingAnnotations().

areValuesCompatibleWithFullReplicationShardings()

template<typename ValueRange , typename ShardingRage >

static bool areValuesCompatibleWithFullReplicationShardings ( ValueRange &&  values, ShardingRage &&  shardings  )

static

Definition at line 528 of file ShardingInterface.cpp.

References isValueCompatibleWithFullReplicationSharding().

Referenced by mlir::shard::partitionFullyReplicatedOperation().

checkOperandAffineExpr()

static FailureOr<llvm::SmallSet<unsigned, 2> > checkOperandAffineExpr ( AffineExpr  expr, unsigned  numDims  )

static

Definition at line 82 of file ShardingInterface.cpp.

References checkOperandAffineExprRecursively(), mlir::detail::enumerate(), and mlir::remark::failed().

Referenced by mlir::shard::detail::defaultGetShardingOption(), and getSharding().

checkOperandAffineExprRecursively()

static LogicalResult checkOperandAffineExprRecursively ( AffineExpr  expr, SmallVectorImpl< bool > &  seenIds  )

static

Definition at line 36 of file ShardingInterface.cpp.

References mlir::Add, mlir::Constant, mlir::DimId, mlir::remark::failed(), mlir::AffineExpr::getKind(), and mlir::Mul.

Referenced by checkOperandAffineExpr().

fromArrayOfVector()

template<typename T >

SmallVector<GridAxesAttr> fromArrayOfVector	(	MLIRContext *	ctxt,
		const SmallVector< SmallVector< T >> &	vec
	)

Definition at line 97 of file ShardingInterface.cpp.

References mlir::detail::DenseArrayAttrImpl< T >::get().

Referenced by getSharding().

getSharding() [1/2]

static FailureOr<Sharding> getSharding ( OpOperand &  opOperand, const ShardingOption &  shardingOption, AffineMap  map  )

static

Definition at line 384 of file ShardingInterface.cpp.

References checkOperandAffineExpr(), mlir::detail::enumerate(), mlir::remark::failed(), fromArrayOfVector(), mlir::IROperand< DerivedT, IRValueT >::get(), mlir::shard::Sharding::get(), mlir::Value::getContext(), mlir::AffineMap::getNumDims(), mlir::AffineMap::getResults(), mlir::Value::getType(), mlir::shard::ShardingOption::grid, mlir::Type::isIntOrIndexOrFloat(), mlir::shard::removeTrailingEmptySubArray(), and mlir::shard::ShardingOption::shardingArray.

getSharding() [2/2]

static Sharding getSharding ( OpResult  result, const ShardingOption &  shardingOption, AffineMap  map, ArrayRef< utils::IteratorType >  loopTypes  )

static

Definition at line 362 of file ShardingInterface.cpp.

References mlir::detail::enumerate(), fromArrayOfVector(), mlir::shard::Sharding::get(), mlir::Value::getContext(), mlir::AffineMap::getResults(), mlir::Value::getType(), mlir::shard::ShardingOption::grid, mlir::shard::removeTrailingEmptySubArray(), and mlir::shard::ShardingOption::shardingArray.

isValueCompatibleWithFullReplicationSharding()

static bool isValueCompatibleWithFullReplicationSharding ( Value  value, const Sharding &  sharding  )

static

Definition at line 517 of file ShardingInterface.cpp.

References mlir::Value::getType(), and mlir::shard::isFullReplication().

Referenced by areValuesCompatibleWithFullReplicationShardings().

updateGridAxisAssignmentForLoopIterators()

static void updateGridAxisAssignmentForLoopIterators ( ArrayRef< GridAxis >  gridAxesAssignmentForTensorAxis, AffineExpr  indexingExpr, SmallVector< std::optional< SmallVector< GridAxis >>> &  gridAxesAssignmentForLoopIterators  )

static

Definition at line 557 of file ShardingInterface.cpp.

References mlir::AffineDimExpr::getPosition().

Referenced by mlir::shard::getGridAxisAssignmentForLoopIterators().