#include "mlir/Dialect/SCF/Transforms/TileUsingInterface.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/SCF/Utils/Utils.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Interfaces/DestinationStyleOpInterface.h"
#include "mlir/Interfaces/TilingInterface.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Debug.h"
#include <optional>

Macros
#define	DEBUG_TYPE "tile-using-interface"

Typedefs
using	YieldTiledValuesFn = std::function< LogicalResult(RewriterBase &rewriter, Location loc, ValueRange ivs, ValueRange newBbArgs, SmallVector< Value > &tiledValues, SmallVector< SmallVector< OpFoldResult > > &resultOffsets, SmallVector< SmallVector< OpFoldResult > > &resultSizes)>
	A function that allows returning additional yielded values during `yieldTiledValuesAndReplace`. More...

Functions
static SmallVector< int64_t >	fillInterchangeVector (ArrayRef< int64_t > interchangeVector, size_t iterationDomainSize)
	Helper method to adjust the interchange vector to match the iteration domain. More...

static bool	tileDividesIterationDomain (Range loopRange)

static OpFoldResult	getBoundedTileSize (OpBuilder &b, Location loc, Range loopRange, Value iv, OpFoldResult tileSize)
	Returns the bounded tile size given the current `iv`, `loopRange` and `tileSize`, i.e., `min(tileSize, range.end() - iv)`. More...

static Operation *	cloneOpAndUpdateDestinationArgs (RewriterBase &rewriter, Operation *op, ValueRange newDestArgs)
	Clones the operation and updates the destination if the operation implements the `DestinationStyleOpInterface`. More...

static LogicalResult	generateLoopNestUsingForOp (RewriterBase &rewriter, Location loc, ArrayRef< Range > loopRanges, ArrayRef< OpFoldResult > tileSizes, ValueRange destinationTensors, YieldTiledValuesFn yieldTiledValuesFn, SmallVector< LoopLikeOpInterface > &loops)
	Generate the tile-loop nest using `scf.for` operation. More...

static LogicalResult	generateLoopNestUsingForallOp (RewriterBase &rewriter, Location loc, ArrayRef< Range > loopRanges, ArrayRef< OpFoldResult > tileSizes, ArrayRef< Attribute > mappingVector, ValueRange destinationTensors, YieldTiledValuesFn tiledBodyFn, SmallVector< LoopLikeOpInterface > &loops)
	Generate the tile-loop nest using `scf.forall` operation. More...

static LogicalResult	generateLoopNest (RewriterBase &rewriter, Location loc, const scf::SCFTilingOptions &options, ArrayRef< Range > loopRanges, ArrayRef< OpFoldResult > tileSizes, ValueRange destinationTensors, YieldTiledValuesFn tiledBodyFn, SmallVector< LoopLikeOpInterface > &loops)
	Generate the tile-loop nest using the loop construct specifed in `options`. More...

template<typename LoopType >
FailureOr< LoopLikeOpInterface >	yieldTiledValuesAndReplaceLoop (LoopType loopOp, RewriterBase &rewriter, ValueRange newInitOperands, YieldTiledValuesFn yieldTiledValuesFn)
	Append the specified additional `newInitOperands` operands to the loops existing `init` operands (or similar), and replace `loopOp` with the new loop that has the additional init operands. More...

template<>
FailureOr< LoopLikeOpInterface >	yieldTiledValuesAndReplaceLoop< scf::ForOp > (scf::ForOp loopOp, RewriterBase &rewriter, ValueRange newInitOperands, YieldTiledValuesFn yieldTiledValuesFn)
	Implementation of `yieldTiledValuesAndReplaceLoop` for `scf.for`. More...

template<>
FailureOr< LoopLikeOpInterface >	yieldTiledValuesAndReplaceLoop< scf::ForallOp > (scf::ForallOp loopOp, RewriterBase &rewriter, ValueRange newInitOperands, YieldTiledValuesFn yieldTiledValuesFn)
	Implementation of `yieldTiledValuesAndReplaceLoop` for `scf.forall` More...

FailureOr< LoopLikeOpInterface >	yieldTiledValuesAndReplaceLoop (LoopLikeOpInterface loopLikeOp, RewriterBase &rewriter, ValueRange newInitOperands, YieldTiledValuesFn yieldTiledValuesFn)
	Implementation of `yieldTiledValuesAndReplaceLoop` for `LoopLikeOpInterface`, that just dispatches to the implementation for each supported loop type. More...

static LogicalResult	addInitOperandsToLoopNest (RewriterBase &rewriter, MutableArrayRef< LoopLikeOpInterface > loops, ValueRange newInitValues, YieldTiledValuesFn getNewTiledYieldsFn)
	Method to add new init values to a loop nest. More...

static std::tuple< OpResult, std::optional< OpOperand * > >	getUntiledProducerFromSliceSource (OpOperand *source, ArrayRef< LoopLikeOpInterface > loops)
	Return the untiled producer whose slice is used in a tiled consumer. More...

Macro Definition Documentation

◆ DEBUG_TYPE

#define DEBUG_TYPE "tile-using-interface"

Definition at line 30 of file TileUsingInterface.cpp.

Typedef Documentation

◆ YieldTiledValuesFn

using YieldTiledValuesFn = std::function<LogicalResult( RewriterBase &rewriter, Location loc, ValueRange ivs, ValueRange newBbArgs, SmallVector<Value> &tiledValues, SmallVector<SmallVector<OpFoldResult> > &resultOffsets, SmallVector<SmallVector<OpFoldResult> > &resultSizes)>

A function that allows returning additional yielded values during yieldTiledValuesAndReplace.

ivs induction variable for the loop.
newBbArgs basic block arguments corresponding to newly added iter_args.
tiledValues the tiled values to return. Must be of same size as newbbArgs, each element of this array is inserted into the corresponding element in newbbArgs.
resultOffsets is of the same size as tiledValues and represents the offsets to use when inserting corresponding element from tiledValues into the element from newBbArgs.
resultSizes is of the same size as tiledValues and represents the size of the corresponding element from tiledValues inserted into the element from newBbArgs. In case the method needs to return failure() the method is expected to clean up any inserted operations.

Definition at line 117 of file TileUsingInterface.cpp.

Function Documentation

◆ addInitOperandsToLoopNest()

static LogicalResult addInitOperandsToLoopNest	(	RewriterBase &	rewriter,
		MutableArrayRef< LoopLikeOpInterface >	loops,
		ValueRange	newInitValues,
		YieldTiledValuesFn	getNewTiledYieldsFn
	)

static

Method to add new init values to a loop nest.

Updates loops in-place with new loops that use the newInitValues. The outer-loops are updated to yield the new result values of the inner loop. For the innermost loop, the call back getNewYields is invoked to get the additional values to yield form the innermost loop.

Definition at line 457 of file TileUsingInterface.cpp.

References mlir::OpBuilder::create(), mlir::failed(), mlir::RewriterBase::mergeBlocks(), mlir::RewriterBase::replaceOp(), mlir::RewriterBase::replaceOpWithNewOp(), mlir::OpBuilder::setInsertionPoint(), mlir::success(), and yieldTiledValuesAndReplaceLoop().

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

◆ cloneOpAndUpdateDestinationArgs()

static Operation* cloneOpAndUpdateDestinationArgs	(	RewriterBase &	rewriter,
		Operation *	op,
		ValueRange	newDestArgs
	)

static

Clones the operation and updates the destination if the operation implements the DestinationStyleOpInterface.

Definition at line 125 of file TileUsingInterface.cpp.

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

◆ fillInterchangeVector()

static SmallVector<int64_t> fillInterchangeVector	(	ArrayRef< int64_t >	interchangeVector,
		size_t	iterationDomainSize
	)

static

Helper method to adjust the interchange vector to match the iteration domain.

Definition at line 47 of file TileUsingInterface.cpp.

◆ generateLoopNest()

static LogicalResult generateLoopNest	(	RewriterBase &	rewriter,
		Location	loc,
		const scf::SCFTilingOptions &	options,
		ArrayRef< Range >	loopRanges,
		ArrayRef< OpFoldResult >	tileSizes,
		ValueRange	destinationTensors,
		YieldTiledValuesFn	tiledBodyFn,
		SmallVector< LoopLikeOpInterface > &	loops
	)

static

Generate the tile-loop nest using the loop construct specifed in options.

options: Tiling options specified.
loopRanges specifies the lb, ub and step of the untiled iteration space.
tileSizes is the tile sizes to use. Zero represent untiled loops.
destinationTensors are the init values to use for the outer most loop.
yieldTiledValuesFn is called to generated the loop body of the inner most loop.
loops is an in-out parameter into which the generated loops are populated.

Definition at line 285 of file TileUsingInterface.cpp.

References mlir::scf::SCFTilingOptions::ForallOp, mlir::scf::SCFTilingOptions::ForOp, generateLoopNestUsingForallOp(), generateLoopNestUsingForOp(), mlir::isZeroIndex(), mlir::RewriterBase::notifyMatchFailure(), and options.

◆ generateLoopNestUsingForallOp()

static LogicalResult generateLoopNestUsingForallOp	(	RewriterBase &	rewriter,
		Location	loc,
		ArrayRef< Range >	loopRanges,
		ArrayRef< OpFoldResult >	tileSizes,
		ArrayRef< Attribute >	mappingVector,
		ValueRange	destinationTensors,
		YieldTiledValuesFn	tiledBodyFn,
		SmallVector< LoopLikeOpInterface > &	loops
	)

static

Generate the tile-loop nest using scf.forall operation.

loopRanges specifies the lb, ub and step of the untiled iteration space.
tileSizes is the tile sizes to use. Zero represent untiled loops.
destinationTensors are the init values to use for the outer most loop.
mappingVector is the mapping attributes to use for loop construction. Can be empty.
yieldTiledValuesFn is called to generated the loop body of the inner most loop.
loops is an in-out parameter into which the generated loops are populated.

Definition at line 221 of file TileUsingInterface.cpp.

References mlir::OpBuilder::create(), mlir::failed(), mlir::Builder::getArrayAttr(), mlir::Builder::getIndexAttr(), mlir::isConstantIntValue(), mlir::RewriterBase::notifyMatchFailure(), mlir::OpBuilder::setInsertionPoint(), mlir::OpBuilder::setInsertionPointToEnd(), and mlir::success().

Referenced by generateLoopNest().

◆ generateLoopNestUsingForOp()

static LogicalResult generateLoopNestUsingForOp	(	RewriterBase &	rewriter,
		Location	loc,
		ArrayRef< Range >	loopRanges,
		ArrayRef< OpFoldResult >	tileSizes,
		ValueRange	destinationTensors,
		YieldTiledValuesFn	yieldTiledValuesFn,
		SmallVector< LoopLikeOpInterface > &	loops
	)

static

Generate the tile-loop nest using scf.for operation.

loopRanges specifies the lb, ub and step of the untiled iteration space.
tileSizes is the tile sizes to use. Zero represent untiled loops.
destinationTensors are the init values to use for the outer most loop.
yieldTiledValuesFn is called to generated the loop body of the inner most loop.
loops is an in-out parameter into which the generated loops are populated.

Definition at line 145 of file TileUsingInterface.cpp.

References mlir::OpBuilder::create(), mlir::failed(), mlir::Builder::getIndexAttr(), mlir::getValueOrCreateConstantIndexOp(), mlir::isConstantIntValue(), mlir::RewriterBase::notifyMatchFailure(), mlir::OpBuilder::setInsertionPointToEnd(), and mlir::success().

Referenced by generateLoopNest().

◆ getBoundedTileSize()

static OpFoldResult getBoundedTileSize	(	OpBuilder &	b,
		Location	loc,
		Range	loopRange,
		Value	iv,
		OpFoldResult	tileSize
	)

static

Returns the bounded tile size given the current iv, loopRange and tileSize, i.e., min(tileSize, range.end() - iv).

Definition at line 79 of file TileUsingInterface.cpp.

References mlir::bindDims(), mlir::bindSymbols(), mlir::AffineMap::get(), mlir::getConstantIntValue(), mlir::Builder::getContext(), mlir::getValueOrCreateConstantIndexOp(), mlir::affine::makeComposedFoldedAffineMin(), mlir::Range::offset, mlir::Range::size, and tileDividesIterationDomain().

◆ getUntiledProducerFromSliceSource()

static std::tuple<OpResult, std::optional<OpOperand *> > getUntiledProducerFromSliceSource	(	OpOperand *	source,
		ArrayRef< LoopLikeOpInterface >	loops
	)

static

Return the untiled producer whose slice is used in a tiled consumer.

The method traverses the tile loop nest (loops) if needed, and returns the iter_args of the outer most that is encountered. Traversing the iter_args indicates that this is a destination operand of the consumer. If there was no loop traversal needed, the second value of the returned tuple is empty.

Definition at line 807 of file TileUsingInterface.cpp.

References mlir::IROperand< DerivedT, IRValueT >::get().

Referenced by mlir::scf::tileAndFuseProducerOfSlice(), and mlir::scf::tileConsumerAndFuseProducersUsingSCF().

◆ tileDividesIterationDomain()

static bool tileDividesIterationDomain ( Range loopRange )

static

Definition at line 64 of file TileUsingInterface.cpp.

References mlir::getConstantIntValue(), mlir::Range::offset, mlir::Range::size, and mlir::Range::stride.

Referenced by getBoundedTileSize().

◆ yieldTiledValuesAndReplaceLoop() [1/2]

FailureOr<LoopLikeOpInterface> yieldTiledValuesAndReplaceLoop	(	LoopLikeOpInterface	loopLikeOp,
		RewriterBase &	rewriter,
		ValueRange	newInitOperands,
		YieldTiledValuesFn	yieldTiledValuesFn
	)

Implementation of yieldTiledValuesAndReplaceLoop for LoopLikeOpInterface, that just dispatches to the implementation for each supported loop type.

Definition at line 437 of file TileUsingInterface.cpp.

References mlir::RewriterBase::notifyMatchFailure(), and yieldTiledValuesAndReplaceLoop().

◆ yieldTiledValuesAndReplaceLoop() [2/2]

template<typename LoopType >

FailureOr<LoopLikeOpInterface> yieldTiledValuesAndReplaceLoop	(	LoopType	loopOp,
		RewriterBase &	rewriter,
		ValueRange	newInitOperands,
		YieldTiledValuesFn	yieldTiledValuesFn
	)

Append the specified additional newInitOperands operands to the loops existing init operands (or similar), and replace loopOp with the new loop that has the additional init operands.

The loop body of this loop is moved over to the new loop. yieldTiledValuesFn is called to get the new tiled values returned, and the offset and sizes at which the tiled value is inserted into the new region iter_args that correspond to the newly added init operands.

Definition at line 321 of file TileUsingInterface.cpp.

References mlir::RewriterBase::notifyMatchFailure().

Referenced by addInitOperandsToLoopNest(), and yieldTiledValuesAndReplaceLoop().

◆ yieldTiledValuesAndReplaceLoop< scf::ForallOp >()

template<>

FailureOr<LoopLikeOpInterface> yieldTiledValuesAndReplaceLoop< scf::ForallOp >	(	scf::ForallOp	loopOp,
		RewriterBase &	rewriter,
		ValueRange	newInitOperands,
		YieldTiledValuesFn	yieldTiledValuesFn
	)

Implementation of yieldTiledValuesAndReplaceLoop for scf.forall

Definition at line 321 of file TileUsingInterface.cpp.

◆ yieldTiledValuesAndReplaceLoop< scf::ForOp >()

template<>

FailureOr<LoopLikeOpInterface> yieldTiledValuesAndReplaceLoop< scf::ForOp >	(	scf::ForOp	loopOp,
		RewriterBase &	rewriter,
		ValueRange	newInitOperands,
		YieldTiledValuesFn	yieldTiledValuesFn
	)

Implementation of yieldTiledValuesAndReplaceLoop for scf.for.

Definition at line 321 of file TileUsingInterface.cpp.

Macros

Typedefs

Functions

Macro Definition Documentation

◆ DEBUG_TYPE

Typedef Documentation

◆ YieldTiledValuesFn

Function Documentation

◆ addInitOperandsToLoopNest()

◆ cloneOpAndUpdateDestinationArgs()

◆ fillInterchangeVector()

◆ generateLoopNest()

◆ generateLoopNestUsingForallOp()

◆ generateLoopNestUsingForOp()

◆ getBoundedTileSize()

◆ getUntiledProducerFromSliceSource()

◆ tileDividesIterationDomain()

◆ yieldTiledValuesAndReplaceLoop() [1/2]

◆ yieldTiledValuesAndReplaceLoop() [2/2]

◆ yieldTiledValuesAndReplaceLoop< scf::ForallOp >()

◆ yieldTiledValuesAndReplaceLoop< scf::ForOp >()