#include "mlir/Dialect/Linalg/Transforms/Transforms.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/Utils/Utils.h"
#include "mlir/Dialect/SCF/Transforms/Transforms.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h"
#include "mlir/Dialect/Tensor/Utils/Utils.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/Matchers.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/ScopeExit.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
#include <type_traits>
#include <utility>

Macros
#define	DEBUG_TYPE "linalg-transforms"

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

#define	DBGSNL() (llvm::dbgs() << "\n")

Functions
static bool	hasAtMostOneResultFunctionOfDim (AffineMap map, int64_t dim)
	Return true if `map` has 0 or 1 result function of AffineDimExpr(dim). More...

static std::optional< int64_t >	getFirstResultIndexFunctionOf (AffineMap map, int64_t dim)
	Return the index of the first result of `map` that is a function of AffineDimExpr(dim), std::nullopt otherwise. More...

static FailureOr< SmallVector< std::optional< int64_t > > >	packLinalgMetadataOnce (SmallVectorImpl< AffineMap > &indexingMaps, SmallVectorImpl< utils::IteratorType > &iteratorTypes, int64_t dim)
	Perform one step of packing of a LinalgOp's metadata along `dim` into the `newDim` at `iteratorTypes.size()` by: More...

static RankedTensorType	permuteShape (RankedTensorType tensorType, ArrayRef< int64_t > permutationVector)
	Return a copy of `tensorType` after permutation by `permutationVector`. More...

static LinalgOp	transposeOneLinalgOperandAndReplace (RewriterBase &rewriter, LinalgOp linalgOp, OpOperand &opOperand, ArrayRef< int64_t > permutation, Value transposedValue)
	Return a new GenericOp obtained by transposing opOperand by the permutation vector: More...

static Value	getPackOpSourceOrPaddedSource (OpBuilder &builder, tensor::PackOp packOp)
	Returns a tensor.pad op if padding value is set. More...

static SmallVector< int64_t >	getPackUnpackNormalizedPerm (int rank, ArrayRef< int64_t > perm)

static SmallVector< int64_t >	getPackUnpackRankReducedPerm (ArrayRef< int64_t > shape, ArrayRef< int64_t > innerDimsPos, ArrayRef< int64_t > outerDimsPerm)

Macro Definition Documentation

◆ DBGS

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

Definition at line 45 of file Transforms.cpp.

◆ DBGSNL

#define DBGSNL ( ) (llvm::dbgs() << "\n")

Definition at line 46 of file Transforms.cpp.

◆ DEBUG_TYPE

#define DEBUG_TYPE "linalg-transforms"

Definition at line 40 of file Transforms.cpp.

Function Documentation

◆ getFirstResultIndexFunctionOf()

static std::optional<int64_t> getFirstResultIndexFunctionOf	(	AffineMap	map,
		int64_t	dim
	)

static

Return the index of the first result of map that is a function of AffineDimExpr(dim), std::nullopt otherwise.

Definition at line 102 of file Transforms.cpp.

References mlir::AffineMap::getNumResults(), mlir::AffineMap::getResult(), and mlir::AffineExpr::isFunctionOfDim().

Referenced by packLinalgMetadataOnce().

◆ getPackOpSourceOrPaddedSource()

static Value getPackOpSourceOrPaddedSource	(	OpBuilder &	builder,
		tensor::PackOp	packOp
	)

static

Returns a tensor.pad op if padding value is set.

Otherwise, returns the source directly. The method assumes that the packOp has static shapes.

Definition at line 1039 of file Transforms.cpp.

References mlir::tensor::createPadHighOp(), mlir::get(), and mlir::getConstantIntValue().

Referenced by mlir::linalg::GeneralizeOuterUnitDimsPackOpPattern::matchAndRewrite().

◆ getPackUnpackNormalizedPerm()

static SmallVector<int64_t> getPackUnpackNormalizedPerm	(	int	rank,
		ArrayRef< int64_t >	perm
	)

static

Definition at line 1079 of file Transforms.cpp.

References mlir::detail::enumerate(), and mlir::invertPermutationVector().

Referenced by getPackUnpackRankReducedPerm().

◆ getPackUnpackRankReducedPerm()

static SmallVector<int64_t> getPackUnpackRankReducedPerm	(	ArrayRef< int64_t >	shape,
		ArrayRef< int64_t >	innerDimsPos,
		ArrayRef< int64_t >	outerDimsPerm
	)

static

Definition at line 1093 of file Transforms.cpp.

References getPackUnpackNormalizedPerm().

Referenced by mlir::linalg::GeneralizeOuterUnitDimsPackOpPattern::matchAndRewrite(), and mlir::linalg::GeneralizeOuterUnitDimsUnPackOpPattern::matchAndRewrite().

◆ hasAtMostOneResultFunctionOfDim()

static bool hasAtMostOneResultFunctionOfDim	(	AffineMap	map,
		int64_t	dim
	)

static

Return true if map has 0 or 1 result function of AffineDimExpr(dim).

Definition at line 87 of file Transforms.cpp.

References mlir::AffineMap::getResults().

Referenced by packLinalgMetadataOnce().

◆ packLinalgMetadataOnce()

static FailureOr<SmallVector<std::optional<int64_t> > > packLinalgMetadataOnce	(	SmallVectorImpl< AffineMap > &	indexingMaps,
		SmallVectorImpl< utils::IteratorType > &	iteratorTypes,
		int64_t	dim
	)

static

Perform one step of packing of a LinalgOp's metadata along dim into the newDim at iteratorTypes.size() by:

Appending iteratorTypes[newDim], equal to iteratorTypes[dim].
Appending a newDim to the domain of every indexing map.
For each operand (i.e. for each map in indexingMaps), perform packing by potentially adding a newDim result to map. The preserved invariant is that iteratorTypes.size() is always equal to map.getNumDims() for every map in indexingMaps.

Update indexingMaps and iteratorTypes inplace as one step of the update. Return a vector that records the optional packing for each operand. Return failure if the packed indexing cannot be represented with a LinalgOp.

Further details:

The current implementation of packing (i.e. data tiling) consists of rewriting a linearized strip-mined form into a higher-dimensional access. e.g. consider an access A[I][f(j, k, l)] and packing by 4; we rewrite I into 4 * i + ii, where 0 <= ii < 4. The access is further rewritten as A[i][f(j, k, l)][ii].

This rewrite into higher dimensional access is not possible for general AffineExpr in Linalg atm, it is restricted to an AffineDimExpr: e.g. consider an access A[I + J][f(j, k, l)] and packing by 4; we rewrite I + J into 4 * i + ii + J, where 0 <= ii < 4. The rewrite of the access would be a form not representable in Linalg: A[i + (ii + J) / 4][f(j, k, l)][(ii + J) % 4]. Note however that as J and ii iterate, the accesses do not have a particular alignment, so packing does not achieve alignment in this case

In the future, we may want to consider a mixed-form that allows some alignment in the presence of multiple accesses: A[I][f(j, k, l)] and B[I + J][f(j, k, l)] And would rewrite accesses as: A[i][f(j, k, l)][ii] and B[4 * i + ii + J][f(j, k, l)]

Definition at line 149 of file Transforms.cpp.

References mlir::failure(), mlir::Builder::getAffineDimExpr(), mlir::AffineMap::getContext(), getFirstResultIndexFunctionOf(), mlir::AffineMap::getNumDims(), mlir::AffineMap::getNumResults(), mlir::AffineMap::getResult(), hasAtMostOneResultFunctionOfDim(), mlir::AffineMap::insertResult(), and mlir::AffineMap::shiftDims().

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

◆ permuteShape()

static RankedTensorType permuteShape	(	RankedTensorType	tensorType,
		ArrayRef< int64_t >	permutationVector
	)

static

Return a copy of tensorType after permutation by permutationVector.

Definition at line 630 of file Transforms.cpp.

References mlir::applyPermutationToVector(), and mlir::RankedTensorType::Builder::setShape().

Referenced by transposeOneLinalgOperandAndReplace().

◆ transposeOneLinalgOperandAndReplace()

static LinalgOp transposeOneLinalgOperandAndReplace	(	RewriterBase &	rewriter,
		LinalgOp	linalgOp,
		OpOperand &	opOperand,
		ArrayRef< int64_t >	permutation,
		Value	transposedValue
	)

static

Return a new GenericOp obtained by transposing opOperand by the permutation vector:

the corresponding indexing map is transposed by permutation
the corresponding operand value is replaced by transposedValue linalgOp is replaced by the return op in the process. Asserts that transposedValue is of the proper transposed ShapedType.

Definition at line 643 of file Transforms.cpp.

References mlir::AffineMap::compose(), mlir::OpBuilder::create(), mlir::IROperand< DerivedT, IRValueT >::get(), mlir::Builder::getContext(), mlir::OpOperand::getOperandNumber(), mlir::detail::IROperandBase::getOwner(), mlir::AffineMap::getPermutationMap(), mlir::Operation::getRegion(), mlir::Value::getType(), mlir::ValueRange::getTypes(), permuteShape(), mlir::RewriterBase::replaceOp(), and mlir::Region::takeBody().

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

Macros

Functions

Macro Definition Documentation

◆ DBGS

◆ DBGSNL

◆ DEBUG_TYPE

Function Documentation

◆ getFirstResultIndexFunctionOf()

◆ getPackOpSourceOrPaddedSource()

◆ getPackUnpackNormalizedPerm()

◆ getPackUnpackRankReducedPerm()

◆ hasAtMostOneResultFunctionOfDim()

◆ packLinalgMetadataOnce()

Further details:

◆ permuteShape()

◆ transposeOneLinalgOperandAndReplace()