MLIR  19.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/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/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 <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 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, 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, 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)
 
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...
 

Macro Definition Documentation

◆ GET_OP_CLASSES

#define GET_OP_CLASSES

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

◆ areTilesAndTiledDimsAllConstant()

template<typename OpTy >
bool areTilesAndTiledDimsAllConstant ( OpTy  op)

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

Definition at line 3960 of file TensorOps.cpp.

◆ 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 3826 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

◆ commonVerifierPackAndUnPackOp()

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

Definition at line 3584 of file TensorOps.cpp.

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

References mlir::Value::getDefiningOp().

◆ foldIdentityOffsetSizeAndStrideOpInterface()

static LogicalResult foldIdentityOffsetSizeAndStrideOpInterface ( OffsetSizeAndStrideOpInterface  op,
ShapedType  shapedType 
)
static

Definition at line 2326 of file TensorOps.cpp.

◆ 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 2494 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 2474 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 3508 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 131 of file TensorOps.cpp.

References mlir::detail::enumerate().

◆ getMixedTilesImpl()

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

Definition at line 3523 of file TensorOps.cpp.

◆ getNumElements()

static int64_t getNumElements ( ShapedType  type)
static

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

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

◆ getStaticTilesImpl()

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

Definition at line 3539 of file TensorOps.cpp.

◆ hasSameInnerOuterAttribute()

static bool hasSameInnerOuterAttribute ( PackOp  packOp,
UnPackOp  unPackOp 
)
static

Definition at line 3991 of file TensorOps.cpp.

◆ haveSameTiles()

static bool haveSameTiles ( PackOp  packOp,
UnPackOp  unPackOp 
)
static

Definition at line 3999 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 4026 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 4260 of file TensorOps.cpp.

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

◆ isLikePadUnPad()

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

Definition at line 4110 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 383 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 4012 of file TensorOps.cpp.

◆ parseInferType()

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

Definition at line 2754 of file TensorOps.cpp.

References mlir::success().

◆ printInferType()

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

Definition at line 2750 of file TensorOps.cpp.

◆ produceSliceErrorMsg()

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

Definition at line 2053 of file TensorOps.cpp.

◆ reifyResultShapesImpl()

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

Definition at line 3496 of file TensorOps.cpp.

◆ 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 1074 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,
int64_t  rank,
StringRef  gatherOrScatter,
StringRef  sourceOrDest 
)
static

Definition at line 1288 of file TensorOps.cpp.

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