ShardingInterface.cpp File Reference

#include "mlir/Dialect/Mesh/Interfaces/ShardingInterface.h"
#include "mlir/Dialect/Mesh/Interfaces/ShardingInterfaceImpl.h"
#include "mlir/Dialect/Mesh/IR/MeshOps.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/Mesh/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)
static LogicalResult	addShardOp (OpBuilder &b, OpResult result, const ShardingOption &shardingOption, AffineMap map, ArrayRef< utils::IteratorType > loopTypes, ArrayRef< ReductionKind > reductionLoopKinds)
static LogicalResult	addShardOp (OpBuilder &b, OpOperand &opOperand, const ShardingOption &shardingOption, AffineMap map)
static bool	isValueCompatibleWithFullReplicationSharding (Value value, MeshShardingAttr sharding)
template<typename ValueRange , typename MeshShardingAttrRage >
static bool	areValuesCompatibleWithFullReplicationShardings (ValueRange &&values, MeshShardingAttrRage &&shardings)
static void	updateMeshAxisAssignmentForLoopIterators (ArrayRef< MeshAxis > meshAxesAssignmentForTensorAxis, AffineExpr indexingExpr, SmallVector< std::optional< SmallVector< MeshAxis >>> &meshAxesAssignmentForLoopIterators)

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 453 of file ShardingInterface.cpp.

References checkOperandAffineExpr(), mlir::OpBuilder::create(), mlir::detail::enumerate(), mlir::failed(), mlir::failure(), mlir::IROperand< DerivedT, IRValueT >::get(), mlir::get(), mlir::Builder::getContext(), mlir::Value::getLoc(), mlir::mesh::getMeshShardingAttr(), mlir::AffineMap::getNumDims(), mlir::detail::IROperandBase::getOwner(), mlir::AffineMap::getResults(), mlir::Value::getType(), mlir::mesh::ShardingOption::mesh, mlir::mesh::removeTrailingEmptySubArray(), mlir::IROperand< DerivedT, IRValueT >::set(), mlir::OpBuilder::setInsertionPoint(), mlir::mesh::ShardingOption::shardingArray, mlir::succeeded(), and mlir::success().

addShardOp() [2/2]

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

static

Definition at line 397 of file ShardingInterface.cpp.

References mlir::OpBuilder::create(), mlir::detail::enumerate(), mlir::get(), mlir::Builder::getContext(), mlir::Value::getLoc(), mlir::mesh::getMeshShardingAttr(), mlir::AffineMap::getResults(), mlir::Value::getType(), mlir::mesh::isReductionLoop(), mlir::mesh::ShardingOption::mesh, mlir::mesh::removeTrailingEmptySubArray(), mlir::Value::replaceAllUsesExcept(), mlir::OpBuilder::setInsertionPointAfterValue(), mlir::mesh::ShardingOption::shardingArray, mlir::succeeded(), and mlir::success().

areValuesCompatibleWithFullReplicationShardings()

template<typename ValueRange , typename MeshShardingAttrRage >

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

static

Definition at line 537 of file ShardingInterface.cpp.

References isValueCompatibleWithFullReplicationSharding().

checkOperandAffineExpr()

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

static

Definition at line 81 of file ShardingInterface.cpp.

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

Referenced by addShardOp().

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::failed(), mlir::failure(), mlir::AffineExpr::getKind(), mlir::Mul, and mlir::success().

Referenced by checkOperandAffineExpr().

isValueCompatibleWithFullReplicationSharding()

static bool isValueCompatibleWithFullReplicationSharding ( Value  value, MeshShardingAttr  sharding  )

static

Definition at line 527 of file ShardingInterface.cpp.

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

Referenced by areValuesCompatibleWithFullReplicationShardings().

updateMeshAxisAssignmentForLoopIterators()

static void updateMeshAxisAssignmentForLoopIterators ( ArrayRef< MeshAxis >  meshAxesAssignmentForTensorAxis, AffineExpr  indexingExpr, SmallVector< std::optional< SmallVector< MeshAxis >>> &  meshAxesAssignmentForLoopIterators  )

static

Definition at line 567 of file ShardingInterface.cpp.

References mlir::AffineDimExpr::getPosition().

Referenced by mlir::mesh::getMeshAxisAssignmentForLoopIterators().