Namespaces
namespace	impl

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.
using	ControlFoldFn = std::function<bool(OpOperand *)>

Functions
void	registerAllExtensions (DialectRegistry &registry)
	Register all extensions of the Tensor dialect.
void	registerShardingInterfaceExternalModels (DialectRegistry &registry)
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.
bool	canFoldIntoConsumerOp (CastOp castOp)
	Determines whether tensor::CastOp casts to a more dynamic version of the source tensor.
bool	canFoldIntoProducerOp (CastOp castOp)
	Determines whether the tensor::CastOp casts to a more static version of the source tensor.
bool	hasFoldableTensorCastOperand (Operation *op)
	Return true if any of the operands of op is a CastOp that can be folded into its consumer, i.e.
SmallVector< Value >	getUpdatedOperandsAfterCastOpFolding (DestinationStyleOpInterface op, SmallVector< Type > &newResTy)
	Assuming that op contains at least one operand that is a foldable CastOp (i.e.
LogicalResult	foldTensorCast (Operation *op)
	Performs folding of any operand of op if it comes from a tensor::CastOp that can be folded.
OpFoldResult	getMixedSize (OpBuilder &builder, Location loc, Value value, int64_t dim)
	Return the dimension of the given tensor value.
SmallVector< OpFoldResult >	getMixedSizes (OpBuilder &builder, Location loc, Value value)
	Return the dimensions of the given tensor value.
Value	createCanonicalRankReducingExtractSliceOp (OpBuilder &b, Location loc, Value tensor, RankedTensorType targetType)
	Create a rank-reducing ExtractSliceOp @[0 .
Value	createCanonicalRankReducingInsertSliceOp (OpBuilder &b, Location loc, Value tensor, Value dest)
	Create a rank-reducing InsertSliceOp @[0 .
FailureOr< Value >	getOrCreateDestination (OpBuilder &b, Location loc, OpResult opResult)
	This is a helper function for DestinationStyleOpInterface.
LogicalResult	getOrCreateDestinations (OpBuilder &b, Location loc, Operation *op, SmallVector< Value > &result)
	This is a helper function for DestinationStyleOpInterface.
bool	isSameTypeWithoutEncoding (Type tp1, Type tp2)
	Tests if types are the same when ignoring encoding on ranked tensors.
void	populateFoldConstantExtractSlicePatterns (RewritePatternSet &patterns, const ControlConstantExtractSliceFusionFn &controlFn=[](ExtractSliceOp op) { return false;})
	Patterns to fold the extract slice op with its constant operand.
void	populateFoldCollapseExtractPatterns (RewritePatternSet &patterns)
	Patterns to fold extracts of a collapse_shaped tensor to an extract of the source tensor.
void	registerInferTypeOpInterfaceExternalModels (mlir::DialectRegistry &registry)
	Registers external models for Infer Type interfaces for tensor ops.
FailureOr< TilingResult >	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.
void	registerTilingInterfaceExternalModels (mlir::DialectRegistry &registry)
	Registers external models for Tiling interface for tensor ops.
void	registerTilingInterfaceExternalModelsForPackUnPackOps (DialectRegistry &registry)
	Similar to the above registeration, but it is only for tensor.pack and tensor.unpack ops.
void	registerValueBoundsOpInterfaceExternalModels (DialectRegistry &registry)
void	registerTransformDialectExtension (DialectRegistry &registry)
void	registerFindPayloadReplacementOpInterfaceExternalModels (DialectRegistry &registry)
void	registerBufferizableOpInterfaceExternalModels (DialectRegistry &registry)
std::unique_ptr<::mlir::Pass >	createFoldTensorSubsetOpsPass ()
void	registerFoldTensorSubsetOpsPass ()
void	registerFoldTensorSubsetOpsPassPass ()
void	registerTensorPasses ()
void	registerRuntimeVerifiableOpInterfaceExternalModels (DialectRegistry &registry)
void	registerSubsetOpInterfaceExternalModels (DialectRegistry &registry)
FailureOr< TilingResult >	replaceExtractSliceWithTiledProducer (OpBuilder &builder, tensor::ExtractSliceOp sliceOp, OpResult producerOp)
	Method to swap an tensor.extract_slice with its producer when the producer implements the TilingInterface.
FailureOr< TilingResult >	replaceInsertSlicesWithTiledConsumer (OpBuilder &builder, ArrayRef< tensor::InsertSliceOp > sliceOps, ArrayRef< OpOperand * > consumerOperands)
	Method to swap tensor.insert_slices with their consumers when the consumer implements the TilingInterface.
void	populateFoldTensorSubsetOpPatterns (RewritePatternSet &patterns)
	Appends patterns for folding tensor subset ops into consumer load/store ops into patterns.
void	populateFoldTensorSubsetIntoVectorTransferPatterns (RewritePatternSet &patterns)
	Appends patterns for folding tensor subset ops into vector transfer ops.
void	populateMergeConsecutiveInsertExtractSlicePatterns (RewritePatternSet &patterns)
	Collects patterns to merge consecutive tensor.insert_slice/extract_slice into one.
void	populateBubbleUpExtractSliceOpPatterns (RewritePatternSet &patterns)
	Appends patterns that are used to bubble up tensor.extract slice op above its producer.
void	populateDropRedundantInsertSliceRankExpansionPatterns (RewritePatternSet &patterns)
	Populates patterns with patterns that drop redundant tensor.insert_slice rank expansions.
void	populateReassociativeReshapeFoldingPatterns (RewritePatternSet &patterns)
	Populates patterns with patterns that fold tensor.expand_shape and tensor.collapse_shape into other ops.
void	populateBubbleUpExpandShapePatterns (RewritePatternSet &patterns)
	Populates patterns with patterns that bubble up tensor.expand_shape through tensor.collapse_shape ops.
void	populateFoldTensorEmptyPatterns (RewritePatternSet &patterns, bool foldSingleUseOnly=false)
	Populates patterns with patterns that fold tensor.empty with its consumers.
void	populateDecomposeTensorConcatPatterns (RewritePatternSet &patterns)
	Populates patterns with patterns that decompose tensor.concat into tensor.empty of a tensor of the concatenated size, followed by a chain of tensor.insert_slice operations on the inputs.
void	populateRewriteAsConstantPatterns (RewritePatternSet &patterns, const ControlFoldFn &controlFn)
	Populates patterns with patterns that replace tensor ops (such as tensor.generate) with constants when possible.
FailureOr< Value >	buildIndependentOp (OpBuilder &b, tensor::PadOp padOp, ValueRange independencies)
	Build a new tensor::PadOp with low/high padding that is independent of all given independencies.
FailureOr< Value >	buildIndependentOp (OpBuilder &b, tensor::EmptyOp emptyOp, ValueRange independencies)
	Build a new tensor::EmptyOp who's dynamic sizes are independent of all given independencies.
LogicalResult	getCollapsedExtractSliceInfo (OpBuilder &b, tensor::ExtractSliceOp sliceOp, ArrayRef< ReassociationIndices > reassociation, SmallVectorImpl< OpFoldResult > &collapsedOffsets, SmallVectorImpl< OpFoldResult > &collapsedSizes, SmallVectorImpl< OpFoldResult > &collapsedStrides)
	Computes the offsets, sizes, and strides needed to build a collapsed sliceOp.
LogicalResult	getExpandedExtractSliceInfo (OpBuilder &b, tensor::ExtractSliceOp sliceOp, ArrayRef< ReassociationIndices > reassociation, ArrayRef< int64_t > expandedShape, SmallVectorImpl< OpFoldResult > &expandedOffsets, SmallVectorImpl< OpFoldResult > &expandedSizes, SmallVectorImpl< OpFoldResult > &expandedStrides)
	Computes the offsets, sizes, and strides needed to build an expanded sliceOp.
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.
PadOp	createPadHighOp (RankedTensorType resType, Value source, Value pad, bool nofold, Location loc, OpBuilder &builder, ValueRange dynOutDims={})
SmallVector< Value >	createDynamicDimValues (OpBuilder &b, Location loc, Value rankedTensor)
FailureOr< RankedTensorType >	computeTransposedType (RankedTensorType rankedTensorType, ArrayRef< int64_t > transposeVector)
	Returns the transposed rankedTensorType if transposeVector is non-empty.
CollapseShapeOp	dropGivenUnitDims (OpBuilder &b, Location loc, Value src, const llvm::SmallBitVector &dropDims)
	Create tensor.collapse_shape to drop unit dimensions in dropDims in tensor src.
bool	isCastLikeInsertSliceOp (InsertSliceOp op)
	A tensor.insert_slice is a cast-like operation if it merely rank-extends the source tensor or inserts the source tensor into a destination tensor with the same shape.
bool	isCastLikeExtractSliceOp (ExtractSliceOp op)
	A tensor.extract_slice is a cast-like operation if it merely rank-reduces unit dimensions of the source tensor or extracts the entire source tensor.

Typedef Documentation

◆ ControlConstantExtractSliceFusionFn

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

Function to control the folding of constant and extract slice.

Definition at line 167 of file Tensor.h.

◆ ControlFoldFn

using mlir::tensor::ControlFoldFn = std::function<bool(OpOperand *)>

Definition at line 99 of file Transforms.h.

Function Documentation

◆ bubbleUpPadSlice()

FailureOr< TilingResult > 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[...]

mlir::linalg

Definition LinalgToStandard.h:24

mlir::tensor

Definition BufferizationTransformOps.h:19

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 88 of file TensorTilingInterfaceImpl.cpp.

References b, mlir::bindDims(), mlir::dispatchIndexOpFoldResults(), mlir::AffineMap::get(), getMixedSize(), mlir::AffineMap::getMultiDimIdentityMap(), mlir::getPrunedAttributeList(), mlir::Operation::getResult(), mlir::getValueOrCreateConstantIndexOp(), indices, mlir::isZeroInteger(), mlir::affine::makeComposedFoldedAffineApply(), mlir::affine::makeComposedFoldedAffineMax(), mlir::affine::makeComposedFoldedAffineMin(), max(), min(), and result.

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

◆ buildIndependentOp() [1/2]

FailureOr< Value > mlir::tensor::buildIndependentOp	(	OpBuilder &	b,
		tensor::EmptyOp	emptyOp,
		ValueRange	independencies )

Build a new tensor::EmptyOp who's dynamic sizes are independent of all given independencies.

If the op is already independent of all independencies, the same EmptyOp result is returned.

Failure indicates the no suitable upper bound for the dynamic sizes could be found.

Definition at line 107 of file IndependenceTransforms.cpp.

References b, and makeIndependent().

◆ buildIndependentOp() [2/2]

FailureOr< Value > mlir::tensor::buildIndependentOp	(	OpBuilder &	b,
		tensor::PadOp	padOp,
		ValueRange	independencies )

Build a new tensor::PadOp with low/high padding that is independent of all given independencies.

If the op is already independent of all independencies, the same PadOp result is returned.

Failure indicates the no suitable upper bound for low/high padding could be found.

Example: scf.for iv = lb to ub step step { high = affine.apply affine_map<(d0)[s0] -> (s0 - d0)> (i)[ub] p = tensor.pad t low[5] high[high] ... ... }

The function builds IR such as: high_new = affine.apply affine_map<()[s0, s1] -> (-s0 + s1)> ()[lb, ub] p_hoistable = tensor.pad t low[5] high[high_new] dim = tensor.dim t, c0 size = affine.apply affine_map<(d0)[s0, s1] -> (-d0 + s0 + s1 + 5)> (iv)[ub, dim] slice = tensor.extract_slice p_hoistable [0] [size] [1]

The slice is returned.

Definition at line 36 of file IndependenceTransforms.cpp.

References b, and makeIndependent().

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

source and result are ranked tensors with same element type and rank.
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:

source and result are ranked tensors with same element type and rank.
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 318 of file TensorOps.cpp.

References preservesStaticInformation().

Referenced by foldTensorCast(), and getUpdatedOperandsAfterCastOpFolding().

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

source and result and ranked tensors with same element type and rank.
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.

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

source and result and ranked tensors with same element type and rank.
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 349 of file TensorOps.cpp.

References preservesStaticInformation().

◆ computeTransposedType()

FailureOr< RankedTensorType > mlir::tensor::computeTransposedType	(	RankedTensorType	rankedTensorType,
		ArrayRef< int64_t >	transposeVector )

Returns the transposed rankedTensorType if transposeVector is non-empty.

Fail if transposeVector is not a permutation matching the tensor rank.

Definition at line 76 of file Utils.cpp.

References mlir::applyPermutationToVector(), mlir::isPermutationVector(), and mlir::RankedTensorType::Builder::setShape().

Referenced by buildPackingLoopNestImpl(), and mlir::linalg::hoistPaddingOnTensors().

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

References b, and getMixedSizes().

Referenced by mlir::linalg::ExtractSliceOfPadTensorSwapPattern::matchAndRewrite(), mlir::linalg::DownscaleConv2DOp::returningMatchAndRewrite(), mlir::linalg::DownscaleDepthwiseConv2DNhwcHwcOp::returningMatchAndRewrite(), and mlir::linalg::DownscaleSizeOneWindowed2DConvolution< Conv2DOp, Conv1DOp >::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 3177 of file TensorOps.cpp.

References b, getMixedSizes(), and mlir::Value::getType().

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

◆ createDynamicDimValues()

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

Definition at line 62 of file Utils.cpp.

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

Referenced by mlir::linalg::splitReductionByScaling().

◆ createFoldTensorSubsetOpsPass()

std::unique_ptr<::mlir::Pass > mlir::tensor::createFoldTensorSubsetOpsPass ( )

We declare an explicit private instantiation because Pass classes should only be visible by the current library.

Definition at line 81 of file FoldTensorSubsetOps.cpp.

References mlir::RewriterBase::notifyMatchFailure().

◆ createPadHighOp()

PadOp mlir::tensor::createPadHighOp	(	RankedTensorType	resType,
		Value	source,
		Value	pad,
		bool	nofold,
		Location	loc,
		OpBuilder &	builder,
		ValueRange	dynOutDims = {} )

Definition at line 23 of file Utils.cpp.

References b, mlir::bindDims(), getMixedSize(), and mlir::affine::makeComposedFoldedAffineApply().

Referenced by getPackOpSourceOrPaddedSource(), and mlir::linalg::makeComposedPadHighOp().

◆ dropGivenUnitDims()

CollapseShapeOp mlir::tensor::dropGivenUnitDims	(	OpBuilder &	b,
		Location	loc,
		Value	src,
		const llvm::SmallBitVector &	dropDims )

Create tensor.collapse_shape to drop unit dimensions in dropDims in tensor src.

Definition at line 95 of file Utils.cpp.

References b, mlir::dropDims(), mlir::Value::getType(), and setBit().

Referenced by mlir::linalg::fuseProducerOfTensor().

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

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

◆ getCollapsedExtractSliceInfo()

LogicalResult mlir::tensor::getCollapsedExtractSliceInfo	(	OpBuilder &	b,
		tensor::ExtractSliceOp	sliceOp,
		ArrayRef< ReassociationIndices >	reassociation,
		SmallVectorImpl< OpFoldResult > &	collapsedOffsets,
		SmallVectorImpl< OpFoldResult > &	collapsedSizes,
		SmallVectorImpl< OpFoldResult > &	collapsedStrides )

Computes the offsets, sizes, and strides needed to build a collapsed sliceOp.

The dimensions to collapse are specified by reassociation.

This fails when the specified collapse cannot be represented by a valid ExtractSliceOp.

Definition at line 463 of file ReshapePatterns.cpp.

References mlir::ValueBoundsConstraintSet::areEqual(), b, mlir::bindDims(), mlir::AffineMap::get(), getMixedSizes(), mlir::getValueOrCreateConstantIndexOp(), indices, mlir::isOneInteger(), mlir::isZeroInteger(), mlir::affine::makeComposedFoldedAffineApply(), mul, and success().

◆ getExpandedExtractSliceInfo()

LogicalResult mlir::tensor::getExpandedExtractSliceInfo	(	OpBuilder &	b,
		tensor::ExtractSliceOp	sliceOp,
		ArrayRef< ReassociationIndices >	reassociation,
		ArrayRef< int64_t >	expandedShape,
		SmallVectorImpl< OpFoldResult > &	expandedOffsets,
		SmallVectorImpl< OpFoldResult > &	expandedSizes,
		SmallVectorImpl< OpFoldResult > &	expandedStrides )

Computes the offsets, sizes, and strides needed to build an expanded sliceOp.

The dimensions to expand are specified by reassociation and expandedShape.

This fails when the specified expansion cannot be represented by a valid ExtractSliceOp.

Definition at line 583 of file ReshapePatterns.cpp.

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

◆ getMixedSize()

OpFoldResult mlir::tensor::getMixedSize	(	OpBuilder &	builder,
		Location	loc,
		Value	value,
		int64_t	dim )

Return the dimension of the given tensor value.

Definition at line 57 of file TensorOps.cpp.

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

Referenced by bubbleUpPadSlice(), createPadHighOp(), getMixedSizes(), mlir::linalg::lowerPack(), mlir::linalg::DecomposeOuterUnitDimsUnPackOpPattern::matchAndRewrite(), and mlir::linalg::DecomposePadOpPattern::matchAndRewrite().

◆ getMixedSizes()

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

Return the dimensions of the given tensor value.

Definition at line 66 of file TensorOps.cpp.

References getMixedSize(), mlir::Value::getType(), and result.

Referenced by mlir::tensor::ExtractSliceFromCollapseHelper::create(), createCanonicalRankReducingExtractSliceOp(), createCanonicalRankReducingInsertSliceOp(), mlir::memref::createCanonicalRankReducingSubViewOp(), mlir::linalg::createDestinationPassingStyleInitOperand(), mlir::vector::createReadOrMaskedRead(), expandValue(), fuseWithReshapeByExpansion(), getCollapsedExtractSliceInfo(), mlir::vector::getMixedSizesXfer(), mlir::linalg::lowerPack(), mlir::linalg::lowerUnPack(), mlir::linalg::DecomposePadOpPattern::matchAndRewrite(), mlir::linalg::reifyResultShapesImpl(), mlir::linalg::rewriteInDestinationPassingStyle(), and simplifyCollapseShapeWithRankReducingExtractSlice().

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

References b, mlir::Value::getDefiningOp(), mlir::OpResult::getResultNumber(), mlir::Value::getType(), and mlir::reifyResultShapes().

Referenced by getOrCreateDestinations().

◆ getOrCreateDestinations()

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

This is a helper function for DestinationStyleOpInterface.

Get or create destinations for every tensor OpResult of the given op.

Definition at line 110 of file TensorOps.cpp.

References b, getOrCreateDestination(), mlir::Operation::getResults(), result, and success().

Referenced by createInitialTensorsForTiling(), mlir::linalg::splitOp(), tileAndFuseFirstExtractUseThroughContainingOpBlockArgument(), and mlir::linalg::tileReductionUsingForall().

◆ getUpdatedOperandsAfterCastOpFolding()

SmallVector< Value > mlir::tensor::getUpdatedOperandsAfterCastOpFolding	(	DestinationStyleOpInterface	op,
		SmallVector< Type > &	newResTy )

Assuming that op contains at least one operand that is a foldable CastOp (i.e.

hasFoldableTensorCastOperand returns true), calculate the updated operands.

Definition at line 365 of file TensorOps.cpp.

References canFoldIntoConsumerOp(), and hasFoldableTensorCastOperand().

Referenced by FoldTensorCastProducerOp::matchAndRewrite(), mlir::linalg::FoldTensorCastPackOp::matchAndRewrite(), and mlir::linalg::FoldTensorCastUnPackOp::matchAndRewrite().

◆ hasFoldableTensorCastOperand()

bool mlir::tensor::hasFoldableTensorCastOperand ( Operation * op )

Return true if any of the operands of op is a CastOp that can be folded into its consumer, i.e.

op. This is effectively a convenience wrapper for canFoldIntoProducerOp.

Definition at line 356 of file TensorOps.cpp.

References mlir::Operation::getOpOperands().

Referenced by foldTensorCastPrecondition(), getUpdatedOperandsAfterCastOpFolding(), mlir::linalg::FoldTensorCastPackOp::matchAndRewrite(), and mlir::linalg::FoldTensorCastUnPackOp::matchAndRewrite().

◆ isCastLikeExtractSliceOp()

bool mlir::tensor::isCastLikeExtractSliceOp ( ExtractSliceOp op )

A tensor.extract_slice is a cast-like operation if it merely rank-reduces unit dimensions of the source tensor or extracts the entire source tensor.

Definition at line 149 of file Utils.cpp.

References mlir::ValueBoundsConstraintSet::areEqual().

◆ isCastLikeInsertSliceOp()

bool mlir::tensor::isCastLikeInsertSliceOp ( InsertSliceOp op )

A tensor.insert_slice is a cast-like operation if it merely rank-extends the source tensor or inserts the source tensor into a destination tensor with the same shape.

Definition at line 125 of file Utils.cpp.

References mlir::ValueBoundsConstraintSet::areEqual().

◆ isSameTypeWithoutEncoding()

bool mlir::tensor::isSameTypeWithoutEncoding	(	Type	tp1,
		Type	tp2 )

Tests if types are the same when ignoring encoding on ranked tensors.

Definition at line 124 of file TensorOps.cpp.

Referenced by verifyTensorReshapeOp().

◆ populateBubbleUpExpandShapePatterns()

void mlir::tensor::populateBubbleUpExpandShapePatterns ( RewritePatternSet & patterns )

Populates patterns with patterns that bubble up tensor.expand_shape through tensor.collapse_shape ops.

Definition at line 731 of file ReshapePatterns.cpp.

References mlir::patterns.

◆ populateBubbleUpExtractSliceOpPatterns()

void mlir::tensor::populateBubbleUpExtractSliceOpPatterns ( RewritePatternSet & patterns )

Appends patterns that are used to bubble up tensor.extract slice op above its producer.

When used as cleanup patterns of tile and fuse, enables fusing the producer with the consumer even if the producer does not implement the tiling interface.

Definition at line 736 of file ReshapePatterns.cpp.

References mlir::patterns.

◆ populateDecomposeTensorConcatPatterns()

void mlir::tensor::populateDecomposeTensorConcatPatterns ( RewritePatternSet & patterns )

Populates patterns with patterns that decompose tensor.concat into tensor.empty of a tensor of the concatenated size, followed by a chain of tensor.insert_slice operations on the inputs.

This is intended to be used as a fallback tensor -> tensor lowering that decomposes concat such that it can be bufferized into a sequence of copies.

Definition at line 46 of file ConcatOpPatterns.cpp.

References mlir::patterns.

◆ populateDropRedundantInsertSliceRankExpansionPatterns()

void mlir::tensor::populateDropRedundantInsertSliceRankExpansionPatterns ( RewritePatternSet & patterns )

Populates patterns with patterns that drop redundant tensor.insert_slice rank expansions.

Definition at line 238 of file MergeConsecutiveInsertExtractSlicePatterns.cpp.

References mlir::patterns.

◆ populateFoldCollapseExtractPatterns()

void mlir::tensor::populateFoldCollapseExtractPatterns ( RewritePatternSet & patterns )

Patterns to fold extracts of a collapse_shaped tensor to an extract of the source tensor.

Definition at line 1437 of file TensorOps.cpp.

References mlir::patterns.

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

References mlir::patterns.

◆ populateFoldTensorEmptyPatterns()

void mlir::tensor::populateFoldTensorEmptyPatterns	(	RewritePatternSet &	patterns,
		bool	foldSingleUseOnly = false )

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

If singleUseOnly is set to "true", only tensor.empty ops with a single use are folded.

Definition at line 130 of file EmptyOpPatterns.cpp.

References mlir::patterns.

Referenced by populateFoldUnitExtentDimsViaReshapesPatterns(), and populateFoldUnitExtentDimsViaSlicesPatterns().

◆ populateFoldTensorSubsetIntoVectorTransferPatterns()

void mlir::tensor::populateFoldTensorSubsetIntoVectorTransferPatterns ( RewritePatternSet & patterns )

Appends patterns for folding tensor subset ops into vector transfer ops.

Definition at line 254 of file FoldTensorSubsetOps.cpp.

References mlir::patterns.

Referenced by populateFoldTensorSubsetOpPatterns().

◆ populateFoldTensorSubsetOpPatterns()

void mlir::tensor::populateFoldTensorSubsetOpPatterns ( RewritePatternSet & patterns )

Appends patterns for folding tensor subset ops into consumer load/store ops into patterns.

(This includes patterns for folding tensor subset ops into vector transfer ops.)

Definition at line 247 of file FoldTensorSubsetOps.cpp.

References mlir::patterns, and populateFoldTensorSubsetIntoVectorTransferPatterns().

◆ populateMergeConsecutiveInsertExtractSlicePatterns()

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

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

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

Definition at line 230 of file MergeConsecutiveInsertExtractSlicePatterns.cpp.

References mlir::patterns.

◆ 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 720 of file ReshapePatterns.cpp.

References mlir::patterns.

◆ populateRewriteAsConstantPatterns()

void mlir::tensor::populateRewriteAsConstantPatterns	(	RewritePatternSet &	patterns,
		const ControlFoldFn &	controlFn )

Populates patterns with patterns that replace tensor ops (such as tensor.generate) with constants when possible.

Definition at line 211 of file RewriteAsConstant.cpp.

References mlir::patterns.

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

References target.

Referenced by canFoldIntoConsumerOp(), and canFoldIntoProducerOp().

◆ registerAllExtensions()

void mlir::tensor::registerAllExtensions ( DialectRegistry & registry )

Register all extensions of the Tensor dialect.

This should generally only be used by tools, or other use cases that really do want all extensions of the dialect. All other cases should prefer to instead register the specific extensions they intend to take advantage of.

Definition at line 14 of file AllExtensions.cpp.

References registerShardingInterfaceExternalModels().

Referenced by mlir::registerAllExtensions().

◆ registerBufferizableOpInterfaceExternalModels()

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

Definition at line 1217 of file BufferizableOpInterfaceImpl.cpp.

References mlir::DialectRegistry::addExtension(), mlir::MLIRContext::loadDialect(), and registerSubsetOpInterfaceExternalModels().

Referenced by mlir::registerAllDialects().

◆ registerFindPayloadReplacementOpInterfaceExternalModels()

void mlir::tensor::registerFindPayloadReplacementOpInterfaceExternalModels ( DialectRegistry & registry )

Definition at line 71 of file TensorTransformOps.cpp.

References mlir::DialectRegistry::addExtension().

Referenced by mlir::registerAllDialects().

◆ registerFoldTensorSubsetOpsPass()

void mlir::tensor::registerFoldTensorSubsetOpsPass ( )

inline

Definition at line 96 of file Passes.h.

◆ registerFoldTensorSubsetOpsPassPass()

void mlir::tensor::registerFoldTensorSubsetOpsPassPass ( )

inline

Definition at line 103 of file Passes.h.

◆ 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 131 of file TensorInferTypeOpInterfaceImpl.cpp.

References mlir::DialectRegistry::addExtension().

Referenced by mlir::registerAllDialects().

◆ registerRuntimeVerifiableOpInterfaceExternalModels()

void mlir::tensor::registerRuntimeVerifiableOpInterfaceExternalModels ( DialectRegistry & registry )

Definition at line 244 of file RuntimeOpVerification.cpp.

References mlir::DialectRegistry::addExtension().

Referenced by mlir::registerAllDialects().

◆ registerShardingInterfaceExternalModels()

void mlir::tensor::registerShardingInterfaceExternalModels ( DialectRegistry & registry )

Definition at line 104 of file ShardingExtensions.cpp.

References mlir::DialectRegistry::addExtension().

Referenced by registerAllExtensions().

◆ registerSubsetOpInterfaceExternalModels()

void mlir::tensor::registerSubsetOpInterfaceExternalModels ( DialectRegistry & registry )

Definition at line 87 of file SubsetInsertionOpInterfaceImpl.cpp.

References mlir::DialectRegistry::addExtension().

Referenced by mlir::registerAllDialects(), and registerBufferizableOpInterfaceExternalModels().

◆ registerTensorPasses()

void mlir::tensor::registerTensorPasses ( )

inline

Definition at line 117 of file Passes.h.

Referenced by mlir::registerAllPasses().

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

◆ registerTransformDialectExtension()

void mlir::tensor::registerTransformDialectExtension ( DialectRegistry & registry )

Definition at line 258 of file TensorTransformOps.cpp.

References mlir::DialectRegistry::addExtensions().

Referenced by mlir::registerAllExtensions().

◆ registerValueBoundsOpInterfaceExternalModels()

void mlir::tensor::registerValueBoundsOpInterfaceExternalModels ( DialectRegistry & registry )

Definition at line 116 of file ValueBoundsOpInterfaceImpl.cpp.

References mlir::DialectRegistry::addExtension().

Referenced by mlir::registerAllDialects().

◆ replaceExtractSliceWithTiledProducer()

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

Method 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 25 of file SwapExtractSliceWithProducerPatterns.cpp.

References mlir::Builder::getIndexAttr(), mlir::OpResult::getOwner(), mlir::OpResult::getResultNumber(), mlir::Builder::getType(), and mlir::isOneInteger().

◆ replaceInsertSlicesWithTiledConsumer()

FailureOr< TilingResult > mlir::tensor::replaceInsertSlicesWithTiledConsumer	(	OpBuilder &	builder,
		ArrayRef< tensor::InsertSliceOp >	sliceOps,
		ArrayRef< OpOperand * >	consumerOperands )

Method to swap tensor.insert_slices with their consumers when the consumer implements the TilingInterface.

The size of sliceOps and consumerOperands is expected to be the same. Every entry in consumerOperands represents a use of the the corresponding entry in sliceOps in the consumer. All entries of consumerOperands is expected to be uses in the same consumer.

Definition at line 61 of file SwapExtractSliceWithProducerPatterns.cpp.

References mlir::OpOperand::getOperandNumber(), and mlir::isOneInteger().

Referenced by tileAndFuseConsumerOfSlicesImpl().

◆ 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 the new op is returned.

Example:

%result = tensor.collapse_shape %0 [[0, 1], [2, 3]]

: tensor<?x1x30x10xf32> to tensor<?x300xf32>

result

Definition LinalgTransformOps.cpp:2097

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 158 of file ExtractSliceFromReshapeUtils.cpp.

References mlir::Builder::getIndexAttr(), getMixedSizes(), mlir::RewriterBase::replaceOp(), and mlir::RewriterBase::replaceOpWithNewOp().

Namespaces

Classes

Typedefs

Functions

Typedef Documentation

◆ ControlConstantExtractSliceFusionFn

◆ ControlFoldFn

Function Documentation

◆ bubbleUpPadSlice()

◆ buildIndependentOp() [1/2]

◆ buildIndependentOp() [2/2]

◆ canFoldIntoConsumerOp()

◆ canFoldIntoProducerOp()

◆ computeTransposedType()

◆ createCanonicalRankReducingExtractSliceOp()

◆ createCanonicalRankReducingInsertSliceOp()

◆ createDynamicDimValues()

◆ createFoldTensorSubsetOpsPass()

◆ createPadHighOp()

◆ dropGivenUnitDims()

◆ foldTensorCast()

◆ getCollapsedExtractSliceInfo()

◆ getExpandedExtractSliceInfo()

◆ getMixedSize()

◆ getMixedSizes()

◆ getOrCreateDestination()

◆ getOrCreateDestinations()

◆ getUpdatedOperandsAfterCastOpFolding()

◆ hasFoldableTensorCastOperand()

◆ isCastLikeExtractSliceOp()

◆ isCastLikeInsertSliceOp()

◆ isSameTypeWithoutEncoding()

◆ populateBubbleUpExpandShapePatterns()

◆ populateBubbleUpExtractSliceOpPatterns()

◆ populateDecomposeTensorConcatPatterns()

◆ populateDropRedundantInsertSliceRankExpansionPatterns()

◆ populateFoldCollapseExtractPatterns()

◆ populateFoldConstantExtractSlicePatterns()

◆ populateFoldTensorEmptyPatterns()

◆ populateFoldTensorSubsetIntoVectorTransferPatterns()

◆ populateFoldTensorSubsetOpPatterns()

◆ populateMergeConsecutiveInsertExtractSlicePatterns()

◆ populateReassociativeReshapeFoldingPatterns()

◆ populateRewriteAsConstantPatterns()

◆ preservesStaticInformation()

◆ registerAllExtensions()

◆ registerBufferizableOpInterfaceExternalModels()

◆ registerFindPayloadReplacementOpInterfaceExternalModels()

◆ registerFoldTensorSubsetOpsPass()

◆ registerFoldTensorSubsetOpsPassPass()

◆ registerInferTypeOpInterfaceExternalModels()

◆ registerRuntimeVerifiableOpInterfaceExternalModels()

◆ registerShardingInterfaceExternalModels()

◆ registerSubsetOpInterfaceExternalModels()

◆ registerTensorPasses()

◆ registerTilingInterfaceExternalModels()

◆ registerTilingInterfaceExternalModelsForPackUnPackOps()

◆ registerTransformDialectExtension()

◆ registerValueBoundsOpInterfaceExternalModels()

◆ replaceExtractSliceWithTiledProducer()

◆ replaceInsertSlicesWithTiledConsumer()

◆ simplifyCollapseShapeWithRankReducingExtractSlice()

Example:

Example:

Unsupported cases: