MLIR  17.0.0git
Classes | Typedefs | Functions
mlir::tensor Namespace Reference

Classes

class  ExtractSliceFromCollapseHelper
 This class assists with generating IR required to materialize an arbitrary-sized slice from the result of a CollapseShapeOp. More...
 

Typedefs

using ControlConstantExtractSliceFusionFn = std::function< bool(ExtractSliceOp)>
 Function to control the folding of constant and extract slice. More...
 

Functions

bool preservesStaticInformation (Type source, Type target)
 Returns true if target is a ranked tensor type that preserves static information available in the source ranked tensor type. More...
 
bool canFoldIntoConsumerOp (CastOp castOp)
 Determines whether tensor::CastOp casts to a more dynamic version of the source tensor. More...
 
bool canFoldIntoProducerOp (CastOp castOp)
 Determines whether the tensor::CastOp casts to a more static version of the source tensor. More...
 
LogicalResult foldTensorCast (Operation *op)
 Performs folding of any operand of op if it comes from a tensor::CastOp that can be folded. More...
 
SmallVector< OpFoldResultgetMixedSizes (OpBuilder &builder, Location loc, Value value)
 Return the dimensions of the given tensor value. More...
 
Value createCanonicalRankReducingExtractSliceOp (OpBuilder &b, Location loc, Value tensor, RankedTensorType targetType)
 Create a rank-reducing ExtractSliceOp @[0 . More...
 
Value createCanonicalRankReducingInsertSliceOp (OpBuilder &b, Location loc, Value tensor, Value dest)
 Create a rank-reducing InsertSliceOp @[0 . More...
 
FailureOr< ValuegetOrCreateDestination (OpBuilder &b, Location loc, OpResult opResult)
 This is a helper function for DestinationStyleOpInterface. More...
 
LogicalResult getOrCreateDestinations (OpBuilder &b, Location loc, Operation *op, SmallVector< Value > &result)
 This is a helper function for DestinationStyleOpInterface. More...
 
void populateFoldConstantExtractSlicePatterns (RewritePatternSet &patterns, const ControlConstantExtractSliceFusionFn &controlFn=[](ExtractSliceOp op) { return false;})
 Patterns to fold the extract slice op with its constant operand. More...
 
void populateSimplifyTensorPack (RewritePatternSet &patterns)
 Patterns to simplify tensor.pack. More...
 
void registerInferTypeOpInterfaceExternalModels (mlir::DialectRegistry &registry)
 Registers external models for Infer Type interfaces for tensor ops. More...
 
OperationbubbleUpPadSlice (OpBuilder &b, tensor::PadOp padOp, ArrayRef< OpFoldResult > offsets, ArrayRef< OpFoldResult > sizes, bool generateZeroSliceGuard=true)
 Bubbles up a slice of this pad by taking the slice first and then performing the padding. More...
 
void registerTilingInterfaceExternalModels (mlir::DialectRegistry &registry)
 Registers external models for Tiling interface for tensor ops. More...
 
void registerTilingInterfaceExternalModelsForPackUnPackOps (DialectRegistry &registry)
 Similar to the above registeration, but it is only for tensor.pack and tensor.unpack ops. More...
 
void registerBufferizableOpInterfaceExternalModels (DialectRegistry &registry)
 
void populateSplitPaddingPatterns (RewritePatternSet &patterns, PatternBenefit baseBenefit=1)
 Populates patterns with patterns to wrap a tensor.pad op with an scf.if op to separate the cases where we don't need padding (all pad sizes are actually zeros) and where we indeed need padding. More...
 
FailureOr< ValuereplaceExtractSliceWithTiledProducer (OpBuilder &builder, tensor::ExtractSliceOp sliceOp, OpResult producerOp)
 Pattern to swap an tensor.extract_slice with its producer when the producer implements the TilingInterface. More...
 
void populateMergeConsecutiveInsertExtractSlicePatterns (RewritePatternSet &patterns)
 Collects patterns to merge consecutive tensor.insert_slice/extract_slice into one. More...
 
void populateReassociativeReshapeFoldingPatterns (RewritePatternSet &patterns)
 Populates patterns with patterns that fold tensor.expand_shape and tensor.collapse_shape into other ops. More...
 
void populateFoldTensorEmptyPatterns (RewritePatternSet &patterns)
 Populates patterns with patterns that fold tensor.empty with tensor. More...
 
void populateFoldIntoPackAndUnpackPatterns (RewritePatternSet &patterns)
 Populates patterns with patterns that fold operations like tensor.pad and tensor.extract_slice into tensor.pack and tensor.unpack operations respectively. More...
 
FailureOr< Operation * > simplifyCollapseShapeWithRankReducingExtractSlice (tensor::CollapseShapeOp op, RewriterBase &rewriter)
 Tries to simplify a tensor.collapse_shape operation by inserting a single rank-reducing tensor.extract_slice operation. More...
 
PadOp createPadHighOp (RankedTensorType type, Value source, Value pad, bool nofold, Location loc, OpBuilder &builder)
 
SmallVector< ValuecreateDynamicDimValues (OpBuilder &b, Location loc, Value rankedTensor)
 
SmallVector< OpFoldResultcreateDimValues (OpBuilder &b, Location loc, Value rankedTensor)
 

Typedef Documentation

◆ ControlConstantExtractSliceFusionFn

using mlir::tensor::ControlConstantExtractSliceFusionFn = typedef std::function<bool(ExtractSliceOp)>

Function to control the folding of constant and extract slice.

Definition at line 146 of file Tensor.h.

Function Documentation

◆ bubbleUpPadSlice()

Operation * mlir::tensor::bubbleUpPadSlice ( OpBuilder b,
tensor::PadOp  padOp,
ArrayRef< OpFoldResult offsets,
ArrayRef< OpFoldResult sizes,
bool  generateZeroSliceGuard = true 
)

Bubbles up a slice of this pad by taking the slice first and then performing the padding.

offsets and strides specifies each dimension's start offset and size for the slice. The slice has unit strides along all dimensions.

Specifically, this function converts:

%0 = tensor.pad %source low[...] high[...] { linalg.yield %cst }
%1 = <extract-slice> %0 offsets=[...], sizes[...]

into

%0 = tensor.extract_slice %source ...
%0 = tensor.pad %0 low[...] high[...] { linalg.yield %cst }

If generateZeroSliceGuard is true, the generated IR will contain logic to guard against the case that we might take a zero-sized slice from the original source. For such cases, we tensor.generate to generate the full tensor.

Definition at line 441 of file TensorTilingInterfaceImpl.cpp.

References mlir::bindDims(), mlir::OpBuilder::create(), mlir::OpBuilder::createOrFold(), mlir::dispatchIndexOpFoldResults(), mlir::AffineMap::get(), mlir::getAsOpFoldResult(), mlir::getConstantIntValue(), mlir::Builder::getContext(), mlir::Builder::getIndexAttr(), mlir::AffineMap::getMultiDimIdentityMap(), mlir::getValueOrCreateConstantIndexOp(), max(), and min().

Referenced by mlir::linalg::ExtractSliceOfPadTensorSwapPattern::matchAndRewrite().

◆ canFoldIntoConsumerOp()

bool mlir::tensor::canFoldIntoConsumerOp ( CastOp  castOp)

Determines whether tensor::CastOp casts to a more dynamic version of the source tensor.

This is useful to fold a tensor.cast into a consuming op and implement canonicalization patterns for ops in different dialects that may consume the results of tensor.cast operations. Such foldable tensor.cast operations are typically inserted as extract_slice ops and are canonicalized, to preserve the type compatibility of their uses.

Returns true when all conditions are met:

  1. source and result are ranked tensors with same element type and rank.
  2. the tensor type has more static information than the result

Example:

%1 = tensor.cast %0 : tensor<8x16xf32> to tensor<?x?xf32>
%2 = consumer %1 ... : tensor<?x?xf32> ...

folds into:

%2 = consumer %0 ... : tensor<8x16xf32> ...

This is useful to fold a tensor.cast into a consuming op and implement canonicalization patterns for ops in different dialects that may consume the results of tensor.cast operations. Such foldable tensor.cast operations are typically inserted as slice ops and are canonicalized, to preserve the type compatibility of their uses.

Returns true when all conditions are met:

  1. source and result are ranked tensors with same element type and rank.
  2. the tensor type has more static information than the result

Example:

%1 = tensor.cast %0 : tensor<8x16xf32> to tensor<?x?xf32>
%2 = consumer %1 ... : tensor<?x?xf32> ...

folds into:

%2 = consumer %0 ... : tensor<8x16xf32> ...

Definition at line 175 of file TensorOps.cpp.

References preservesStaticInformation().

Referenced by foldTensorCast(), CollapseShapeOpMemRefCastFolder::matchAndRewrite(), and FoldTensorCastProducerOp::matchAndRewrite().

◆ canFoldIntoProducerOp()

bool mlir::tensor::canFoldIntoProducerOp ( CastOp  castOp)

Determines whether the tensor::CastOp casts to a more static version of the source tensor.

This is useful to fold into a producing op and implement canonicaliation patterns with the tensor.cast op as the root, but producer being from different dialects. Returns true when all conditions are met:

  1. source and result and ranked tensors with same element type and rank.
  2. the result type has more static information than the source.

Example:

%1 = producer ... : tensor<?x?xf32>
%2 = tensor.cast %1 : tensor<?x?xf32> to tensor<8x16xf32>

can be canonicalized to :

%2 = producer ... : tensor<8x16xf32>

Not all ops might be canonicalizable this way, but for those that can be, this method provides a check that it is worth doing the canonicalization.

Definition at line 205 of file TensorOps.cpp.

References preservesStaticInformation().

◆ createCanonicalRankReducingExtractSliceOp()

Value mlir::tensor::createCanonicalRankReducingExtractSliceOp ( OpBuilder b,
Location  loc,
Value  tensor,
RankedTensorType  targetType 
)

Create a rank-reducing ExtractSliceOp @[0 .

. 0] with strides [1 .. 1] and appropriate sizes (i.e. tensor.getSizes()) to reduce the rank of tensor to that of targetType.

Definition at line 2077 of file TensorOps.cpp.

References mlir::Type::cast(), mlir::OpBuilder::createOrFold(), mlir::Builder::getIndexAttr(), getMixedSizes(), and mlir::Value::getType().

Referenced by mlir::linalg::DownscaleSizeOneWindowed2DConvolution< Conv2DOp, Conv1DOp >::returningMatchAndRewrite(), and mlir::linalg::DownscaleDepthwiseConv2DNhwcHwcOp::returningMatchAndRewrite().

◆ createCanonicalRankReducingInsertSliceOp()

Value mlir::tensor::createCanonicalRankReducingInsertSliceOp ( OpBuilder b,
Location  loc,
Value  tensor,
Value  dest 
)

Create a rank-reducing InsertSliceOp @[0 .

. 0] with strides [1 .. 1] and appropriate sizes (i.e. dest.getSizes()). The result is a new tensor with rank increased to that of dest, obtained by inserting tensor into dest at the canonical [0 .. 0] position.

Definition at line 2428 of file TensorOps.cpp.

References mlir::Type::cast(), mlir::OpBuilder::createOrFold(), mlir::Builder::getIndexAttr(), getMixedSizes(), and mlir::Value::getType().

Referenced by mlir::linalg::DownscaleSizeOneWindowed2DConvolution< Conv2DOp, Conv1DOp >::returningMatchAndRewrite(), and mlir::linalg::DownscaleDepthwiseConv2DNhwcHwcOp::returningMatchAndRewrite().

◆ createDimValues()

SmallVector< OpFoldResult > mlir::tensor::createDimValues ( OpBuilder b,
Location  loc,
Value  rankedTensor 
)

◆ createDynamicDimValues()

SmallVector< Value > mlir::tensor::createDynamicDimValues ( OpBuilder b,
Location  loc,
Value  rankedTensor 
)

◆ createPadHighOp()

PadOp mlir::tensor::createPadHighOp ( RankedTensorType  type,
Value  source,
Value  pad,
bool  nofold,
Location  loc,
OpBuilder builder 
)

◆ foldTensorCast()

LogicalResult mlir::tensor::foldTensorCast ( Operation op)

Performs folding of any operand of op if it comes from a tensor::CastOp that can be folded.

Definition at line 214 of file TensorOps.cpp.

References canFoldIntoConsumerOp(), mlir::Operation::getOpOperands(), and mlir::success().

◆ getMixedSizes()

SmallVector< OpFoldResult > mlir::tensor::getMixedSizes ( OpBuilder builder,
Location  loc,
Value  value 
)

◆ getOrCreateDestination()

FailureOr< Value > mlir::tensor::getOrCreateDestination ( OpBuilder b,
Location  loc,
OpResult  opResult 
)

This is a helper function for DestinationStyleOpInterface.

If there is a destination operand for the given OpResult, return that operand. Otherwise, return an empty tensor (tensor.empty) with the shape of the OpResult. Dynamic dimensions are queried via ReifyRankedShapedTypeOpInterface.

Definition at line 64 of file TensorOps.cpp.

References mlir::OpBuilder::create(), mlir::Type::dyn_cast(), mlir::failed(), mlir::failure(), mlir::getAsOpFoldResult(), mlir::Value::getDefiningOp(), mlir::Builder::getIndexAttr(), mlir::OpResult::getResultNumber(), mlir::Value::getType(), and mlir::OpBuilder::setInsertionPoint().

Referenced by getOrCreateDestinations(), and mlir::scf::yieldReplacementForFusedProducer().

◆ getOrCreateDestinations()

LogicalResult mlir::tensor::getOrCreateDestinations ( OpBuilder b,
Location  loc,
Operation op,
SmallVector< Value > &  result 
)

◆ populateFoldConstantExtractSlicePatterns()

void mlir::tensor::populateFoldConstantExtractSlicePatterns ( RewritePatternSet patterns,
const ControlConstantExtractSliceFusionFn controlFn = [](ExtractSliceOp op) {   return false; } 
)

Patterns to fold the extract slice op with its constant operand.

Definition at line 1989 of file TensorOps.cpp.

References mlir::RewritePatternSet::add(), and mlir::RewritePatternSet::getContext().

◆ populateFoldIntoPackAndUnpackPatterns()

void mlir::tensor::populateFoldIntoPackAndUnpackPatterns ( RewritePatternSet patterns)

Populates patterns with patterns that fold operations like tensor.pad and tensor.extract_slice into tensor.pack and tensor.unpack operations respectively.

Definition at line 81 of file FoldIntoPackAndUnpackPatterns.cpp.

References mlir::RewritePatternSet::getContext(), and mlir::RewritePatternSet::insert().

◆ populateFoldTensorEmptyPatterns()

void mlir::tensor::populateFoldTensorEmptyPatterns ( RewritePatternSet patterns)

Populates patterns with patterns that fold tensor.empty with tensor.

[extract_slice|cast|expand_shape|collapse_shape].

Definition at line 73 of file EmptyOpPatterns.cpp.

References mlir::RewritePatternSet::add(), and mlir::RewritePatternSet::getContext().

Referenced by mlir::linalg::populateFoldUnitExtentDimsViaReshapesPatterns(), and mlir::linalg::populateFoldUnitExtentDimsViaSlicesPatterns().

◆ populateMergeConsecutiveInsertExtractSlicePatterns()

void mlir::tensor::populateMergeConsecutiveInsertExtractSlicePatterns ( RewritePatternSet patterns)

Collects patterns to merge consecutive tensor.insert_slice/extract_slice into one.

These patterns are in in this separate entry point because the bufferization is sensitive over IR structure, particularly those tensor.extract_slice and tensor.insert_slice ops for creating the slices.

Definition at line 79 of file MergeConsecutiveInsertExtractSlicePatterns.cpp.

References mlir::RewritePatternSet::add(), and mlir::RewritePatternSet::getContext().

◆ populateReassociativeReshapeFoldingPatterns()

void mlir::tensor::populateReassociativeReshapeFoldingPatterns ( RewritePatternSet patterns)

Populates patterns with patterns that fold tensor.expand_shape and tensor.collapse_shape into other ops.

Definition at line 84 of file ReshapePatterns.cpp.

References mlir::RewritePatternSet::add(), and mlir::RewritePatternSet::getContext().

◆ populateSimplifyTensorPack()

void mlir::tensor::populateSimplifyTensorPack ( RewritePatternSet patterns)

Patterns to simplify tensor.pack.

Definition at line 3157 of file TensorOps.cpp.

References mlir::RewritePatternSet::add(), and mlir::RewritePatternSet::getContext().

◆ populateSplitPaddingPatterns()

void mlir::tensor::populateSplitPaddingPatterns ( RewritePatternSet patterns,
PatternBenefit  baseBenefit = 1 
)

Populates patterns with patterns to wrap a tensor.pad op with an scf.if op to separate the cases where we don't need padding (all pad sizes are actually zeros) and where we indeed need padding.

Definition at line 92 of file SplitPaddingPatterns.cpp.

References mlir::RewritePatternSet::add(), and mlir::RewritePatternSet::getContext().

◆ preservesStaticInformation()

bool mlir::tensor::preservesStaticInformation ( Type  source,
Type  target 
)

Returns true if target is a ranked tensor type that preserves static information available in the source ranked tensor type.

Definition at line 127 of file TensorOps.cpp.

References mlir::Type::dyn_cast().

Referenced by canFoldIntoConsumerOp(), and canFoldIntoProducerOp().

◆ registerBufferizableOpInterfaceExternalModels()

void mlir::tensor::registerBufferizableOpInterfaceExternalModels ( DialectRegistry registry)

◆ registerInferTypeOpInterfaceExternalModels()

void mlir::tensor::registerInferTypeOpInterfaceExternalModels ( mlir::DialectRegistry registry)

Registers external models for Infer Type interfaces for tensor ops.

Currently, it registers:

  • ReifyRankedShapedTypeOpInterface for tensor.collapse_shape.
  • ReifyRankedShapedTypeOpInterface for tensor.expand_shape.

Unfortunately, a "normal" internal registration is not possible at the moment, because of the dependency of the interface implementation for these ops on affine.apply and Affine dialect already depends on TensorOps. In order to break the cyclic dependency (TensorOps->AffineOps->TensorOps) the implementation is moved to a separate library.

Definition at line 196 of file TensorInferTypeOpInterfaceImpl.cpp.

References mlir::DialectRegistry::addExtension().

Referenced by mlir::registerAllDialects().

◆ registerTilingInterfaceExternalModels()

void mlir::tensor::registerTilingInterfaceExternalModels ( mlir::DialectRegistry registry)

Registers external models for Tiling interface for tensor ops.

Currently, it registers:

  • TilingInterface for tensor.pad, tensor.pack, and tensor.unpack.

Unfortunately, a "normal" internal registration is not possible at the moment, because of the dependency of the interface implementation for these ops on affine.apply and Affine dialect already depends on TensorOps. In order to break the cyclic dependency (TensorOps->AffineOps->TensorOps) the implementation is moved to a separate library.

Definition at line 649 of file TensorTilingInterfaceImpl.cpp.

References mlir::DialectRegistry::addExtension().

Referenced by mlir::registerAllDialects().

◆ registerTilingInterfaceExternalModelsForPackUnPackOps()

void mlir::tensor::registerTilingInterfaceExternalModelsForPackUnPackOps ( DialectRegistry registry)

Similar to the above registeration, but it is only for tensor.pack and tensor.unpack ops.

Definition at line 658 of file TensorTilingInterfaceImpl.cpp.

References mlir::DialectRegistry::addExtension().

◆ replaceExtractSliceWithTiledProducer()

FailureOr< Value > mlir::tensor::replaceExtractSliceWithTiledProducer ( OpBuilder builder,
tensor::ExtractSliceOp  sliceOp,
OpResult  producerOp 
)

Pattern to swap an tensor.extract_slice with its producer when the producer implements the TilingInterface.

The pattern itself does not provide a mechanism to control where the application happens. With use of transform dialect that control is done within the transform dialect. Other use cases can inherit from this pattern and add necessary controls.

Definition at line 23 of file SwapExtractSliceWithProducerPatterns.cpp.

References mlir::failed(), mlir::failure(), mlir::OpResult::getOwner(), and mlir::OpResult::getResultNumber().

Referenced by mlir::scf::tileAndFuseProducerOfSlice().

◆ simplifyCollapseShapeWithRankReducingExtractSlice()

FailureOr< Operation * > mlir::tensor::simplifyCollapseShapeWithRankReducingExtractSlice ( tensor::CollapseShapeOp  op,
RewriterBase rewriter 
)

Tries to simplify a tensor.collapse_shape operation by inserting a single rank-reducing tensor.extract_slice operation.

The extract_slice op will either take the place of the source, allowing for a new, simpler collapse_shape op to replace op, or the collapse_shape op will be completely replaced by the extract_slice result. Either way, op is replaced and new new op is returned.

Example:

%result = tensor.collapse_shape %0 [[0, 1], [2, 3]]
: tensor<?x1x30x10xf32> to tensor<?x300xf32>

can be transformed to

%tmp = tensor.extract_slice %0 [0, 0, 0, 0]
[0, %dim1, 30, 30]
[1, 1, 1 1]
: tensor<?x1x30x10xf32> to tensor<?x30x10xf32>
%result = tensor.collapse_shape %tmp [[0], [1, 2]]
: tensor<?x30x10xf32> to tensor<?x300xf32>

Example:

%result = tensor.collapse_shape %1 [[0, 1], [2]]
: tensor<?x1x30xf32> to tensor<?x30xf32>

can be transformed to

%result = tensor.extract_slice %1 [0, 0, 0]
[%dim2, 1, 30]
[1, 1, 1]
: tensor<?x1x30xf32> to tensor<?x30xf32>

Unsupported cases:

This transform doesn't yet support reducing the rank of the reassociation indices, which would require inserting a tensor.expand_shape op similar to the following example:

%result = tensor.collapse_shape %0 [[0, 1], [2, 3]]
: tensor<1x1x30x10xf32> to tensor<1x300xf32>

can be transformed to

%tmp = tensor.extract_slice %0 [0, 0, 0, 0]
[0, 1, 30, 30]
[1, 1, 1 1]
: tensor<1x1x30x10xf32> to tensor<30x10xf32>
%result0 = tensor.collapse_shape %tmp [[0, 1]]
: tensor<30x10xf32> to tensor<300xf32>
%result1 = tensor.expand_shape %tmp [[0, 1], [2]] :... tensor<1x300xf32>

Definition at line 185 of file ExtractSliceFromReshapeUtils.cpp.

References mlir::OpBuilder::create(), mlir::failed(), mlir::failure(), mlir::Builder::getIndexAttr(), getShapeDimSizes(), mlir::RewriterBase::replaceOp(), and mlir::RewriterBase::replaceOpWithNewOp().