MLIR  16.0.0git
Classes | Macros | Typedefs | Enumerations | Functions
Vectorization.cpp File Reference
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/RegionUtils.h"
#include "llvm/ADT/ScopeExit.h"
#include "llvm/ADT/Sequence.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
#include <type_traits>
+ Include dependency graph for Vectorization.cpp:

Go to the source code of this file.

Classes

struct  VectorizationResult
 
struct  GenericPadOpVectorizationPattern
 Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp and InsertSliceOp. More...
 
struct  VectorizePadOpUserPattern< OpTy >
 Base pattern for rewriting tensor::PadOps whose result is consumed by a given operation type OpTy. More...
 
struct  PadOpVectorizationWithTransferReadPattern
 Rewrite use of tensor::PadOp result in TransferReadOp. More...
 
struct  PadOpVectorizationWithTransferWritePattern
 Rewrite use of tensor::PadOp result in TransferWriteOp. More...
 
struct  PadOpVectorizationWithInsertSlicePattern
 Rewrite use of tensor::PadOp result in InsertSliceOp. More...
 
struct  VectorizeConvolution
 

Macros

#define DEBUG_TYPE   "linalg-vectorization"
 
#define DBGS()   (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
 
#define LDBG(X)   LLVM_DEBUG(DBGS() << X)
 

Typedefs

using CustomVectorizationPrecondition = std::function< LogicalResult(Operation *)>
 
using CustomVectorizationHook = std::function< VectorizationResult(Operation *, const BlockAndValueMapping &)>
 

Enumerations

enum class  Conv1DOpOrder { Ncw , Nwc }
 Helper enum to represent conv1d input traversal order. More...
 
enum  VectorizationStatus { Failure = 0 , NoReplace , NewOp }
 Helper data structure to represent the result of vectorization. More...
 

Functions

static FailureOr< Operation * > vectorizeConvolution (OpBuilder &b, LinalgOp convOp)
 Try to vectorize convOp as a convolution. More...
 
template<typename OpType >
static OpType getSingleOpOfType (Block &block)
 Return the unique instance of OpType in block if it is indeed unique. More...
 
static AffineMap reindexIndexingMap (AffineMap map)
 Given an indexing map coming from a LinalgOp indexing, restricted to a projectedPermutation, compress the unused dimensions to serve as a permutation_map for a vector transfer operation. More...
 
static OperationmatchLinalgReduction (OpOperand *outputOperand)
 Check whether outputOperand is a reduction with a single combiner operation. More...
 
static Value broadcastIfNeeded (OpBuilder &b, Value value, ArrayRef< int64_t > shape)
 Broadcast value to a vector of shape if possible. More...
 
static OperationbuildMultiDimReduce (OpBuilder &b, Operation *reduceOp, Value valueToReduce, Value acc, const SmallVector< bool > &reductionMask)
 Create MultiDimReductionOp to compute the reduction for reductionOp. More...
 
static SmallVector< bool > getReductionMask (LinalgOp linalgOp)
 
static Value buildVectorWrite (OpBuilder &b, Value value, OpOperand *outputOperand)
 Build a vector.transfer_write of value into outputOperand at indices set to all 0; where outputOperand is an output operand of the LinalgOp currently being vectorized. More...
 
static VectorizationResult vectorizeLinalgYield (OpBuilder &b, Operation *op, const BlockAndValueMapping &bvm, LinalgOp linalgOp, SmallVectorImpl< Value > &newResults)
 Helper function to vectorize the terminator of a linalgOp. More...
 
static VectorizationResult vectorizeLinalgIndex (OpBuilder &b, Operation *op, LinalgOp linalgOp)
 Helper function to vectorize the index operations of a linalgOp. More...
 
static LogicalResult tensorExtractVectorizationPrecondition (Operation *op)
 Helper function to check if the tensor.extract can be vectorized by the custom hook vectorizeTensorExtract. More...
 
static VectorizationResult vectorizeTensorExtract (OpBuilder &b, Operation *op, LinalgOp linalgOp, const BlockAndValueMapping &bvm)
 Helper function to vectorize the tensor.extract operations. More...
 
static OperationreduceIfNeeded (OpBuilder &b, LinalgOp linalgOp, Operation *op, Value reduceValue, Value initialValue, const BlockAndValueMapping &bvm)
 Emit reduction operations if the shapes of the value to reduce is different that the result shape. More...
 
static VectorizationResult vectorizeOneOp (OpBuilder &b, LinalgOp linalgOp, Operation *op, const BlockAndValueMapping &bvm, ArrayRef< CustomVectorizationHook > customVectorizationHooks)
 Generic vectorization for a single operation op, given already vectorized operands carried by bvm. More...
 
static LogicalResult vectorizeAsLinalgGeneric (OpBuilder &b, LinalgOp linalgOp, SmallVectorImpl< Value > &newResults)
 Generic vectorization function that rewrites the body of a linalgOp into vector form. More...
 
static LogicalResult reductionPreconditions (LinalgOp op)
 
static LogicalResult vectorizeStaticLinalgOpPrecondition (linalg::LinalgOp op, ArrayRef< CustomVectorizationPrecondition > customPreconditions)
 
static int64_t getIntFromAttr (Attribute attr)
 Helper function that retrieves the value of an IntegerAttr. More...
 
static SmallVector< ValueofrToIndexValues (OpBuilder &builder, Location loc, ArrayRef< OpFoldResult > ofrs)
 Given an ArrayRef of OpFoldResults, return a vector of Values. More...
 
static bool mayExistInterleavedUses (Operation *firstOp, Operation *secondOp, ValueRange values)
 Check whether there is any interleaved use of any values between firstOp and secondOp. More...
 
static memref::SubViewOp getSubViewUseIfUnique (Value v)
 Return the unique subview use of v if it is indeed unique, null otherwise. More...
 
template<int N>
static void bindShapeDims (ShapedType shapedType)
 
template<int N, typename IntTy , typename... IntTy2>
static void bindShapeDims (ShapedType shapedType, IntTy &val, IntTy2 &...vals)
 
template<typename... IntTy>
static void bindShapeDims (ShapedType shapedType, IntTy &...vals)
 Bind a pack of int& to the leading dimensions of shapedType.getShape(). More...
 

Macro Definition Documentation

◆ DBGS

#define DBGS ( )    (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")

Definition at line 45 of file Vectorization.cpp.

◆ DEBUG_TYPE

#define DEBUG_TYPE   "linalg-vectorization"

Definition at line 43 of file Vectorization.cpp.

◆ LDBG

#define LDBG (   X)    LLVM_DEBUG(DBGS() << X)

Definition at line 46 of file Vectorization.cpp.

Typedef Documentation

◆ CustomVectorizationHook

Definition at line 244 of file Vectorization.cpp.

◆ CustomVectorizationPrecondition

Definition at line 238 of file Vectorization.cpp.

Enumeration Type Documentation

◆ Conv1DOpOrder

enum Conv1DOpOrder
strong

Helper enum to represent conv1d input traversal order.

Enumerator
Ncw 
Nwc 

Definition at line 95 of file Vectorization.cpp.

◆ VectorizationStatus

Helper data structure to represent the result of vectorization.

In certain specific cases, like terminators, we do not want to propagate/

Enumerator
Failure 

Op failed to vectorize.

NoReplace 

Op vectorized and custom function took care of replacement logic.

NewOp 

Op vectorized into a new Op whose results will replace original Op's results.

Definition at line 102 of file Vectorization.cpp.

Function Documentation

◆ bindShapeDims() [1/3]

template<int N>
static void bindShapeDims ( ShapedType  shapedType)
static

Definition at line 1383 of file Vectorization.cpp.

Referenced by bindShapeDims().

◆ bindShapeDims() [2/3]

template<typename... IntTy>
static void bindShapeDims ( ShapedType  shapedType,
IntTy &...  vals 
)
static

Bind a pack of int& to the leading dimensions of shapedType.getShape().

Definition at line 1393 of file Vectorization.cpp.

◆ bindShapeDims() [3/3]

template<int N, typename IntTy , typename... IntTy2>
static void bindShapeDims ( ShapedType  shapedType,
IntTy &  val,
IntTy2 &...  vals 
)
static

Definition at line 1386 of file Vectorization.cpp.

References bindShapeDims().

◆ broadcastIfNeeded()

static Value broadcastIfNeeded ( OpBuilder b,
Value  value,
ArrayRef< int64_t >  shape 
)
static

Broadcast value to a vector of shape if possible.

Return value otherwise.

Definition at line 166 of file Vectorization.cpp.

References mlir::OpBuilder::createOrFold(), mlir::getElementTypeOrSelf(), mlir::OpBuilder::getInsertionPoint(), mlir::vector::isBroadcastableTo(), mlir::vector::Success, and value.

Referenced by buildVectorWrite().

◆ buildMultiDimReduce()

static Operation* buildMultiDimReduce ( OpBuilder b,
Operation reduceOp,
Value  valueToReduce,
Value  acc,
const SmallVector< bool > &  reductionMask 
)
static

Create MultiDimReductionOp to compute the reduction for reductionOp.

This assumes that reductionOp has two operands and one of them is the reduction initial value.

Definition at line 183 of file Vectorization.cpp.

References mlir::OpBuilder::create(), mlir::linalg::getCombinerOpKind(), and mlir::Operation::getLoc().

Referenced by reduceIfNeeded().

◆ buildVectorWrite()

static Value buildVectorWrite ( OpBuilder b,
Value  value,
OpOperand outputOperand 
)
static

Build a vector.transfer_write of value into outputOperand at indices set to all 0; where outputOperand is an output operand of the LinalgOp currently being vectorized.

If dest has null rank, build an memref.store. Return the produced value or null if no value is produced.

Definition at line 201 of file Vectorization.cpp.

References mlir::applyPermutationMap(), broadcastIfNeeded(), mlir::OpBuilder::create(), mlir::IROperand< DerivedT, IRValueT >::get(), mlir::getElementTypeOrSelf(), mlir::detail::IROperandBase::getOwner(), mlir::Operation::getResult(), mlir::Operation::getResults(), mlir::Value::getType(), mlir::inversePermutation(), LDBG, reindexIndexingMap(), and value.

Referenced by vectorizeLinalgYield().

◆ getIntFromAttr()

static int64_t getIntFromAttr ( Attribute  attr)
static

Helper function that retrieves the value of an IntegerAttr.

Definition at line 756 of file Vectorization.cpp.

References mlir::Attribute::cast().

Referenced by ofrToIndexValues().

◆ getReductionMask()

static SmallVector<bool> getReductionMask ( LinalgOp  linalgOp)
static

Definition at line 192 of file Vectorization.cpp.

References mlir::linalg::isReductionIterator().

Referenced by reduceIfNeeded().

◆ getSingleOpOfType()

template<typename OpType >
static OpType getSingleOpOfType ( Block block)
static

Return the unique instance of OpType in block if it is indeed unique.

Return null if none or more than 1 instances exist.

Definition at line 55 of file Vectorization.cpp.

◆ getSubViewUseIfUnique()

static memref::SubViewOp getSubViewUseIfUnique ( Value  v)
static

Return the unique subview use of v if it is indeed unique, null otherwise.

Definition at line 1226 of file Vectorization.cpp.

References mlir::Value::getUses().

Referenced by mlir::linalg::LinalgCopyVTRForwardingPattern::matchAndRewrite(), and mlir::linalg::LinalgCopyVTWForwardingPattern::matchAndRewrite().

◆ matchLinalgReduction()

static Operation* matchLinalgReduction ( OpOperand outputOperand)
static

Check whether outputOperand is a reduction with a single combiner operation.

Return the combiner operation of the reduction. Return nullptr otherwise. Multiple reduction operations would impose an ordering between reduction dimensions and is currently unsupported in Linalg. This limitation is motivated by the fact that e.g. min(max(X)) != max(min(X))

Definition at line 150 of file Vectorization.cpp.

References mlir::OpOperand::getOperandNumber(), mlir::detail::IROperandBase::getOwner(), and mlir::matchReduction().

Referenced by reductionPreconditions().

◆ mayExistInterleavedUses()

static bool mayExistInterleavedUses ( Operation firstOp,
Operation secondOp,
ValueRange  values 
)
static

Check whether there is any interleaved use of any values between firstOp and secondOp.

Conservatively return true if any op or value is in a different block.

Definition at line 1199 of file Vectorization.cpp.

References mlir::Operation::getBlock(), mlir::Operation::isBeforeInBlock(), and LDBG.

Referenced by mlir::linalg::LinalgCopyVTRForwardingPattern::matchAndRewrite(), and mlir::linalg::LinalgCopyVTWForwardingPattern::matchAndRewrite().

◆ ofrToIndexValues()

static SmallVector<Value> ofrToIndexValues ( OpBuilder builder,
Location  loc,
ArrayRef< OpFoldResult ofrs 
)
static

Given an ArrayRef of OpFoldResults, return a vector of Values.

IntegerAttrs are converted to ConstantIndexOps. Other attribute types are not supported.

Definition at line 763 of file Vectorization.cpp.

References mlir::OpBuilder::create(), and getIntFromAttr().

Referenced by PadOpVectorizationWithInsertSlicePattern::rewriteUser(), and GenericPadOpVectorizationPattern::tryVectorizeCopy().

◆ reduceIfNeeded()

static Operation* reduceIfNeeded ( OpBuilder b,
LinalgOp  linalgOp,
Operation op,
Value  reduceValue,
Value  initialValue,
const BlockAndValueMapping bvm 
)
static

Emit reduction operations if the shapes of the value to reduce is different that the result shape.

Definition at line 374 of file Vectorization.cpp.

References buildMultiDimReduce(), mlir::Type::dyn_cast(), getReductionMask(), mlir::Value::getType(), and mlir::BlockAndValueMapping::lookup().

Referenced by vectorizeOneOp().

◆ reductionPreconditions()

static LogicalResult reductionPreconditions ( LinalgOp  op)
static

◆ reindexIndexingMap()

static AffineMap reindexIndexingMap ( AffineMap  map)
static

Given an indexing map coming from a LinalgOp indexing, restricted to a projectedPermutation, compress the unused dimensions to serve as a permutation_map for a vector transfer operation.

For example, given a linalg op such as:

%0 = linalg.generic {
indexing_maps = affine_map<(d0, d1, d2, d3, d4) -> (d4, d0, d2)>,
indexing_maps = affine_map<(d0, d1, d2, d3, d4) -> (d1, d3)>
}
ins(%0 : tensor<2x3x4xf32>)
outs(%1 : tensor<5x6xf32>)

the iteration domain size of the linalg op is 3x5x4x6x2. The first affine map is reindexed to affine_map<(d0, d1, d2) -> (d2, d0, d1)>, the second affine map is reindexed to affine_map<(d0, d1) -> (d0, d1)>.

Definition at line 85 of file Vectorization.cpp.

References mlir::AffineMap::isProjectedPermutation().

Referenced by buildVectorWrite(), and vectorizeAsLinalgGeneric().

◆ tensorExtractVectorizationPrecondition()

static LogicalResult tensorExtractVectorizationPrecondition ( Operation op)
static

Helper function to check if the tensor.extract can be vectorized by the custom hook vectorizeTensorExtract.

Definition at line 311 of file Vectorization.cpp.

References mlir::failure(), and mlir::success().

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

◆ vectorizeAsLinalgGeneric()

static LogicalResult vectorizeAsLinalgGeneric ( OpBuilder b,
LinalgOp  linalgOp,
SmallVectorImpl< Value > &  newResults 
)
static

Generic vectorization function that rewrites the body of a linalgOp into vector form.

Generic vectorization proceeds as follows:

  1. Verify the linalgOp has one non-empty region.
  2. Values defined above the region are mapped to themselves and will be broadcasted on a per-need basis by their consumers.
  3. Each region argument is vectorized into a vector.transfer_read (or 0-d load). TODO: Reuse opportunities for RAR dependencies. 4a. Register CustomVectorizationHook for YieldOp to capture the results. 4b. Register CustomVectorizationHook for IndexOp to access the iteration indices.
  4. Iteratively call vectorizeOneOp on the region operations.

When broadcastToMaximalCommonShape is set to true, eager broadcasting is performed to the maximal common vector size implied by the linalgOp iteration space. This eager broadcasting is introduced in the permutation_map of the vector.transfer_read operations. The eager broadcasting makes it trivial to detrmine where broadcast, transposes and reductions should occur, without any bookkeeping. The tradeoff is that, in the absence of good canonicalizations, the amount of work increases. This is not deemed a problem as we expect canonicalizations and foldings to aggressively clean up the useless work.

Definition at line 509 of file Vectorization.cpp.

References mlir::Type::cast(), mlir::AffineMap::compose(), mlir::OpBuilder::create(), mlir::failure(), Failure, mlir::BlockArgument::getArgNumber(), mlir::getElementTypeOrSelf(), mlir::Block::getOperations(), mlir::Operation::getResults(), mlir::Value::getType(), mlir::getUsedValuesDefinedAbove(), mlir::inverseAndBroadcastProjectedPermutation(), mlir::inversePermutation(), LDBG, mlir::BlockAndValueMapping::map(), NewOp, VectorizationResult::newOp, reindexIndexingMap(), VectorizationResult::status, mlir::success(), vectorizeLinalgIndex(), vectorizeLinalgYield(), vectorizeOneOp(), and vectorizeTensorExtract().

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

◆ vectorizeConvolution()

static FailureOr< Operation * > vectorizeConvolution ( OpBuilder b,
LinalgOp  convOp 
)
static

Try to vectorize convOp as a convolution.

Helper function to vectorize a LinalgOp with convolution semantics.

Definition at line 1879 of file Vectorization.cpp.

Referenced by VectorizeConvolution::matchAndRewrite(), and mlir::linalg::vectorize().

◆ vectorizeLinalgIndex()

static VectorizationResult vectorizeLinalgIndex ( OpBuilder b,
Operation op,
LinalgOp  linalgOp 
)
static

Helper function to vectorize the index operations of a linalgOp.

Return VectorizationStatus::NewOp to signal the vectorization algorithm that it should map the produced operations. This function is meant to be used as a CustomVectorizationHook.

Definition at line 277 of file Vectorization.cpp.

References mlir::OpBuilder::create(), Failure, mlir::Builder::getIndexType(), mlir::Builder::getIndexVectorAttr(), and NewOp.

Referenced by vectorizeAsLinalgGeneric().

◆ vectorizeLinalgYield()

static VectorizationResult vectorizeLinalgYield ( OpBuilder b,
Operation op,
const BlockAndValueMapping bvm,
LinalgOp  linalgOp,
SmallVectorImpl< Value > &  newResults 
)
static

Helper function to vectorize the terminator of a linalgOp.

New result vector values are appended to newResults. Return VectorizationStatus::NoReplace to signal the vectorization algorithm that it should not try to map produced operations and instead return the results using the newResults vector making them available to the vectorization algorithm for RAUW. This function is meant to be used as a CustomVectorizationHook.

Definition at line 255 of file Vectorization.cpp.

References buildVectorWrite(), mlir::detail::enumerate(), Failure, mlir::BlockAndValueMapping::lookup(), and NoReplace.

Referenced by vectorizeAsLinalgGeneric().

◆ vectorizeOneOp()

static VectorizationResult vectorizeOneOp ( OpBuilder b,
LinalgOp  linalgOp,
Operation op,
const BlockAndValueMapping bvm,
ArrayRef< CustomVectorizationHook customVectorizationHooks 
)
static

Generic vectorization for a single operation op, given already vectorized operands carried by bvm.

Vectorization occurs as follows:

  1. Try to apply any of the customVectorizationHooks and return its result on success.
  2. Clone any constant in the current scope without vectorization: each consumer of the constant will later determine the shape to which the constant needs to be broadcast to.
  3. Fail on any remaining non ElementwiseMappable op. It is the purpose of the customVectorizationHooks to cover such cases.
  4. Clone op in vector form to a vector of shape prescribed by the first operand of maximal rank. Other operands have smaller rank and are broadcast accordingly. It is assumed this broadcast is always legal, otherwise, it means one of the customVectorizationHooks is incorrect.

This function assumes all operands of op have been vectorized and are in the bvm mapping. As a consequence, this function is meant to be called on a topologically-sorted list of ops. This function does not update bvm but returns a VectorizationStatus that instructs the caller what bvm update needs to occur.

Definition at line 410 of file Vectorization.cpp.

References mlir::OpBuilder::clone(), mlir::OpBuilder::create(), Failure, mlir::Operation::getAttrs(), mlir::OperationName::getIdentifier(), mlir::Operation::getLoc(), mlir::Operation::getName(), mlir::Operation::getOperands(), mlir::Operation::getResultTypes(), mlir::OpTrait::hasElementwiseMappableTraits(), LDBG, mlir::BlockAndValueMapping::lookup(), mlir::matchReduction(), NewOp, reduceIfNeeded(), and VectorizationResult::status.

Referenced by vectorizeAsLinalgGeneric().

◆ vectorizeStaticLinalgOpPrecondition()

static LogicalResult vectorizeStaticLinalgOpPrecondition ( linalg::LinalgOp  op,
ArrayRef< CustomVectorizationPrecondition customPreconditions 
)
static

◆ vectorizeTensorExtract()

static VectorizationResult vectorizeTensorExtract ( OpBuilder b,
Operation op,
LinalgOp  linalgOp,
const BlockAndValueMapping bvm 
)
static

Helper function to vectorize the tensor.extract operations.

Returns VectorizationStatus::NewOp to signal the vectorization algorithm that it should map the produced operations. This function is meant to be used as a CustomVectorizationHook.

Definition at line 337 of file Vectorization.cpp.

References mlir::OpBuilder::create(), Failure, mlir::DenseIntElementsAttr::get(), mlir::Builder::getI1Type(), mlir::Builder::getZeroAttr(), mlir::BlockAndValueMapping::lookup(), and NewOp.

Referenced by vectorizeAsLinalgGeneric().