HoistPadding.cpp File Reference

#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/Dialect/Utils/IndexingUtils.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"

Go to the source code of this file.

Macros
#define	DEBUG_TYPE   "hoist-padding"
#define	DBGS()   (dbgs() << '[' << DEBUG_TYPE << "] ")

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. More...
static void	getEnclosingLoopsUntil (tensor::PadOp padOp, scf::ForOp untilLoop, SmallVector< scf::ForOp > &reverseEnclosingLoops)
	Return at most nLevels of immediately enclosing scf::ForOp loops. More...
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). More...
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. More...
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. More...
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. More...
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. More...

Macro Definition Documentation

DBGS

#define DBGS

(

)

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

Definition at line 36 of file HoistPadding.cpp.

DEBUG_TYPE

#define DEBUG_TYPE   "hoist-padding"

Definition at line 34 of file HoistPadding.cpp.

Function Documentation

buildLoopIterationCount()

static 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 514 of file HoistPadding.cpp.

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

Referenced by replaceByPackingResult().

buildPackingLoopNestImpl() [1/2]

static 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 677 of file HoistPadding.cpp.

References buildPackingLoopNestImpl(), mlir::tensor::computeTransposedType(), mlir::OpBuilder::create(), DBGS, mlir::failed(), mlir::failure(), mlir::get(), and mlir::OpBuilder::setInsertionPoint().

buildPackingLoopNestImpl() [2/2]

static 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 541 of file HoistPadding.cpp.

References mlir::IROperand< DerivedT, IRValueT >::get(), mlir::BlockArgument::getOwner(), mlir::Block::getParentOp(), and mlir::IRMapping::map().

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

computeBackwardSlice()

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

static

Definition at line 111 of file HoistPadding.cpp.

References mlir::SliceOptions::filter.

debugPrintBackwardSlice()

static void debugPrintBackwardSlice ( SetVector< Operation * > &  backwardSlice )

static

Definition at line 55 of file HoistPadding.cpp.

References DBGS.

debugPrintLoopInShortForm()

static bool debugPrintLoopInShortForm ( Operation *  op )

static

Definition at line 43 of file HoistPadding.cpp.

Referenced by getAtMostNEnclosingLoops(), and getEnclosingLoopsUntil().

getAtMostNEnclosingLoops()

static 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 74 of file HoistPadding.cpp.

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

getEnclosingLoopsUntil()

static 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 93 of file HoistPadding.cpp.

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

isDefinedOutsideOrConstant()

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

static

Definition at line 502 of file HoistPadding.cpp.

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

Referenced by buildLoopIterationCount().

padThroughLoopIterArg()

static 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 797 of file HoistPadding.cpp.

References mlir::OpBuilder::create(), 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()

static 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 875 of file HoistPadding.cpp.

References buildLoopIterationCount(), mlir::linalg::detail::PackingResult::clonedLoopIvs, copy(), mlir::OpBuilder::create(), DBGS, mlir::affine::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()

static 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 747 of file HoistPadding.cpp.

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

Referenced by padThroughLoopIterArg().