#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributeInterfaces.h"
#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Interfaces/DestinationStyleOpInterface.h"
#include "mlir/Interfaces/LoopLikeInterface.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallBitVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/MathExtras.h"
#include <algorithm>
#include <optional>
#include "mlir/Dialect/Tensor/IR/TensorOps.cpp.inc"

Classes
struct	SliceReturnTypeCanonicalizer
	Return the canonical type of the result of an extract_slice op. More...

struct	SliceCanonicalizer
	A canonicalizer wrapper to replace ExtractSliceOps. More...

struct	FoldTensorCastPackOp
	Folds a tensor.cast op into a consuming tensor::PackOp op if the `tensor.cast` has source that is more static than the consuming op. More...

struct	FoldTensorCastProducerOp
	Folds a tensor.cast op into a consuming DestinationStyleOpInterface op if the `tensor.cast` has source that is more static than the consuming op. More...

Macros
#define	GET_OP_CLASSES

Functions
static llvm::SmallBitVector	getDroppedDims (ArrayRef< int64_t > reducedShape, ArrayRef< OpFoldResult > mixedSizes)
	Compute the dropped dimensions of a rank-reducing tensor.extract_slice op or rank-extending tensor.insert_slice op. More...

static RankedTensorType	foldDynamicToStaticDimSizes (RankedTensorType type, ValueRange dynamicSizes, SmallVector< Value > &foldedDynamicSizes)
	Given a ranked tensor type and a range of values that defines its dynamic dimension sizes, turn all dynamic sizes that have a constant value into static dimension sizes. More...

static TensorType	joinShapes (TensorType one, TensorType two)
	Compute a TensorType that has the joined shape knowledge of the two given TensorTypes. More...

static OpFoldResult	reshapeConstantSource (DenseElementsAttr source, TensorType result, std::optional< Attribute > cst=std::nullopt)
	Try to remove a tensor operation if it would only reshape a constant. More...

static LogicalResult	verifyGatherOrScatterDims (Operation *op, ArrayRef< int64_t > dims, ArrayRef< int64_t > indices, int64_t rank, StringRef gatherOrScatter, StringRef sourceOrDest)

static int64_t	getNumElements (ShapedType type)

template<typename TensorReshapeOp , bool isExpansion = std::is_same< TensorReshapeOp, ExpandShapeOp>::value>
static LogicalResult	verifyTensorReshapeOp (TensorReshapeOp op, RankedTensorType expandedType, RankedTensorType collapsedType)

static LogicalResult	produceSliceErrorMsg (SliceVerificationResult result, Operation *op, RankedTensorType expectedType)

static LogicalResult	foldIdentityOffsetSizeAndStrideOpInterface (OffsetSizeAndStrideOpInterface op, ShapedType shapedType)

static Value	foldExtractAfterInsertSlice (ExtractSliceOp extractOp)
	If we have an ExtractSliceOp consuming an InsertSliceOp with the same slice, we can return the InsertSliceOp's source directly. More...

static SliceVerificationResult	verifyInsertSliceOp (RankedTensorType srcType, RankedTensorType dstType, ArrayRef< int64_t > staticOffsets, ArrayRef< int64_t > staticSizes, ArrayRef< int64_t > staticStrides, RankedTensorType *expectedType=nullptr)
	Rank-reducing type verification for both InsertSliceOp and ParallelInsertSliceOp. More...

static LogicalResult	foldInsertAfterInsertSlice (InsertSliceOp insertOp)
	If we have two consecutive InsertSliceOp writing to the same slice, we can mutate the second InsertSliceOp's destination to the first one's. More...

static Value	foldInsertAfterExtractSlice (InsertSliceOp insertOp)
	Folds round-trip extract/insert slice op pairs. More...

void	printInferType (OpAsmPrinter &printer, Operation *op, Value optOperand, Type typeToInfer, Type typeToInferFrom)

ParseResult	parseInferType (OpAsmParser &parser, std::optional< OpAsmParser::UnresolvedOperand > optOperand, Type &typeToInfer, Type typeToInferFrom)

template<typename OpTy >
static LogicalResult	reifyResultShapesImpl (OpTy op, OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes)

template<typename OpTy >
static DenseMap< int64_t, OpFoldResult >	getDimAndTileMappingImpl (OpTy op)

template<typename OpTy >
static SmallVector< OpFoldResult >	getMixedTilesImpl (OpTy op)

template<typename OpTy >
static SmallVector< int64_t >	getStaticTilesImpl (OpTy op)

static bool	isInvalidPackingPosSpecification (ArrayRef< int64_t > dimsPos, size_t rank)
	Returns true if `dimsPos` is invalid. More...

static bool	areAllInBound (ArrayRef< int64_t > sourceShape, ArrayRef< int64_t > limitShape)
	Returns true if the dimension of `sourceShape` is smaller than the dimension of the `limitShape`. More...

template<typename OpTy >
static LogicalResult	commonVerifierPackAndUnPackOp (OpTy packOrUnPack)

template<typename OpTy >
static PackOrUnPackTransposeResult	commonPermutationOfPackAndUnPackOp (OpTy packOrUnPackOp, ArrayRef< int64_t > innerPermutation, ArrayRef< int64_t > outerPermutation)

static SmallVector< int64_t >	asShapeWithAnyValueAsDynamic (ArrayRef< OpFoldResult > ofrs)
	Converts OpFoldResults to int64_t shape entries, unconditionally mapping all Value's to kDynamic, even if they are arith.constant values. More...

static SmallVector< int64_t >	getPackOpResultTypeShape (ArrayRef< int64_t > sourceShape, ArrayRef< int64_t > innerTileSizes, ArrayRef< int64_t > innerDimsPos, ArrayRef< int64_t > outerDimsPerm)
	Helper for PackOp::{getResultShape,inferPackedType}. More...

template<typename OpTy >
bool	areTilesAndTiledDimsAllConstant (OpTy op)
	Returns true if the tiles and the tiled dims are constant. More...

static bool	hasSameInnerOuterAttribute (PackOp packOp, UnPackOp unPackOp)

static bool	haveSameTiles (PackOp packOp, UnPackOp unPackOp)

static bool	paddingIsNotNeeded (PackOp op)
	Returns true if the pack op does not need a padding value. More...

static bool	inferStaticShape (PackOp packOp, SmallVectorImpl< int64_t > &srcShape, SmallVectorImpl< int64_t > &destShape)
	Returns true if the `srcShape` or `destShape` is different from the one in `packOp` and populates each with the inferred static shape. More...

template<typename PackOrUnpackOp >
static bool	isLikePadUnPad (PackOrUnpackOp packOp, RankedTensorType packedTensorType)

static bool	inferStaticShape (UnPackOp op, SmallVectorImpl< int64_t > &srcShape, SmallVectorImpl< int64_t > &destShape)
	Returns true if the `srcShape` or `destShape` is different from the one in `op` and populates each with the inferred static shape. More...

bool	foldTensorCastPrecondition (DestinationStyleOpInterface op)

static SmallVector< Value >	getNewOperands (DestinationStyleOpInterface op, SmallVector< Type > &newResTy)

Macro Definition Documentation

◆ GET_OP_CLASSES

#define GET_OP_CLASSES

Definition at line 4930 of file TensorOps.cpp.

Function Documentation

◆ areAllInBound()

static bool areAllInBound	(	ArrayRef< int64_t >	sourceShape,
		ArrayRef< int64_t >	limitShape
	)

static

Returns true if the dimension of sourceShape is smaller than the dimension of the limitShape.

Definition at line 3925 of file TensorOps.cpp.

Referenced by commonVerifierPackAndUnPackOp().

◆ areTilesAndTiledDimsAllConstant()

template<typename OpTy >

bool areTilesAndTiledDimsAllConstant ( OpTy op )

Returns true if the tiles and the tiled dims are constant.

Definition at line 4325 of file TensorOps.cpp.

References mlir::getConstantIntValue(), and mlir::tile().

◆ asShapeWithAnyValueAsDynamic()

static SmallVector<int64_t> asShapeWithAnyValueAsDynamic ( ArrayRef< OpFoldResult > ofrs )

static

Converts OpFoldResults to int64_t shape entries, unconditionally mapping all Value's to kDynamic, even if they are arith.constant values.

Definition at line 4191 of file TensorOps.cpp.

References mlir::getConstantIntValue().

◆ commonPermutationOfPackAndUnPackOp()

template<typename OpTy >

static PackOrUnPackTransposeResult commonPermutationOfPackAndUnPackOp	(	OpTy	packOrUnPackOp,
		ArrayRef< int64_t >	innerPermutation,
		ArrayRef< int64_t >	outerPermutation
	)

static

Definition at line 4036 of file TensorOps.cpp.

References mlir::applyPermutationToVector(), and mlir::isPermutationVector().

◆ commonVerifierPackAndUnPackOp()

template<typename OpTy >

static LogicalResult commonVerifierPackAndUnPackOp ( OpTy packOrUnPack )

static

Definition at line 3940 of file TensorOps.cpp.

References areAllInBound(), mlir::Operation::emitError(), mlir::isConstantIntValue(), isInvalidPackingPosSpecification(), and mlir::tile().

◆ foldDynamicToStaticDimSizes()

static RankedTensorType foldDynamicToStaticDimSizes	(	RankedTensorType	type,
		ValueRange	dynamicSizes,
		SmallVector< Value > &	foldedDynamicSizes
	)

static

Given a ranked tensor type and a range of values that defines its dynamic dimension sizes, turn all dynamic sizes that have a constant value into static dimension sizes.

Definition at line 180 of file TensorOps.cpp.

References mlir::get(), and mlir::getConstantIntValue().

◆ foldExtractAfterInsertSlice()

static Value foldExtractAfterInsertSlice ( ExtractSliceOp extractOp )

static

If we have an ExtractSliceOp consuming an InsertSliceOp with the same slice, we can return the InsertSliceOp's source directly.

Definition at line 2601 of file TensorOps.cpp.

References mlir::Value::getDefiningOp().

◆ foldIdentityOffsetSizeAndStrideOpInterface()

static LogicalResult foldIdentityOffsetSizeAndStrideOpInterface	(	OffsetSizeAndStrideOpInterface	op,
		ShapedType	shapedType
	)

static

Definition at line 2579 of file TensorOps.cpp.

References mlir::getConstantIntValue().

◆ foldInsertAfterExtractSlice()

static Value foldInsertAfterExtractSlice ( InsertSliceOp insertOp )

static

Folds round-trip extract/insert slice op pairs.

Example:

%0 = tensor.extract_slice %val[0, 0, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1]

%1 = tensor.insert_slice %0 into %val[0, 0, 0, 0] [1, 1, 2, 4] [1, 1, 1, 1]

can be folded into val.

Definition at line 2747 of file TensorOps.cpp.

References mlir::Value::getDefiningOp().

◆ foldInsertAfterInsertSlice()

static LogicalResult foldInsertAfterInsertSlice ( InsertSliceOp insertOp )

static

If we have two consecutive InsertSliceOp writing to the same slice, we can mutate the second InsertSliceOp's destination to the first one's.

Example:

%0 = tensor.insert_slice %slice0 into %input[0, 0] [64, 64] [1, 1]

%1 = tensor.insert_slice %slice1 into %0[0, 0] [64, 64] [1, 1]

folds into:

%1 = tensor.insert_slice %slice1 into %input[0, 0] [64, 64] [1, 1]

This pattern works with both InsertSliceOp and ParallelInsertSliceOp.

Definition at line 2727 of file TensorOps.cpp.

◆ foldTensorCastPrecondition()

bool foldTensorCastPrecondition ( DestinationStyleOpInterface op )

Definition at line 4757 of file TensorOps.cpp.

Referenced by FoldTensorCastProducerOp::matchAndRewrite(), and FoldTensorCastPackOp::matchAndRewrite().

◆ getDimAndTileMappingImpl()

template<typename OpTy >

static DenseMap<int64_t, OpFoldResult> getDimAndTileMappingImpl ( OpTy op )

static

Definition at line 3864 of file TensorOps.cpp.

◆ getDroppedDims()

static llvm::SmallBitVector getDroppedDims	(	ArrayRef< int64_t >	reducedShape,
		ArrayRef< OpFoldResult >	mixedSizes
	)

static

Compute the dropped dimensions of a rank-reducing tensor.extract_slice op or rank-extending tensor.insert_slice op.

Definition at line 136 of file TensorOps.cpp.

References mlir::detail::enumerate().

◆ getMixedTilesImpl()

template<typename OpTy >

static SmallVector<OpFoldResult> getMixedTilesImpl ( OpTy op )

static

Definition at line 3879 of file TensorOps.cpp.

References mlir::Builder::getI64IntegerAttr().

◆ getNewOperands()

static SmallVector<Value> getNewOperands	(	DestinationStyleOpInterface	op,
		SmallVector< Type > &	newResTy
	)

static

Definition at line 4777 of file TensorOps.cpp.

References mlir::tensor::canFoldIntoConsumerOp().

Referenced by FoldTensorCastProducerOp::matchAndRewrite(), and FoldTensorCastPackOp::matchAndRewrite().

◆ getNumElements()

static int64_t getNumElements ( ShapedType type )

static

Definition at line 1604 of file TensorOps.cpp.

◆ getPackOpResultTypeShape()

static SmallVector<int64_t> getPackOpResultTypeShape	(	ArrayRef< int64_t >	sourceShape,
		ArrayRef< int64_t >	innerTileSizes,
		ArrayRef< int64_t >	innerDimsPos,
		ArrayRef< int64_t >	outerDimsPerm
	)

static

Helper for PackOp::{getResultShape,inferPackedType}.

Returns the shape of the packed type. Having a shared helper helps implement these two methods in a way that ensures that they agree on which dimensions are dynamic.

Definition at line 4206 of file TensorOps.cpp.

References mlir::applyPermutationToVector(), and mlir::detail::enumerate().

◆ getStaticTilesImpl()

template<typename OpTy >

static SmallVector<int64_t> getStaticTilesImpl ( OpTy op )

static

Definition at line 3895 of file TensorOps.cpp.

References mlir::dispatchIndexOpFoldResults().

◆ hasSameInnerOuterAttribute()

static bool hasSameInnerOuterAttribute	(	PackOp	packOp,
		UnPackOp	unPackOp
	)

static

Definition at line 4356 of file TensorOps.cpp.

References mlir::isIdentityPermutation().

◆ haveSameTiles()

static bool haveSameTiles	(	PackOp	packOp,
		UnPackOp	unPackOp
	)

static

Definition at line 4370 of file TensorOps.cpp.

References mlir::isEqualConstantIntOrValue().

◆ inferStaticShape() [1/2]

static bool inferStaticShape	(	PackOp	packOp,
		SmallVectorImpl< int64_t > &	srcShape,
		SmallVectorImpl< int64_t > &	destShape
	)

static

Returns true if the srcShape or destShape is different from the one in packOp and populates each with the inferred static shape.

Definition at line 4397 of file TensorOps.cpp.

References mlir::invertPermutationVector().

◆ inferStaticShape() [2/2]

static bool inferStaticShape	(	UnPackOp	op,
		SmallVectorImpl< int64_t > &	srcShape,
		SmallVectorImpl< int64_t > &	destShape
	)

static

Returns true if the srcShape or destShape is different from the one in op and populates each with the inferred static shape.

Definition at line 4657 of file TensorOps.cpp.

References mlir::invertPermutationVector().

◆ isInvalidPackingPosSpecification()

static bool isInvalidPackingPosSpecification	(	ArrayRef< int64_t >	dimsPos,
		size_t	rank
	)

static

Returns true if dimsPos is invalid.

It is invalid when: a) It contains duplicate. b) At least one dimension is out of bound (dimPos is >= 0 and < rank). c) The number of elements in dimsPos is > than rank.

Definition at line 3908 of file TensorOps.cpp.

Referenced by commonVerifierPackAndUnPackOp().

◆ isLikePadUnPad()

template<typename PackOrUnpackOp >

static bool isLikePadUnPad	(	PackOrUnpackOp	packOp,
		RankedTensorType	packedTensorType
	)

static

Definition at line 4490 of file TensorOps.cpp.

◆ joinShapes()

static TensorType joinShapes	(	TensorType	one,
		TensorType	two
	)

static

Compute a TensorType that has the joined shape knowledge of the two given TensorTypes.

The element types need to match.

Definition at line 386 of file TensorOps.cpp.

References mlir::get(), mlir::TensorType::getElementType(), and mlir::TensorType::hasRank().

◆ paddingIsNotNeeded()

static bool paddingIsNotNeeded ( PackOp op )

static

Returns true if the pack op does not need a padding value.

Definition at line 4383 of file TensorOps.cpp.

◆ parseInferType()

ParseResult parseInferType	(	OpAsmParser &	parser,
		std::optional< OpAsmParser::UnresolvedOperand >	optOperand,
		Type &	typeToInfer,
		Type	typeToInferFrom
	)

Definition at line 3032 of file TensorOps.cpp.

◆ printInferType()

void printInferType	(	OpAsmPrinter &	printer,
		Operation *	op,
		Value	optOperand,
		Type	typeToInfer,
		Type	typeToInferFrom
	)

Definition at line 3028 of file TensorOps.cpp.

◆ produceSliceErrorMsg()

static LogicalResult produceSliceErrorMsg	(	SliceVerificationResult	result,
		Operation *	op,
		RankedTensorType	expectedType
	)

static

Definition at line 2306 of file TensorOps.cpp.

References mlir::ElemTypeMismatch, mlir::Operation::emitError(), mlir::RankTooLarge, mlir::SizeMismatch, and mlir::Success.

◆ reifyResultShapesImpl()

template<typename OpTy >

static LogicalResult reifyResultShapesImpl	(	OpTy	op,
		OpBuilder &	builder,
		ReifiedRankedShapedTypeDims &	reifiedReturnShapes
	)

static

Definition at line 3852 of file TensorOps.cpp.

References mlir::tensor::getMixedSizes().

◆ reshapeConstantSource()

static OpFoldResult reshapeConstantSource	(	DenseElementsAttr	source,
		TensorType	result,
		std::optional< Attribute >	cst = `std::nullopt`
	)

static

Try to remove a tensor operation if it would only reshape a constant.

Removes the op and replaces the constant with a new constant of the result shape. When an optional cst attribute is passed, it is reshaped only if the splat value matches the value in the attribute.

Definition at line 1127 of file TensorOps.cpp.

References mlir::DenseElementsAttr::getSplatValue(), mlir::DenseElementsAttr::isSplat(), and mlir::DenseElementsAttr::resizeSplat().

◆ verifyGatherOrScatterDims()

static LogicalResult verifyGatherOrScatterDims	(	Operation *	op,
		ArrayRef< int64_t >	dims,
		ArrayRef< int64_t >	indices,
		int64_t	rank,
		StringRef	gatherOrScatter,
		StringRef	sourceOrDest
	)

static

Definition at line 1346 of file TensorOps.cpp.

References mlir::Operation::emitOpError().

◆ verifyInsertSliceOp()

static SliceVerificationResult verifyInsertSliceOp	(	RankedTensorType	srcType,
		RankedTensorType	dstType,
		ArrayRef< int64_t >	staticOffsets,
		ArrayRef< int64_t >	staticSizes,
		ArrayRef< int64_t >	staticStrides,
		RankedTensorType *	expectedType = `nullptr`
	)

static

Rank-reducing type verification for both InsertSliceOp and ParallelInsertSliceOp.

Definition at line 2688 of file TensorOps.cpp.

References mlir::isRankReducedType().

◆ verifyTensorReshapeOp()

template<typename TensorReshapeOp , bool isExpansion = std::is_same< TensorReshapeOp, ExpandShapeOp>::value>

static LogicalResult verifyTensorReshapeOp	(	TensorReshapeOp	op,
		RankedTensorType	expandedType,
		RankedTensorType	collapsedType
	)

static

Definition at line 1828 of file TensorOps.cpp.

References mlir::tensor::isSameTypeWithoutEncoding(), and mlir::verifyReshapeLikeTypes().

Classes

Macros

Functions

Macro Definition Documentation

◆ GET_OP_CLASSES

Function Documentation

◆ areAllInBound()

◆ areTilesAndTiledDimsAllConstant()

◆ asShapeWithAnyValueAsDynamic()

◆ commonPermutationOfPackAndUnPackOp()

◆ commonVerifierPackAndUnPackOp()

◆ foldDynamicToStaticDimSizes()

◆ foldExtractAfterInsertSlice()

◆ foldIdentityOffsetSizeAndStrideOpInterface()

◆ foldInsertAfterExtractSlice()

◆ foldInsertAfterInsertSlice()

◆ foldTensorCastPrecondition()

◆ getDimAndTileMappingImpl()

◆ getDroppedDims()

◆ getMixedTilesImpl()

◆ getNewOperands()

◆ getNumElements()

◆ getPackOpResultTypeShape()

◆ getStaticTilesImpl()

◆ hasSameInnerOuterAttribute()

◆ haveSameTiles()

◆ inferStaticShape() [1/2]

◆ inferStaticShape() [2/2]

◆ isInvalidPackingPosSpecification()

◆ isLikePadUnPad()

◆ joinShapes()

◆ paddingIsNotNeeded()

◆ parseInferType()

◆ printInferType()

◆ produceSliceErrorMsg()

◆ reifyResultShapesImpl()

◆ reshapeConstantSource()

◆ verifyGatherOrScatterDims()

◆ verifyInsertSliceOp()

◆ verifyTensorReshapeOp()