#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/Linalg/IR/RelayoutOpInterface.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/PatternMatch.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Interfaces/DestinationStyleOpInterface.h"
#include "mlir/Interfaces/InferIntRangeInterface.h"
#include "mlir/Interfaces/LoopLikeInterface.h"
#include "mlir/Interfaces/Utils/InferIntRangeCommon.h"
#include "mlir/Interfaces/ViewLikeInterface.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/Casting.h"
#include "llvm/Support/MathExtras.h"
#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	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.
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.
static TensorType	joinShapes (TensorType one, TensorType two)
	Compute a TensorType that has the joined shape knowledge of the two given TensorTypes.
static Value	foldExtractAfterInsert (ExtractOp extractOp)
	If we have an ExtractOp consuming an InsertOp with the same indices, we can return the InsertOp's scalar directly.
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.
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.
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.
static Value	foldInsertAfterExtractSlice (InsertSliceOp insertOp)
	Folds round-trip extract/insert slice op pairs.
static bool	foldTensorCastPrecondition (DestinationStyleOpInterface op)

Macro Definition Documentation

◆ GET_OP_CLASSES

#define GET_OP_CLASSES

Definition at line 4133 of file TensorOps.cpp.

Function Documentation

◆ foldDynamicToStaticDimSizes()

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

◆ foldExtractAfterInsert()

Value foldExtractAfterInsert ( ExtractOp extractOp )

static

If we have an ExtractOp consuming an InsertOp with the same indices, we can return the InsertOp's scalar directly.

Definition at line 1366 of file TensorOps.cpp.

References b, mlir::getAsOpFoldResult(), and mlir::Value::getDefiningOp().

◆ foldExtractAfterInsertSlice()

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 2747 of file TensorOps.cpp.

References b.

◆ foldIdentityOffsetSizeAndStrideOpInterface()

LogicalResult foldIdentityOffsetSizeAndStrideOpInterface	(	OffsetSizeAndStrideOpInterface	op,
		ShapedType	shapedType )

static

Definition at line 2725 of file TensorOps.cpp.

References b, mlir::getConstantIntValue(), and success().

◆ foldInsertAfterExtractSlice()

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 2905 of file TensorOps.cpp.

References b, and mlir::Value::getDefiningOp().

◆ foldInsertAfterInsertSlice()

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 2885 of file TensorOps.cpp.

References b, and success().

◆ foldTensorCastPrecondition()

bool foldTensorCastPrecondition ( DestinationStyleOpInterface op )

static

Definition at line 4055 of file TensorOps.cpp.

References mlir::tensor::hasFoldableTensorCastOperand().

Referenced by FoldTensorCastProducerOp::matchAndRewrite().

◆ getDroppedDims()

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.

◆ getNumElements()

int64_t getNumElements ( ShapedType type )

static

Definition at line 1810 of file TensorOps.cpp.

◆ joinShapes()

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 416 of file TensorOps.cpp.

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

◆ produceSliceErrorMsg()

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

static

Definition at line 2432 of file TensorOps.cpp.

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

◆ verifyGatherOrScatterDims()

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

static

Definition at line 1552 of file TensorOps.cpp.

References mlir::Operation::emitOpError(), indices, and success().

◆ verifyInsertSliceOp()

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 2834 of file TensorOps.cpp.

References mlir::isRankReducedType().

◆ verifyTensorReshapeOp()

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

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

static

Definition at line 2037 of file TensorOps.cpp.

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

Classes

Macros

Functions

Macro Definition Documentation

◆ GET_OP_CLASSES

Function Documentation

◆ foldDynamicToStaticDimSizes()

◆ foldExtractAfterInsert()

◆ foldExtractAfterInsertSlice()

◆ foldIdentityOffsetSizeAndStrideOpInterface()

◆ foldInsertAfterExtractSlice()

◆ foldInsertAfterInsertSlice()

◆ foldTensorCastPrecondition()

◆ getDroppedDims()

◆ getNumElements()

◆ joinShapes()

◆ produceSliceErrorMsg()

◆ verifyGatherOrScatterDims()

◆ verifyInsertSliceOp()

◆ verifyTensorReshapeOp()