#include "mlir/Dialect/Vector/Utils/VectorUtils.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/Arith/Utils/Utils.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/IntegerSet.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/Support/DebugLog.h"
#include "llvm/Support/InterleavedRange.h"

Macros
#define	DEBUG_TYPE "vector-utils"

Functions
static bool	areDimsTransposedIn2DSlice (int64_t dim0, int64_t dim1, ArrayRef< int64_t > transp)
	Given the n-D transpose pattern 'transp', return true if 'dim0' and 'dim1' should be transposed with each other within the context of their 2D transposition slice.
static AffineMap	makePermutationMap (ArrayRef< Value > indices, const DenseMap< Operation *, unsigned > &enclosingLoopToVectorDim)
	Constructs a permutation map from memref indices to vector dimension.
template<typename T>
static SetVector< Operation * >	getParentsOfType (Block *block)
	Implementation detail that walks up the parents and records the ones with the specified type.
static SetVector< Operation * >	getEnclosingforOps (Block *block)
	Returns the enclosing AffineForOp, from closest to farthest.
static bool	isMaskTriviallyFoldable (SmallVector< OpFoldResult > &maskSizes, SmallVector< Value > &indices, ArrayRef< int64_t > baseShape, ArrayRef< int64_t > maskShape)
	Determines whether a mask for xfer_read/write is trivially "all true".

Macro Definition Documentation

◆ DEBUG_TYPE

#define DEBUG_TYPE "vector-utils"

Definition at line 35 of file VectorUtils.cpp.

Function Documentation

◆ areDimsTransposedIn2DSlice()

bool areDimsTransposedIn2DSlice	(	int64_t	dim0,
		int64_t	dim1,
		ArrayRef< int64_t >	transp )

static

Given the n-D transpose pattern 'transp', return true if 'dim0' and 'dim1' should be transposed with each other within the context of their 2D transposition slice.

Example 1: dim0 = 0, dim1 = 2, transp = [2, 1, 0] Return true: dim0 and dim1 are transposed within the context of their 2D transposition slice ([1, 0]).

Example 2: dim0 = 0, dim1 = 1, transp = [2, 1, 0] Return true: dim0 and dim1 are transposed within the context of their 2D transposition slice ([1, 0]). Paradoxically, note how dim1 (1) is not transposed within the full context of the transposition.

Example 3: dim0 = 0, dim1 = 1, transp = [2, 0, 1] Return false: dim0 and dim1 are not transposed within the context of their 2D transposition slice ([0, 1]). Paradoxically, note how dim0 (0) and dim1 (1) are transposed within the full context of the of the transposition.

Definition at line 68 of file VectorUtils.cpp.

Referenced by mlir::vector::isTranspose2DSlice().

◆ getEnclosingforOps()

SetVector< Operation * > getEnclosingforOps ( Block * block )

static

Returns the enclosing AffineForOp, from closest to farthest.

Definition at line 179 of file VectorUtils.cpp.

References getParentsOfType().

◆ getParentsOfType()

template<typename T>

SetVector< Operation * > getParentsOfType ( Block * block )

static

Implementation detail that walks up the parents and records the ones with the specified type.

TODO: could also be implemented as a collect parents followed by a filter and made available outside this file.

Definition at line 165 of file VectorUtils.cpp.

References mlir::Block::getParentOp().

Referenced by getEnclosingforOps().

◆ isMaskTriviallyFoldable()

bool isMaskTriviallyFoldable	(	SmallVector< OpFoldResult > &	maskSizes,
		SmallVector< Value > &	indices,
		ArrayRef< int64_t >	baseShape,
		ArrayRef< int64_t >	maskShape )

static

Determines whether a mask for xfer_read/write is trivially "all true".

Given all the inputs required to generate a mask (mask sizes and shapes), and an xfer_read/write operation (indices and the source/destination tensor shape), determines whether the corresponding mask would be trivially foldable (i.e., trivially "all true").

Use this method to avoid generating spurious masks and relying on vectorization post-processing to remove them.

Pre-conditions for a mask to be trivially foldable:

All involved shapes (mask + destination tensor) are static.
All indices are constant.
All mask sizes are constant (including arith.constant).

If the pre-conditions are met, the method checks for each destination dimension d: (1) destDimSize[rankDiff + d] <= maskShape[d] (2) destDimSize[rankDiff + d] <= index[d] + maskSize[d]

rankDiff = rank(dest) - rank(mask).

This method takes a conservative view: it may return false even if the mask is technically foldable.

EXAMPLE 1 (trivially foldable, all shapes match, mask sizes match the shape of the dest tensor): c0 = arith.constant 0 : index mask = vector.create_mask 5, 1 vector.mask mask { vector.transfer_write vecToStore_1, dest{[c0, c0] {in_bounds = [true, true]} : vector<5x1xi32>, tensor<5x1xi32> }

EXAMPLE 2 (not trivially foldable - vector shape exceeds the tensor shape, mask is required to avoid out-of-bounds write): c0 = arith.constant 0 : index mask = vector.create_mask 5, 1 vector.mask mask { vector.transfer_write vecToStore_2, dest[c0, c0] {in_bounds = [true, true]} : vector<8x1xi32>, tensor<5x1xi32> }

Definition at line 358 of file VectorUtils.cpp.

References mlir::getConstantIntValue(), indices, mlir::m_ConstantInt(), and mlir::matchPattern().

Referenced by mlir::vector::createReadOrMaskedRead(), and mlir::vector::createWriteOrMaskedWrite().

◆ makePermutationMap()

AffineMap makePermutationMap	(	ArrayRef< Value >	indices,
		const DenseMap< Operation *, unsigned > &	enclosingLoopToVectorDim )

static

Constructs a permutation map from memref indices to vector dimension.

The implementation uses the knowledge of the mapping of enclosing loop to vector dimension. enclosingLoopToVectorDim carries this information as a map with:

keys representing "vectorized enclosing loops";
values representing the corresponding vector dimension. The algorithm traverses "vectorized enclosing loops" and extracts the at-most-one MemRef index that is invariant along said loop. This index is guaranteed to be at most one by construction: otherwise the MemRef is not vectorizable. If this invariant index is found, it is added to the permutation_map at the proper vector dimension. If no index is found to be invariant, 0 is added to the permutation_map and corresponds to a vector broadcast along that dimension.

Returns an empty AffineMap if enclosingLoopToVectorDim is empty, signalling that no permutation map can be constructed given enclosingLoopToVectorDim.

Examples can be found in the documentation of makePermutationMap, in the header file.

Definition at line 126 of file VectorUtils.cpp.

References mlir::AffineMap::get(), mlir::getAffineConstantExpr(), mlir::getAffineDimExpr(), mlir::affine::getInvariantAccesses(), and indices.

Referenced by vectorizeAffineLoad(), and vectorizeAffineStore().

Macros

Functions

Macro Definition Documentation

◆ DEBUG_TYPE

Function Documentation

◆ areDimsTransposedIn2DSlice()

◆ getEnclosingforOps()

◆ getParentsOfType()

◆ isMaskTriviallyFoldable()

◆ makePermutationMap()