#include "mlir/Analysis/Presburger/IntegerRelation.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/Transforms/Transforms.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/Utils/Utils.h"
#include "mlir/IR/AsmState.h"
#include "mlir/IR/Dominance.h"
#include "mlir/IR/Matchers.h"
#include "mlir/Interfaces/DestinationStyleOpInterface.h"
#include "mlir/Transforms/LoopInvariantCodeMotionUtils.h"
#include "mlir/Transforms/RegionUtils.h"
#include "llvm/Support/Debug.h"

Macros
#define	DEBUG_TYPE "hoist-padding"
#define	DBGS()

Functions
static bool	debugPrintLoopInShortForm (Operation *op)
static void	debugPrintBackwardSlice (SetVector< Operation * > &backwardSlice)
static void	getAtMostNEnclosingLoops (tensor::PadOp padOp, int nLevels, SmallVector< scf::ForOp > &reverseEnclosingLoops)
	Return at most nLevels of immediately enclosing scf::ForOp loops.
static void	getEnclosingLoopsUntil (tensor::PadOp padOp, scf::ForOp untilLoop, SmallVector< scf::ForOp > &reverseEnclosingLoops)
	Return at most nLevels of immediately enclosing scf::ForOp loops.
static void	computeBackwardSlice (tensor::PadOp padOp, scf::ForOp outermostEnclosingForOp, SetVector< Operation * > &backwardSlice)
static bool	isDefinedOutsideOrConstant (scf::ForOp outer, Value v)
static Value	buildLoopIterationCount (RewriterBase &rewriter, scf::ForOp outer, scf::ForOp forOp)
	Return the current iteration number in the loop (iv - lb).ceilDiv(step).
static FailureOr< PackingResult >	buildPackingLoopNestImpl (RewriterBase &rewriter, IRMapping &bvm, tensor::PadOp opToHoist, ArrayRef< int64_t > transposeVector, RankedTensorType transposedTensorType, tensor::EmptyOp emptyOp, const HoistPaddingAnalysis &analysis)
static FailureOr< PackingResult >	buildPackingLoopNestImpl (RewriterBase &rewriter, IRMapping &bvm, tensor::PadOp opToHoist, ArrayRef< int64_t > transposeVector, const HoistPaddingAnalysis &analysis)
	Build the packing loop nest required to hoist opToHoist above outermostEnclosingForOp.
static bool	tracesBackToExpectedValue (tensor::ExtractSliceOp extractSliceOp, Value expectedSource)
	Return true if we can walk back the use-def chain from extractSliceOp to expectedSource going through DestinationStyleOpInterface inits only.
static tensor::ExtractSliceOp	padThroughLoopIterArg (RewriterBase &rewriter, Value paddedValueBeforeHoisting, Value hoistedPackedTensor, tensor::ExtractSliceOp outerSliceOp, scf::ForOp forOp)
	If the original consumer of outerSliceOp was a forOp (i.e.
static Value	replaceByPackingResult (RewriterBase &rewriter, const IRMapping &bvm, tensor::PadOp opToHoist, RankedTensorType transposedTensorType, const HoistPaddingAnalysis &analysis, const PackingResult &packingResult)
	Produce a tensor extracted from the packingResult.

Macro Definition Documentation

◆ DBGS

#define DBGS ( )

Value:

(dbgs() << '[' << DEBUG_TYPE << "] ")

DEBUG_TYPE

#define DEBUG_TYPE

Definition LocalAliasAnalysis.cpp:32

Definition at line 35 of file HoistPadding.cpp.

Referenced by debugPrintBackwardSlice(), getAtMostNEnclosingLoops(), and getEnclosingLoopsUntil().

◆ DEBUG_TYPE

#define DEBUG_TYPE "hoist-padding"

Definition at line 33 of file HoistPadding.cpp.

Function Documentation

◆ buildLoopIterationCount()

Value buildLoopIterationCount	(	RewriterBase &	rewriter,
		scf::ForOp	outer,
		scf::ForOp	forOp )

static

Return the current iteration number in the loop (iv - lb).ceilDiv(step).

The returned Value is guaranteed not to depend on any loop comprised in [outer, forOp]. Return null if such a loop-independent quantity cannot be computed.

Definition at line 518 of file HoistPadding.cpp.

References mlir::bindDims(), mlir::bindSymbols(), mlir::OpBuilder::createOrFold(), mlir::Value::getLoc(), and isDefinedOutsideOrConstant().

Referenced by buildPackingLoopNestImpl(), and replaceByPackingResult().

◆ buildPackingLoopNestImpl() [1/2]

FailureOr< PackingResult > buildPackingLoopNestImpl	(	RewriterBase &	rewriter,
		IRMapping &	bvm,
		tensor::PadOp	opToHoist,
		ArrayRef< int64_t >	transposeVector,
		const HoistPaddingAnalysis &	analysis )

static

Build the packing loop nest required to hoist opToHoist above outermostEnclosingForOp.

The loop nest is built just before outermostEnclosingForOp.

Definition at line 683 of file HoistPadding.cpp.

References buildPackingLoopNestImpl(), mlir::tensor::computeTransposedType(), DBGS, and mlir::OpBuilder::setInsertionPoint().

◆ buildPackingLoopNestImpl() [2/2]

FailureOr< PackingResult > buildPackingLoopNestImpl	(	RewriterBase &	rewriter,
		IRMapping &	bvm,
		tensor::PadOp	opToHoist,
		ArrayRef< int64_t >	transposeVector,
		RankedTensorType	transposedTensorType,
		tensor::EmptyOp	emptyOp,
		const HoistPaddingAnalysis &	analysis )

static

Definition at line 545 of file HoistPadding.cpp.

References buildLoopIterationCount(), mlir::OpBuilder::clone(), DBGS, mlir::IROperand< DerivedT, IRValueT >::get(), mlir::scf::getForInductionVarOwner(), mlir::Builder::getIndexAttr(), mlir::BlockArgument::getOwner(), mlir::Block::getParentOp(), inserted(), mlir::IRMapping::lookup(), mlir::IRMapping::lookupOrDefault(), mlir::IRMapping::map(), mlir::OpBuilder::setInsertionPointToEnd(), and mlir::OpBuilder::setInsertionPointToStart().

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

◆ computeBackwardSlice()

void computeBackwardSlice	(	tensor::PadOp	padOp,
		scf::ForOp	outermostEnclosingForOp,
		SetVector< Operation * > &	backwardSlice )

static

Definition at line 110 of file HoistPadding.cpp.

References mlir::DominanceInfo::dominates(), mlir::SliceOptions::filter, mlir::getBackwardSlice(), mlir::getUsedValuesDefinedAbove(), mlir::SliceOptions::inclusive, and result.

◆ debugPrintBackwardSlice()

void debugPrintBackwardSlice ( SetVector< Operation * > & backwardSlice )

static

Definition at line 54 of file HoistPadding.cpp.

References DBGS, and debugPrintLoopInShortForm().

◆ debugPrintLoopInShortForm()

bool debugPrintLoopInShortForm ( Operation * op )

static

Definition at line 42 of file HoistPadding.cpp.

References mlir::Operation::getParentOfType().

Referenced by debugPrintBackwardSlice(), getAtMostNEnclosingLoops(), and getEnclosingLoopsUntil().

◆ getAtMostNEnclosingLoops()

void getAtMostNEnclosingLoops	(	tensor::PadOp	padOp,
		int	nLevels,
		SmallVector< scf::ForOp > &	reverseEnclosingLoops )

static

Return at most nLevels of immediately enclosing scf::ForOp loops.

Stops at the first parent that is not an scf::ForOp. Multi-loops such as scf.parallel or linalg.tiled_loop are not modeled atm. Control-flow and other containing ops with regions are not modeled atm.

Definition at line 73 of file HoistPadding.cpp.

References DBGS, debugPrintLoopInShortForm(), and mlir::Operation::getParentOp().

◆ getEnclosingLoopsUntil()

void getEnclosingLoopsUntil	(	tensor::PadOp	padOp,
		scf::ForOp	untilLoop,
		SmallVector< scf::ForOp > &	reverseEnclosingLoops )

static

Return at most nLevels of immediately enclosing scf::ForOp loops.

Stops at the first parent that is not an scf::ForOp. Multi-loops such as scf.parallel or linalg.tiled_loop are not modeled atm. Control-flow and other containing ops with regions are not modeled atm.

Definition at line 92 of file HoistPadding.cpp.

References DBGS, debugPrintLoopInShortForm(), and mlir::Operation::getParentOp().

◆ isDefinedOutsideOrConstant()

bool isDefinedOutsideOrConstant	(	scf::ForOp	outer,
		Value	v )

static

Definition at line 506 of file HoistPadding.cpp.

References mlir::m_Constant(), and mlir::matchPattern().

Referenced by buildLoopIterationCount().

◆ padThroughLoopIterArg()

tensor::ExtractSliceOp padThroughLoopIterArg	(	RewriterBase &	rewriter,
		Value	paddedValueBeforeHoisting,
		Value	hoistedPackedTensor,
		tensor::ExtractSliceOp	outerSliceOp,
		scf::ForOp	forOp )

static

If the original consumer of outerSliceOp was a forOp (i.e.

through an iter arg), propagate the hoistedPackedTensor value through the same iter arg. TODO: for multiple loops we need to track the use to the innermost loop.

Match:

%outerSliceOp = tensor.extract_slice ..
%f = scf.for ... iter_args(%arg0 = %outerSliceOp) {
  %hoistedPackedTensor = tensor.pad %arg0
  %1 = compute %hoistedPackedTensor
  %2 = tensor.extract_slice %1
  scf.yield %2
}

and rewrite as:

%outerSliceOp = tensor.extract_slice ..
%hoistedPackedTensor = tensor.pad %outerSliceOp
%f = scf.for ... iter_args(%arg0 = %hoistedPackedTensor) {
  %1 = compute %arg0
  scf.yield %1
}
%2 = tensor.extract_slice %forOp

Return null when no rewrite happened.

Definition at line 802 of file HoistPadding.cpp.

References b, DBGS, mlir::RewriterBase::finalizeOpModification(), mlir::Value::getDefiningOp(), mlir::Value::getLoc(), mlir::RewriterBase::replaceAllUsesWith(), mlir::OpBuilder::setInsertionPointAfter(), mlir::RewriterBase::startOpModification(), and tracesBackToExpectedValue().

Referenced by replaceByPackingResult().

◆ replaceByPackingResult()

Value replaceByPackingResult	(	RewriterBase &	rewriter,
		const IRMapping &	bvm,
		tensor::PadOp	opToHoist,
		RankedTensorType	transposedTensorType,
		const HoistPaddingAnalysis &	analysis,
		const PackingResult &	packingResult )

static

Produce a tensor extracted from the packingResult.

This can be used as a replacement for opToHoist in callers.

Definition at line 880 of file HoistPadding.cpp.

References buildLoopIterationCount(), mlir::linalg::detail::PackingResult::clonedLoopIvs, DBGS, mlir::scf::getForInductionVarOwner(), mlir::Builder::getIndexAttr(), mlir::IRMapping::lookup(), padThroughLoopIterArg(), mlir::OpBuilder::setInsertionPoint(), mlir::linalg::detail::PackingResult::sizes, and mlir::linalg::detail::PackingResult::strides.

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

◆ tracesBackToExpectedValue()

bool tracesBackToExpectedValue	(	tensor::ExtractSliceOp	extractSliceOp,
		Value	expectedSource )

static

Return true if we can walk back the use-def chain from extractSliceOp to expectedSource going through DestinationStyleOpInterface inits only.

This is a poor man's analysis that is sufficient to check the extractSliceOp the matches tensor.pad we want to hoist. In the future, it will be easier to ensure this with a matching symmetric tensor.unpad op.

Definition at line 753 of file HoistPadding.cpp.

References DBGS, and mlir::Value::getDefiningOp().

Referenced by padThroughLoopIterArg().

Macros

Functions

Macro Definition Documentation

◆ DBGS

◆ DEBUG_TYPE

Function Documentation

◆ buildLoopIterationCount()

◆ buildPackingLoopNestImpl() [1/2]

◆ buildPackingLoopNestImpl() [2/2]

◆ computeBackwardSlice()

◆ debugPrintBackwardSlice()

◆ debugPrintLoopInShortForm()

◆ getAtMostNEnclosingLoops()

◆ getEnclosingLoopsUntil()

◆ isDefinedOutsideOrConstant()

◆ padThroughLoopIterArg()

◆ replaceByPackingResult()

◆ tracesBackToExpectedValue()