MLIR  20.0.0git
Macros | Functions
Transforms.cpp File Reference
#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>

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::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)
 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)
 

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

If padding value is set, returns a tensor.pad Op for the source tensor, with the output shape matching the output of packOp.

Otherwise, returns the source directly.

This method assumes that all outer dims for this pack Op are 1.

Definition at line 1026 of file Transforms.cpp.

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

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

◆ getPackUnpackNormalizedPerm()

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

◆ getPackUnpackRankReducedPerm()

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

◆ 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:

  1. Appending iteratorTypes[newDim], equal to iteratorTypes[dim].
  2. Appending a newDim to the domain of every indexing map.
  3. 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::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 619 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 632 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().