TensorOps.cpp File Reference

#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/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/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"

Go to the source code of this file.

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. 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 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	verifyInBoundsSlice (Operation *op, RankedTensorType tensorType, ArrayRef< int64_t > staticOffsets, ArrayRef< int64_t > staticSizes, ArrayRef< int64_t > staticStrides)
	Verify that the offsets/sizes/strides-style access into the given tensor is in-bounds. More...
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)
bool	foldTensorCastPrecondition (DestinationStyleOpInterface op)

Macro Definition Documentation

GET_OP_CLASSES

#define GET_OP_CLASSES

Definition at line 4009 of file TensorOps.cpp.

Function Documentation

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

References mlir::Value::getDefiningOp().

foldIdentityOffsetSizeAndStrideOpInterface()

static LogicalResult foldIdentityOffsetSizeAndStrideOpInterface ( OffsetSizeAndStrideOpInterface  op, ShapedType  shapedType  )

static

Definition at line 2647 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 2822 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 2802 of file TensorOps.cpp.

foldTensorCastPrecondition()

bool foldTensorCastPrecondition

(

DestinationStyleOpInterface

op

)

Definition at line 3931 of file TensorOps.cpp.

References mlir::tensor::hasFoldableTensorCastOperand().

Referenced by FoldTensorCastProducerOp::matchAndRewrite().

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

References mlir::detail::enumerate().

getNumElements()

static int64_t getNumElements ( ShapedType  type )

static

Definition at line 1628 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 418 of file TensorOps.cpp.

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

parseInferType()

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

Definition at line 3107 of file TensorOps.cpp.

printInferType()

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

Definition at line 3103 of file TensorOps.cpp.

produceSliceErrorMsg()

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

static

Definition at line 2335 of file TensorOps.cpp.

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

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

References mlir::Operation::emitOpError().

verifyInBoundsSlice()

static LogicalResult verifyInBoundsSlice ( Operation *  op, RankedTensorType  tensorType, ArrayRef< int64_t >  staticOffsets, ArrayRef< int64_t >  staticSizes, ArrayRef< int64_t >  staticStrides  )

static

Verify that the offsets/sizes/strides-style access into the given tensor is in-bounds.

Only static information is verified.

Definition at line 2357 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 2756 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 1856 of file TensorOps.cpp.

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