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

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

(

)

Value:

(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]

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

static

Definition at line 473 of file ShardingInterface.cpp.

References b, mlir::shard::getSharding(), mlir::shard::maybeInsertSourceShardingAnnotation(), and success().

addShardOp() [2/2]

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 b, mlir::shard::getSharding(), mlir::shard::maybeInsertTargetShardingAnnotation(), result, and success().

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

areValuesCompatibleWithFullReplicationShardings()

template<typename ValueRange, typename ShardingRage>

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

static

Definition at line 528 of file ShardingInterface.cpp.

References isValueCompatibleWithFullReplicationSharding().

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

checkOperandAffineExpr()

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

static

Definition at line 82 of file ShardingInterface.cpp.

References checkOperandAffineExprRecursively().

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

checkOperandAffineExprRecursively()

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

static

Definition at line 36 of file ShardingInterface.cpp.

References mlir::Add, checkOperandAffineExprRecursively(), mlir::Constant, mlir::DimId, mlir::AffineExpr::getKind(), lhs, mlir::Mul, rhs, and success().

Referenced by checkOperandAffineExpr(), and checkOperandAffineExprRecursively().

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< int16_t >::get().

Referenced by getSharding(), and getSharding().

getSharding() [1/2]

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

static

Definition at line 384 of file ShardingInterface.cpp.

References checkOperandAffineExpr(), 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]

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

static

Definition at line 362 of file ShardingInterface.cpp.

References fromArrayOfVector(), mlir::shard::Sharding::get(), mlir::AffineMap::getResults(), mlir::shard::ShardingOption::grid, mlir::shard::removeTrailingEmptySubArray(), result, and mlir::shard::ShardingOption::shardingArray.

isValueCompatibleWithFullReplicationSharding()

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

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().