MeshShardingInterfaceImpl.cpp File Reference

#include "mlir/Dialect/Linalg/Transforms/MeshShardingInterfaceImpl.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h"
#include "mlir/Dialect/Mesh/IR/MeshOps.h"
#include "mlir/Dialect/Mesh/Interfaces/ShardingInterface.h"
#include "mlir/Dialect/Mesh/Interfaces/ShardingInterfaceImpl.h"
#include "mlir/Dialect/Mesh/Transforms/Transforms.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/DialectRegistry.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/ImplicitLocOpBuilder.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/IR/Value.h"
#include "mlir/Interfaces/TilingInterface.h"
#include "mlir/Support/LogicalResult.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/TypeSwitch.h"
#include <iterator>
#include <optional>
#include <utility>
#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"

Go to the source code of this file.

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

Macros
#define	GET_OP_LIST

Typedefs
using	mlir::linalg::MeshAxis = mesh::MeshAxis
using	mlir::linalg::ReductionKind = mesh::ReductionKind
using	mlir::linalg::MeshShardingAttr = mesh::MeshShardingAttr
using	mlir::linalg::ShardingArray = mesh::ShardingArray
using	mlir::linalg::MeshOp = mesh::MeshOp

Functions
static ReductionKind	mlir::linalg::getReductionKind (Operation *op)
static std::optional< Operation * >	mlir::linalg::getCombinerOp (LinalgOp op)
static ReductionKind	mlir::linalg::getReductionKindOfLinalgOp (LinalgOp op)
static MeshOp	mlir::linalg::getMesh (Operation *op, ArrayRef< MeshShardingAttr > operandShardings, ArrayRef< MeshShardingAttr > resultShardings, SymbolTableCollection &symbolTable)
static Value	mlir::linalg::createDestinationPassingStyleInitOperand (LinalgOp op, Value spmdizedOperand, ArrayRef< MeshAxis > reductionMeshAxes, MeshOp meshOp, ImplicitLocOpBuilder &builder)
static SmallVector< Value >	mlir::linalg::createDestinationPassingStyleInitOperands (LinalgOp op, MeshOp meshOp, ArrayRef< Value > spmdizedOperands, ArrayRef< MeshAxis > reductionMeshAxes, IRMapping &spmdizationMap, ImplicitLocOpBuilder &builder)
static void	mlir::linalg::createAllReduceForResultWithoutPartialSharding (Value unshardedLinalgOpResult, ArrayRef< MeshAxis > opReductionMeshAxes, MeshShardingAttr resultSharding, ReductionKind reductionKind, IRMapping &spmdizationMap, ImplicitLocOpBuilder &builder)
static void	mlir::linalg::createAllReduceForResultsWithoutPartialShardings (LinalgOp unshardedOp, ArrayRef< MeshAxis > opReductionMeshAxes, ArrayRef< MeshShardingAttr > resultShardings, IRMapping &spmdizationMap, ImplicitLocOpBuilder &builder)
static void	mlir::linalg::spmdizeLinalgOpWithShardedReduction (LinalgOp op, ArrayRef< Value > spmdizedOperands, ArrayRef< MeshShardingAttr > operandShardings, ArrayRef< MeshShardingAttr > resultShardings, ArrayRef< utils::IteratorType > loopIteratorTypes, ArrayRef< SmallVector< MeshAxis >> meshAxisAssignmentForLoopIterators, IRMapping &spmdizationMap, SymbolTableCollection &symbolTable, ImplicitLocOpBuilder &builder)
template<typename OpType >
static void	mlir::linalg::registerOne (MLIRContext *ctx)
template<typename... OpTypes>
static void	mlir::linalg::registerAll (MLIRContext *ctx)
	Variadic helper function. More...
void	mlir::linalg::registerMeshShardingInterfaceExternalModels (DialectRegistry &registry)

Macro Definition Documentation

GET_OP_LIST

#define GET_OP_LIST