MLIR  17.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/Analysis/DependenceAnalysis.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/HoistPadding.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>
+ Include dependency graph for Transforms.cpp:

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 FailureOr< ValuepadOperandToSmallestStaticBoundingBox (OpBuilder &b, linalg::LinalgOp opToPad, OpOperand *opOperand, ArrayRef< int64_t > paddingDimensions, ArrayRef< Attribute > paddingValues, ArrayRef< bool > packPaddings)
 Pad the opOperand in the paddingDimensions using the padding value and the nofold flag found in paddingValues and packPaddings, respectively. More...
 
static SmallVector< utils::IteratorType > getNParallelLoopsAttrs (unsigned nParallelLoops)
 
static Value getPackOpSourceOrPaddedSource (OpBuilder &builder, tensor::PackOp packOp)
 Returns a tensor.pad op if padding value is set. More...
 
static SmallVector< int64_t > getPackUnpackNormalizedInnerPerm (int rank, ArrayRef< int64_t > innerDimsPos)
 
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...
 

Macro Definition Documentation

◆ DBGS

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

Definition at line 47 of file Transforms.cpp.

◆ DBGSNL

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

Definition at line 48 of file Transforms.cpp.

◆ DEBUG_TYPE

#define DEBUG_TYPE   "linalg-transforms"

Definition at line 42 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 913 of file Transforms.cpp.

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

Referenced by packLinalgMetadataOnce().

◆ getNParallelLoopsAttrs()

static SmallVector<utils::IteratorType> getNParallelLoopsAttrs ( unsigned  nParallelLoops)
static

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

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

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

◆ getPackUnpackNormalizedInnerPerm()

static SmallVector<int64_t> getPackUnpackNormalizedInnerPerm ( int  rank,
ArrayRef< int64_t >  innerDimsPos 
)
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 898 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 960 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(), mlir::AffineExpr::isa(), and mlir::AffineMap::shiftDims().

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

◆ padOperandToSmallestStaticBoundingBox()

static FailureOr<Value> padOperandToSmallestStaticBoundingBox ( OpBuilder b,
linalg::LinalgOp  opToPad,
OpOperand opOperand,
ArrayRef< int64_t >  paddingDimensions,
ArrayRef< Attribute paddingValues,
ArrayRef< bool >  packPaddings 
)
static

Pad the opOperand in the paddingDimensions using the padding value and the nofold flag found in paddingValues and packPaddings, respectively.

Exit early and return the opOperand value if the shape dimensions that match paddingDimensions have a static size and the nofold flag is not set. Otherwise, try to pad the shape dimensions that match the iterator dimensions paddingDimensions and return the tensor::PadOp result if padding succeeds or failure otherwise.

Definition at line 77 of file Transforms.cpp.

References mlir::Value::cast(), mlir::OpBuilder::create(), DBGS, mlir::Attribute::dyn_cast(), mlir::detail::enumerate(), mlir::failed(), mlir::failure(), mlir::IROperand< DerivedT, IRValueT >::get(), mlir::linalg::getConstantUpperBoundForIndex(), mlir::Value::getDefiningOp(), mlir::getElementTypeOrSelf(), mlir::OpOperand::getOperandNumber(), mlir::OpResult::getResultNumber(), mlir::AffineMap::getResults(), mlir::Builder::getType(), and mlir::linalg::makeComposedPadHighOp().

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

◆ permuteShape()

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

Return a copy of tensorType after permutation by permutationVector.

Definition at line 1173 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 1186 of file Transforms.cpp.

References mlir::Type::cast(), 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().