MLIR  16.0.0git
Classes | Macros | Functions
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/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/BlockAndValueMapping.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributeInterfaces.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Interfaces/DestinationStyleOpInterface.h"
#include "mlir/Support/MathExtras.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallBitVector.h"
#include "llvm/ADT/StringRef.h"
#include <algorithm>
#include "mlir/Dialect/Tensor/IR/TensorOps.cpp.inc"
+ Include dependency graph for TensorOps.cpp:

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...
 

Macros

#define GET_OP_CLASSES
 

Functions

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, int64_t rank, StringRef gatherOrScatter, StringRef sourceOrDest)
 
static int64_t getNumElements (ShapedType type)
 
static RankedTensorType computeTensorReshapeCollapsedType (RankedTensorType type, ArrayRef< AffineMap > reassociation)
 Compute the RankedTensorType obtained by applying reassociation to type. More...
 
static bool isSameTypesWithoutEncoding (Type tp1, Type tp2)
 
template<typename TensorReshapeOp , bool isExpansion = std::is_same< TensorReshapeOp, ExpandShapeOp>::value>
static LogicalResult verifyTensorReshapeOp (TensorReshapeOp op, RankedTensorType expandedType, RankedTensorType collapsedType)
 
template<typename OpTy >
static LogicalResult produceSliceErrorMsg (SliceVerificationResult result, OpTy op, Type 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 (ShapedType srcType, ShapedType dstType, ArrayRef< int64_t > staticOffsets, ArrayRef< int64_t > staticSizes, ArrayRef< int64_t > staticStrides, ShapedType *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, 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, OpFoldResultgetDimAndTileMappingImpl (OpTy op)
 
template<typename OpTy >
static SmallVector< OpFoldResultgetMixedTilesImpl (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)
 
static bool areNotFullTiles (ArrayRef< int64_t > inputShape, DenseMap< int64_t, OpFoldResult > const &dimAndTileMapping)
 Check if we have enough static information to catch undefined behavior when the tile size does not divide perfectly the dimension of the input tensor. More...
 
template<typename OpTy >
bool areTilesAndTiledDimsAllConstant (OpTy op)
 Returns true if the tiles and the tiled dims are constant. More...
 

Macro Definition Documentation

◆ GET_OP_CLASSES

#define GET_OP_CLASSES

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

Referenced by commonVerifierPackAndUnPackOp().

◆ areNotFullTiles()

static bool areNotFullTiles ( ArrayRef< int64_t >  inputShape,
DenseMap< int64_t, OpFoldResult > const &  dimAndTileMapping 
)
static

Check if we have enough static information to catch undefined behavior when the tile size does not divide perfectly the dimension of the input tensor.

Definition at line 3284 of file TensorOps.cpp.

References mlir::getConstantIntValue().

◆ areTilesAndTiledDimsAllConstant()

template<typename OpTy >
bool areTilesAndTiledDimsAllConstant ( OpTy  op)

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

Definition at line 3354 of file TensorOps.cpp.

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

◆ commonVerifierPackAndUnPackOp()

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

◆ computeTensorReshapeCollapsedType()

static RankedTensorType computeTensorReshapeCollapsedType ( RankedTensorType  type,
ArrayRef< AffineMap reassociation 
)
static

Compute the RankedTensorType obtained by applying reassociation to type.

Definition at line 1331 of file TensorOps.cpp.

References mlir::isReassociationValid().

Referenced by verifyTensorReshapeOp().

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

References mlir::Value::getDefiningOp().

◆ foldIdentityOffsetSizeAndStrideOpInterface()

static LogicalResult foldIdentityOffsetSizeAndStrideOpInterface ( OffsetSizeAndStrideOpInterface  op,
ShapedType  shapedType 
)
static

Definition at line 2044 of file TensorOps.cpp.

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

◆ 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 2216 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 2196 of file TensorOps.cpp.

References mlir::failure(), and mlir::success().

◆ getDimAndTileMappingImpl()

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

Definition at line 3069 of file TensorOps.cpp.

References value.

◆ getMixedTilesImpl()

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

Definition at line 3084 of file TensorOps.cpp.

References mlir::Builder::getI64IntegerAttr(), and value.

◆ getNumElements()

static int64_t getNumElements ( ShapedType  type)
static

Definition at line 1252 of file TensorOps.cpp.

Referenced by getPushConstantVariable(), and mlirElementsAttrGetNumElements().

◆ getStaticTilesImpl()

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

Definition at line 3100 of file TensorOps.cpp.

References mlir::dispatchIndexOpFoldResults(), and value.

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

Referenced by commonVerifierPackAndUnPackOp().

◆ isSameTypesWithoutEncoding()

static bool isSameTypesWithoutEncoding ( Type  tp1,
Type  tp2 
)
static

Definition at line 1370 of file TensorOps.cpp.

References mlir::Type::dyn_cast().

Referenced by verifyTensorReshapeOp().

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

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

◆ parseInferType()

ParseResult parseInferType ( OpAsmParser parser,
Optional< OpAsmParser::UnresolvedOperand optOperand,
Type typeToInfer,
Type  typeToInferFrom 
)

Definition at line 2472 of file TensorOps.cpp.

References mlir::success().

◆ printInferType()

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

Definition at line 2469 of file TensorOps.cpp.

◆ produceSliceErrorMsg()

template<typename OpTy >
static LogicalResult produceSliceErrorMsg ( SliceVerificationResult  result,
OpTy  op,
Type  expectedType 
)
static

◆ reifyResultShapesImpl()

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

Definition at line 3055 of file TensorOps.cpp.

References mlir::OpBuilder::createOrFold(), mlir::success(), and value.

◆ verifyGatherOrScatterDims()

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

Definition at line 986 of file TensorOps.cpp.

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

◆ verifyInsertSliceOp()

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

Rank-reducing type verification for both InsertSliceOp and ParallelInsertSliceOp.

Definition at line 2157 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