Transforms.cpp File Reference (original) (raw)

Go to the source code of this file.

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::string	stringifyReassocIndices (ReassociationIndicesRef ri)

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, linalg::PackOp packOp)
	If padding value is set, returns a tensor.pad Op for the source tensor, with the output shape matching the output of packOp. 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)

◆ DBGS

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

◆ DBGSNL

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

◆ DEBUG_TYPE

#define DEBUG_TYPE "linalg-transforms"

◆ getFirstResultIndexFunctionOf()

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

◆ getPackOpSourceOrPaddedSource()

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

◆ getPackUnpackNormalizedPerm()

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

◆ getPackUnpackRankReducedPerm()

◆ hasAtMostOneResultFunctionOfDim()

static bool hasAtMostOneResultFunctionOfDim ( AffineMap map, int64_t dim )	static

◆ packLinalgMetadataOnce()

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 154 of file Transforms.cpp.

References 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

◆ stringifyReassocIndices()

◆ 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 633 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().