MLIR  19.0.0git
Classes | Macros | Typedefs | Enumerations | Functions
Vectorization.cpp File Reference
#include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Analysis/SliceAnalysis.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/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/Tensor/Utils/Utils.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Interfaces/MaskableOpInterface.h"
#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/RegionUtils.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/Sequence.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/ADT/iterator_range.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/MathExtras.h"
#include "llvm/Support/raw_ostream.h"
#include <optional>
#include <type_traits>

Go to the source code of this file.

Classes

struct  VectorizationState
 Contains the vectorization state and related methods used across the vectorization process of a given operation. More...
 
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 << "\n")
 

Typedefs

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

Enumerations

enum class  Conv1DOpOrder { W , 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...
 
enum  VectorMemoryAccessKind { ScalarBroadcast , Contiguous , Gather }
 

Functions

static FailureOr< Operation * > vectorizeConvolution (RewriterBase &rewriter, LinalgOp convOp, ArrayRef< int64_t > inputVecSizes={}, ArrayRef< bool > inputVecScalableFlags={}, bool flatten1DDepthwiseConv=false)
 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 SmallVector< ValueextractConvInputSlices (RewriterBase &rewriter, Location loc, Value input, int64_t nSize, int64_t wSize, int64_t cSize, int64_t kwSize, int strideW, int dilationW, int64_t wSizeStep, bool isSingleChanneled)
 Helper function to extract the input slices after filter is unrolled along kw. More...
 
static SmallVector< ValueextractConvFilterSlices (RewriterBase &rewriter, Location loc, Value filter, int64_t kwSize)
 Helper function to extract the filter slices after filter is unrolled along kw. More...
 
static SmallVector< ValueextractConvResultSlices (RewriterBase &rewriter, Location loc, Value res, int64_t nSize, int64_t wSize, int64_t fSize, int64_t wSizeStep, bool isSingleChanneled)
 Helper function to extract the result slices after filter is unrolled along kw. More...
 
static Value insertConvResultSlices (RewriterBase &rewriter, Location loc, Value res, int64_t wSize, int64_t wSizeStep, SmallVectorImpl< Value > &resVals, bool isSingleChanneled)
 Helper function to insert the computed result slices. 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, Type dstType)
 Broadcast value to a vector of shape if possible. More...
 
static OperationbuildMultiDimReduce (OpBuilder &b, Operation *reduceOp, Value valueToReduce, Value acc, ArrayRef< bool > dimsToMask)
 Create MultiDimReductionOp to compute the reduction for reductionOp. More...
 
static SmallVector< bool > getDimsToReduce (LinalgOp linalgOp)
 
static Value buildVectorWrite (RewriterBase &rewriter, Value value, OpOperand *outputOperand, VectorizationState &state)
 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 (RewriterBase &rewriter, Operation *op, const IRMapping &bvm, VectorizationState &state, LinalgOp linalgOp, SmallVectorImpl< Value > &newResults)
 Helper function to vectorize the terminator of a linalgOp. More...
 
static VectorizationResult vectorizeLinalgIndex (RewriterBase &rewriter, VectorizationState &state, Operation *op, LinalgOp linalgOp)
 Helper function to vectorize the index operations of a linalgOp. More...
 
static LogicalResult tensorExtractVectorizationPrecondition (Operation *op, bool vectorizeNDExtract)
 Helper function to check if the tensor.extract can be vectorized by the custom hook vectorizeTensorExtract. More...
 
static Value calculateGatherOffset (RewriterBase &rewriter, VectorizationState &state, tensor::ExtractOp extractOp, const IRMapping &bvm)
 Calculates the offsets ($index_vec) for vector.gather operations generated from tensor.extract. More...
 
static bool isLoopInvariantIdx (LinalgOp &linalgOp, Value &val)
 Checks whether /p val can be used for calculating a loop invariant index. More...
 
static bool isContiguousLoadIdx (LinalgOp &linalgOp, Value &val, bool &foundIndexOp)
 Check whether val could be used for calculating the trailing index for a contiguous load operation. More...
 
static VectorMemoryAccessKind getTensorExtractMemoryAccessPattern (tensor::ExtractOp extractOp, LinalgOp &linalgOp)
 Infer the memory access pattern for the input ExtractOp. More...
 
static VectorizationResult vectorizeTensorExtract (RewriterBase &rewriter, VectorizationState &state, Operation *op, LinalgOp linalgOp, const IRMapping &bvm)
 Helper function to vectorize the tensor.extract operations. More...
 
static OperationreduceIfNeeded (OpBuilder &b, LinalgOp linalgOp, Operation *op, Value reduceValue, Value initialValue, const IRMapping &bvm)
 Emit reduction operations if the shapes of the value to reduce is different that the result shape. More...
 
static VectorizationResult vectorizeOneOp (RewriterBase &rewriter, VectorizationState &state, LinalgOp linalgOp, Operation *op, const IRMapping &bvm, ArrayRef< CustomVectorizationHook > customVectorizationHooks)
 Generic vectorization for a single operation op, given already vectorized operands carried by bvm. More...
 
static LogicalResult vectorizeAsLinalgGeneric (RewriterBase &rewriter, VectorizationState &state, LinalgOp linalgOp, SmallVectorImpl< Value > &newResults)
 Generic vectorization function that rewrites the body of a linalgOp into vector form. More...
 
static SmallVector< int64_t > getTiledPackShape (tensor::PackOp packOp, ArrayRef< int64_t > destShape)
 Given a tensor::PackOp, return the dest shape before any packing permutations. More...
 
static Value createReadOrMaskedRead (OpBuilder &builder, Location loc, Value source, ArrayRef< int64_t > readShape, Value padValue, bool doMasking=true)
 Create a TransferReadOp from source with static shape readShape. More...
 
static OperationcreateWriteOrMaskedWrite (OpBuilder &builder, Location loc, Value input, SmallVector< OpFoldResult > destSizes, ArrayRef< int64_t > inputVectorSizes)
 Given an input, the mixed destSizes, and the vector sizes for vectorization, create an empty destination tensor and create a TransferWriteOp from the input to the empty tensor. More...
 
static LogicalResult vectorizeAsTensorPackOp (RewriterBase &rewriter, tensor::PackOp packOp, ArrayRef< int64_t > inputVectorSizes, SmallVectorImpl< Value > &newResults)
 Vectorize tensor::PackOp with (1) static innerTiles (2) constant padding value and (3) input vector sizes into: masked_transfer_read->shape_cast->transpose->transfer_write_in_bounds As in the following example: pack = tensor.pack src inner_dims_pos = [2, 1] inner_tiles = [16, 2] into dst : tensor<32x8x16xf32> -> tensor<32x4x1x16x2xf32> More...
 
static LogicalResult vectorizeAsTensorUnpackOp (RewriterBase &rewriter, tensor::UnPackOp unpackOp, ArrayRef< int64_t > inputVectorSizes, SmallVectorImpl< Value > &newResults)
 Vectorize a tensor::UnPackOp to these 4 Ops: Vector::TransferReadOp - Reads a vector from the source tensor vector::TransposeOp - Transpose the Source tensor ShapeCastOp - Reshape the data based on the target. More...
 
static LogicalResult vectorizeAsTensorPadOp (RewriterBase &rewriter, tensor::PadOp padOp, ArrayRef< int64_t > inputVectorSizes, SmallVectorImpl< Value > &newResults)
 Vectorize a padOp with (1) static result type, (2) constant padding value and (3) all-zero lowPad to transfer_write_in_bounds(transfer_read_masked(pad_source, pad_value)). More...
 
static LogicalResult reductionPreconditions (LinalgOp op)
 
static LogicalResult vectorizeDynamicConvOpPrecondition (linalg::LinalgOp conv, bool flatten1DDepthwiseConv)
 
static LogicalResult vectorizeDynamicLinalgOpPrecondition (linalg::LinalgOp op, bool flatten1DDepthwiseConv)
 
static LogicalResult isValidMaskedInputVector (ArrayRef< int64_t > shape, ArrayRef< int64_t > inputVectorSizes)
 Returns success if inputVectorSizes is a valid masking configuraion for given shape, i.e., it meets: More...
 
static LogicalResult vectorizeUnPackOpPrecondition (tensor::UnPackOp unpackOp, ArrayRef< int64_t > inputVectorSizes)
 Need to check if the inner-tiles are static/constant. More...
 
static LogicalResult vectorizeLinalgOpPrecondition (LinalgOp linalgOp, ArrayRef< int64_t > inputVectorSizes, bool vectorizeNDExtract, bool flatten1DDepthwiseConv)
 
static LogicalResult vectorizePackOpPrecondition (tensor::PackOp packOp, ArrayRef< int64_t > inputVectorSizes)
 
static LogicalResult vectorizePadOpPrecondition (tensor::PadOp padOp, ArrayRef< int64_t > inputVectorSizes)
 
static LogicalResult vectorizeScalableVectorPrecondition (Operation *op, ArrayRef< int64_t > inputVectorSizes, ArrayRef< bool > inputScalableVecDims)
 Preconditions for scalable vectors. More...
 
static void convertAffineApply (RewriterBase &rewriter, LinalgOp linalgOp)
 Converts affine.apply Ops to arithmetic operations. More...
 
static int64_t getIntFromAttr (Attribute attr)
 Helper function that retrieves the value of an IntegerAttr. More...
 
static SmallVector< ValueofrToIndexValues (RewriterBase &rewriter, 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 52 of file Vectorization.cpp.

◆ DEBUG_TYPE

#define DEBUG_TYPE   "linalg-vectorization"

Definition at line 50 of file Vectorization.cpp.

◆ LDBG

#define LDBG (   X)    LLVM_DEBUG(DBGS() << X << "\n")

Definition at line 53 of file Vectorization.cpp.

Typedef Documentation

◆ CustomVectorizationHook

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

Definition at line 653 of file Vectorization.cpp.

◆ CustomVectorizationPrecondition

using CustomVectorizationPrecondition = std::function<LogicalResult(Operation *, bool)>

Definition at line 647 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 481 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 489 of file Vectorization.cpp.

◆ VectorMemoryAccessKind

Enumerator
ScalarBroadcast 
Contiguous 
Gather 

Definition at line 796 of file Vectorization.cpp.

Function Documentation

◆ bindShapeDims() [1/3]

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

Definition at line 2782 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 2792 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 2785 of file Vectorization.cpp.

References bindShapeDims().

◆ broadcastIfNeeded()

static Value broadcastIfNeeded ( OpBuilder b,
Value  value,
Type  dstType 
)
static

Broadcast value to a vector of shape if possible.

Return value otherwise.

Definition at line 554 of file Vectorization.cpp.

References mlir::OpBuilder::createOrFold(), mlir::OpBuilder::getInsertionPoint(), mlir::Value::getType(), mlir::vector::isBroadcastableTo(), and mlir::vector::Success.

◆ buildMultiDimReduce()

static Operation* buildMultiDimReduce ( OpBuilder b,
Operation reduceOp,
Value  valueToReduce,
Value  acc,
ArrayRef< bool >  dimsToMask 
)
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.buildMultiDimReduce

Definition at line 571 of file Vectorization.cpp.

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

◆ buildVectorWrite()

static Value buildVectorWrite ( RewriterBase rewriter,
Value  value,
OpOperand outputOperand,
VectorizationState state 
)
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 591 of file Vectorization.cpp.

◆ calculateGatherOffset()

static Value calculateGatherOffset ( RewriterBase rewriter,
VectorizationState state,
tensor::ExtractOp  extractOp,
const IRMapping bvm 
)
static

Calculates the offsets ($index_vec) for vector.gather operations generated from tensor.extract.

The offset is calculated as follows (example using scalar values):

offset = extractOp.indices[0] for (i = 1; i < numIndices; i++) offset = extractOp.dimSize[i] * offset + extractOp.indices[i];

For tensor<45 x 80 x 15 x f32> and index [1, 2, 3], this leads to: offset = ( ( 1 ) * 80 + 2 ) * 15 + 3

Definition at line 765 of file Vectorization.cpp.

◆ convertAffineApply()

static void convertAffineApply ( RewriterBase rewriter,
LinalgOp  linalgOp 
)
static

Converts affine.apply Ops to arithmetic operations.

Definition at line 2011 of file Vectorization.cpp.

◆ createReadOrMaskedRead()

static Value createReadOrMaskedRead ( OpBuilder builder,
Location  loc,
Value  source,
ArrayRef< int64_t >  readShape,
Value  padValue,
bool  doMasking = true 
)
static

Create a TransferReadOp from source with static shape readShape.

If the vector type for the read is not the same as the type of source, then a mask is created on the read. If doMasking parameter is set to false we update the inBounds attribute instead of masking.

Definition at line 1417 of file Vectorization.cpp.

References mlir::OpBuilder::create(), mlir::get(), mlir::Builder::getI1Type(), mlir::tensor::getMixedSizes(), mlir::Operation::getResult(), mlir::Value::getType(), and mlir::vector::maskOperation().

Referenced by vectorizeAsTensorPackOp(), vectorizeAsTensorPadOp(), and vectorizeAsTensorUnpackOp().

◆ createWriteOrMaskedWrite()

static Operation* createWriteOrMaskedWrite ( OpBuilder builder,
Location  loc,
Value  input,
SmallVector< OpFoldResult destSizes,
ArrayRef< int64_t >  inputVectorSizes 
)
static

Given an input, the mixed destSizes, and the vector sizes for vectorization, create an empty destination tensor and create a TransferWriteOp from the input to the empty tensor.

If the destination shape is not the same as the inputVectorSizes for the first rank(inputVectorSizes) dims, then create a mask for the write.

Definition at line 1458 of file Vectorization.cpp.

References mlir::OpBuilder::create(), mlir::get(), mlir::Builder::getI1Type(), mlir::Value::getType(), and mlir::vector::maskOperation().

Referenced by vectorizeAsTensorPackOp(), vectorizeAsTensorPadOp(), and vectorizeAsTensorUnpackOp().

◆ extractConvFilterSlices()

static SmallVector<Value> extractConvFilterSlices ( RewriterBase rewriter,
Location  loc,
Value  filter,
int64_t  kwSize 
)
static

Helper function to extract the filter slices after filter is unrolled along kw.

Definition at line 116 of file Vectorization.cpp.

References mlir::OpBuilder::create().

◆ extractConvInputSlices()

static SmallVector<Value> extractConvInputSlices ( RewriterBase rewriter,
Location  loc,
Value  input,
int64_t  nSize,
int64_t  wSize,
int64_t  cSize,
int64_t  kwSize,
int  strideW,
int  dilationW,
int64_t  wSizeStep,
bool  isSingleChanneled 
)
static

Helper function to extract the input slices after filter is unrolled along kw.

Definition at line 81 of file Vectorization.cpp.

References mlir::OpBuilder::create().

◆ extractConvResultSlices()

static SmallVector<Value> extractConvResultSlices ( RewriterBase rewriter,
Location  loc,
Value  res,
int64_t  nSize,
int64_t  wSize,
int64_t  fSize,
int64_t  wSizeStep,
bool  isSingleChanneled 
)
static

Helper function to extract the result slices after filter is unrolled along kw.

Definition at line 132 of file Vectorization.cpp.

References mlir::OpBuilder::create().

◆ getDimsToReduce()

static SmallVector<bool> getDimsToReduce ( LinalgOp  linalgOp)
static

Definition at line 580 of file Vectorization.cpp.

References mlir::linalg::isReductionIterator().

◆ getIntFromAttr()

static int64_t getIntFromAttr ( Attribute  attr)
static

Helper function that retrieves the value of an IntegerAttr.

Definition at line 2160 of file Vectorization.cpp.

Referenced by ofrToIndexValues().

◆ 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 65 of file Vectorization.cpp.

References mlir::Block::walk().

◆ 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 2632 of file Vectorization.cpp.

References mlir::Value::getUses().

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

◆ getTensorExtractMemoryAccessPattern()

static VectorMemoryAccessKind getTensorExtractMemoryAccessPattern ( tensor::ExtractOp  extractOp,
LinalgOp &  linalgOp 
)
static

Infer the memory access pattern for the input ExtractOp.

Based on the operation shapes and indices (usually based on the iteration space of the parent linalgOp operation), decides whether the input ExtractOp is a contiguous load (including a broadcast of a scalar) or a gather load.

Note that it is always safe to use gather load operations for contiguous loads (albeit slow), but not vice-versa. When in doubt, bail out and assume that extractOp is a gather load.

Definition at line 918 of file Vectorization.cpp.

References Contiguous, mlir::detail::enumerate(), Gather, isContiguousLoadIdx(), isLoopInvariantIdx(), LDBG, and ScalarBroadcast.

◆ getTiledPackShape()

static SmallVector<int64_t> getTiledPackShape ( tensor::PackOp  packOp,
ArrayRef< int64_t >  destShape 
)
static

Given a tensor::PackOp, return the dest shape before any packing permutations.

Definition at line 1408 of file Vectorization.cpp.

References mlir::applyPermutation(), and mlir::tensor::getPackInverseDestPerm().

Referenced by vectorizeAsTensorPackOp().

◆ insertConvResultSlices()

static Value insertConvResultSlices ( RewriterBase rewriter,
Location  loc,
Value  res,
int64_t  wSize,
int64_t  wSizeStep,
SmallVectorImpl< Value > &  resVals,
bool  isSingleChanneled 
)
static

Helper function to insert the computed result slices.

Definition at line 158 of file Vectorization.cpp.

References mlir::OpBuilder::create().

◆ isContiguousLoadIdx()

static bool isContiguousLoadIdx ( LinalgOp &  linalgOp,
Value val,
bool &  foundIndexOp 
)
static

Check whether val could be used for calculating the trailing index for a contiguous load operation.

There are currently 3 types of values that are allowed here:

  1. loop-invariant values,
  2. values that increment by 1 with every loop iteration,
  3. results of basic arithmetic operations (linear and continuous) involving 1., 2. and 3. This method returns True if indeed only such values are used in calculating val.

Additionally, the trailing index for a contiguous load operation should increment by 1 with every loop iteration, i.e. be based on:

  • linalg.index <dim> , where <dim> is the trailing dim of the iteration space. foundIndexOp is updated to true when such an op is found.

Definition at line 860 of file Vectorization.cpp.

References mlir::Value::getDefiningOp().

Referenced by getTensorExtractMemoryAccessPattern().

◆ isLoopInvariantIdx()

static bool isLoopInvariantIdx ( LinalgOp &  linalgOp,
Value val 
)
static

Checks whether /p val can be used for calculating a loop invariant index.

Definition at line 799 of file Vectorization.cpp.

References mlir::Value::getDefiningOp().

Referenced by getTensorExtractMemoryAccessPattern().

◆ isValidMaskedInputVector()

static LogicalResult isValidMaskedInputVector ( ArrayRef< int64_t >  shape,
ArrayRef< int64_t >  inputVectorSizes 
)
static

Returns success if inputVectorSizes is a valid masking configuraion for given shape, i.e., it meets:

  1. The numbers of elements in both array are equal.
  2. inputVectorSizes does not have dynamic dimensions.
  3. All the values in inputVectorSizes are greater than or equal to static sizes in shape.

Definition at line 1791 of file Vectorization.cpp.

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

Referenced by vectorizeLinalgOpPrecondition(), vectorizePackOpPrecondition(), vectorizePadOpPrecondition(), and vectorizeUnPackOpPrecondition().

◆ 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 538 of file Vectorization.cpp.

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

◆ 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 2605 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 ( RewriterBase rewriter,
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 2167 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 IRMapping bvm 
)
static

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

Definition at line 1133 of file Vectorization.cpp.

◆ reductionPreconditions()

static LogicalResult reductionPreconditions ( LinalgOp  op)
static

Definition at line 1720 of file Vectorization.cpp.

Referenced by vectorizeLinalgOpPrecondition().

◆ 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 471 of file Vectorization.cpp.

References mlir::compressUnusedDims(), and mlir::AffineMap::isProjectedPermutation().

◆ tensorExtractVectorizationPrecondition()

static LogicalResult tensorExtractVectorizationPrecondition ( Operation op,
bool  vectorizeNDExtract 
)
static

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

Definition at line 731 of file Vectorization.cpp.

Referenced by vectorizeLinalgOpPrecondition().

◆ vectorizeAsLinalgGeneric()

static LogicalResult vectorizeAsLinalgGeneric ( RewriterBase rewriter,
VectorizationState state,
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. 4rewriter. 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 1283 of file Vectorization.cpp.

◆ vectorizeAsTensorPackOp()

static LogicalResult vectorizeAsTensorPackOp ( RewriterBase rewriter,
tensor::PackOp  packOp,
ArrayRef< int64_t >  inputVectorSizes,
SmallVectorImpl< Value > &  newResults 
)
static

Vectorize tensor::PackOp with (1) static innerTiles (2) constant padding value and (3) input vector sizes into: masked_transfer_read->shape_cast->transpose->transfer_write_in_bounds As in the following example: pack = tensor.pack src inner_dims_pos = [2, 1] inner_tiles = [16, 2] into dst : tensor<32x8x16xf32> -> tensor<32x4x1x16x2xf32>

This pack would be vectorized to:

load = vector.mask mask { vector.transfer_read arg0[c0, c0, c0], cst {in_bounds = [true, true, true]} : tensor<32x7x16xf32>, vector<32x8x16xf32> } : vector<32x8x16xi1> -> vector<32x8x16xf32> shape_cast = vector.shape_cast load : vector<32x8x16xf32> to vector<32x4x2x1x16xf32> transpose = vector.transpose shape_cast, [0, 1, 3, 4, 2] : vector<32x4x2x1x16xf32> to vector<32x4x1x16x2xf32> write = vector.transfer_write transpose, empty[c0_0, c0_0, c0_0, c0_0, c0_0] {in_bounds = [true, true, true, true, true]} : vector<32x4x1x16x2xf32>, tensor<32x4x1x16x2xf32>

If the (3) input vector sizes are not provided, the vector sizes are determined by the result tensor shape. Also, we update the inBounds attribute instead of masking.

Definition at line 1520 of file Vectorization.cpp.

References mlir::applyPermutationToVector(), mlir::OpBuilder::create(), createReadOrMaskedRead(), createWriteOrMaskedWrite(), mlir::detail::enumerate(), mlir::get(), mlir::tensor::getPackInverseDestPerm(), getTiledPackShape(), mlir::Builder::getZeroAttr(), mlir::invertPermutationVector(), mlir::OpBuilder::setInsertionPoint(), mlir::succeeded(), and mlir::success().

◆ vectorizeAsTensorPadOp()

static LogicalResult vectorizeAsTensorPadOp ( RewriterBase rewriter,
tensor::PadOp  padOp,
ArrayRef< int64_t >  inputVectorSizes,
SmallVectorImpl< Value > &  newResults 
)
static

Vectorize a padOp with (1) static result type, (2) constant padding value and (3) all-zero lowPad to transfer_write_in_bounds(transfer_read_masked(pad_source, pad_value)).

Definition at line 1694 of file Vectorization.cpp.

References createReadOrMaskedRead(), createWriteOrMaskedWrite(), mlir::Operation::getResult(), mlir::OpBuilder::setInsertionPoint(), mlir::succeeded(), and mlir::success().

◆ vectorizeAsTensorUnpackOp()

static LogicalResult vectorizeAsTensorUnpackOp ( RewriterBase rewriter,
tensor::UnPackOp  unpackOp,
ArrayRef< int64_t >  inputVectorSizes,
SmallVectorImpl< Value > &  newResults 
)
static

Vectorize a tensor::UnPackOp to these 4 Ops: Vector::TransferReadOp - Reads a vector from the source tensor vector::TransposeOp - Transpose the Source tensor ShapeCastOp - Reshape the data based on the target.

vector::TransferWriteOp. - Write the result vector back to the destination tensor

Definition at line 1591 of file Vectorization.cpp.

References mlir::applyPermutationToVector(), mlir::OpBuilder::create(), createReadOrMaskedRead(), createWriteOrMaskedWrite(), mlir::detail::divideCeil(), mlir::detail::enumerate(), mlir::LogicalResult::failed(), mlir::failure(), mlir::get(), mlir::Operation::getResult(), mlir::tensor::getUnPackInverseSrcPerm(), mlir::Builder::getZeroAttr(), LDBG, mlir::OpBuilder::setInsertionPoint(), and mlir::success().

◆ vectorizeConvolution()

static FailureOr< Operation * > vectorizeConvolution ( RewriterBase rewriter,
LinalgOp  convOp,
ArrayRef< int64_t >  inputVecSizes = {},
ArrayRef< bool >  inputVecScalableFlags = {},
bool  flatten1DDepthwiseConv = false 
)
static

Try to vectorize convOp as a convolution.

Helper function to vectorize a LinalgOp with convolution semantics.

Definition at line 3607 of file Vectorization.cpp.

◆ vectorizeDynamicConvOpPrecondition()

static LogicalResult vectorizeDynamicConvOpPrecondition ( linalg::LinalgOp  conv,
bool  flatten1DDepthwiseConv 
)
static

Definition at line 1740 of file Vectorization.cpp.

References mlir::failure(), mlir::Value::getType(), LDBG, and mlir::success().

◆ vectorizeDynamicLinalgOpPrecondition()

static LogicalResult vectorizeDynamicLinalgOpPrecondition ( linalg::LinalgOp  op,
bool  flatten1DDepthwiseConv 
)
static

Definition at line 1768 of file Vectorization.cpp.

Referenced by vectorizeLinalgOpPrecondition().

◆ vectorizeLinalgIndex()

static VectorizationResult vectorizeLinalgIndex ( RewriterBase rewriter,
VectorizationState state,
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 688 of file Vectorization.cpp.

◆ vectorizeLinalgOpPrecondition()

static LogicalResult vectorizeLinalgOpPrecondition ( LinalgOp  linalgOp,
ArrayRef< int64_t >  inputVectorSizes,
bool  vectorizeNDExtract,
bool  flatten1DDepthwiseConv 
)
static

◆ vectorizeLinalgYield()

static VectorizationResult vectorizeLinalgYield ( RewriterBase rewriter,
Operation op,
const IRMapping bvm,
VectorizationState state,
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 664 of file Vectorization.cpp.

◆ vectorizeOneOp()

static VectorizationResult vectorizeOneOp ( RewriterBase rewriter,
VectorizationState state,
LinalgOp  linalgOp,
Operation op,
const IRMapping 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 1169 of file Vectorization.cpp.

◆ vectorizePackOpPrecondition()

static LogicalResult vectorizePackOpPrecondition ( tensor::PackOp  packOp,
ArrayRef< int64_t >  inputVectorSizes 
)
static

◆ vectorizePadOpPrecondition()

static LogicalResult vectorizePadOpPrecondition ( tensor::PadOp  padOp,
ArrayRef< int64_t >  inputVectorSizes 
)
static

◆ vectorizeScalableVectorPrecondition()

static LogicalResult vectorizeScalableVectorPrecondition ( Operation op,
ArrayRef< int64_t >  inputVectorSizes,
ArrayRef< bool >  inputScalableVecDims 
)
static

Preconditions for scalable vectors.

Definition at line 1964 of file Vectorization.cpp.

◆ vectorizeTensorExtract()

static VectorizationResult vectorizeTensorExtract ( RewriterBase rewriter,
VectorizationState state,
Operation op,
LinalgOp  linalgOp,
const IRMapping 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 1009 of file Vectorization.cpp.

◆ vectorizeUnPackOpPrecondition()

static LogicalResult vectorizeUnPackOpPrecondition ( tensor::UnPackOp  unpackOp,
ArrayRef< int64_t >  inputVectorSizes 
)
static

Need to check if the inner-tiles are static/constant.

Definition at line 1821 of file Vectorization.cpp.

References mlir::failed(), mlir::failure(), isValidMaskedInputVector(), LDBG, and mlir::success().