FusionOnTensors.cpp File Reference

#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Passes.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/Support/LLVM.h"

Include dependency graph for FusionOnTensors.cpp:

Go to the source code of this file.

Functions
static SmallVector< int64_t >	getTiledSliceDims (OpOperand *consumerOperand, ArrayRef< int64_t > tiledLoopDims)
	Returns the tiled slice dimensions given the tiled consumer loop dimensions. More...
static SmallVector< int64_t >	getTiledProducerLoops (OpResult producerResult, ArrayRef< int64_t > tiledSliceDimIndices)
	Given a vector of tiledSliceDimIndices that represent the tiled dimensions of the producer result slice returns the tiled producer loop dimensions. More...
static LinalgOp	getTiledProducer (OpBuilder &b, OpResult producerResult, tensor::ExtractSliceOp sliceOp, ArrayRef< int64_t > tiledSliceDimIndices, ArrayRef< int64_t > tiledProducerLoopIndices, OpOperand *iterArg)
	Returns the producer fused in place of sliceOp. More...

Function Documentation

getTiledProducer()

static LinalgOp getTiledProducer ( OpBuilder &  b, OpResult  producerResult, tensor::ExtractSliceOp  sliceOp, ArrayRef< int64_t >  tiledSliceDimIndices, ArrayRef< int64_t >  tiledProducerLoopIndices, OpOperand *  iterArg  )

static

Returns the producer fused in place of sliceOp.

Tile the producer operands along the tiledSliceDimIndices and clone the producer. Consider the case of fusion of an output tensor:

%1 = producer ins(...) outs(%0)
%2 = consumer ins(...) outs(%1)

When consumer is tiled, %1 appears in the loop iter_args:

%1 = producer ins(...) outs(%0)
%2 = scf.for ... iter_args(%1) .. (%bbarg) {
  %t1 = tensor.extract_slice %bbarg[..]
  %t2 = consumer ins(...) outs(%t1)
  %r = tensor.insert_slice %t2, %bbarg[...]
}

Fusing %1 into the loop requires updating iter_args(%1) to iter_args(%0):

%2 = scf.for ... iter_args(%0) .. (%bbarg) {
  %t0 = tensor.extract_slice %bbarg[..]
  %t1 = producer ins(...) outs(%t0)
  %t2 = consumer ins(...) outs(%t1)
  %r = tensor.insert_slice %t2, %bbarg[...]
}

This transformation is only valid if bbarg is exclusively used by the output ExtractSliceOp / InsertSliceOp pair, which is checked by the fuseProducer method. TODO: instead of check and failure, insert new iter_args each time a producer is fused into a consumer and fold away unused iter_args.

omitPartialTileCheck=

Definition at line 128 of file FusionOnTensors.cpp.

References mlir::IROperand< DerivedT, IRValueT >::get(), mlir::Builder::getIndexAttr(), mlir::OpOperand::getOperandNumber(), mlir::OpResult::getOwner(), mlir::OpResult::getResultNumber(), mlir::linalg::makeTiledShapes(), mlir::linalg::offsetIndices(), mlir::IROperand< DerivedT, IRValueT >::set(), and mlir::OpBuilder::setInsertionPointAfter().

Referenced by mlir::linalg::TileLoopNest::fuseProducer().

getTiledProducerLoops()

static SmallVector<int64_t> getTiledProducerLoops ( OpResult  producerResult, ArrayRef< int64_t >  tiledSliceDimIndices  )

static

Given a vector of tiledSliceDimIndices that represent the tiled dimensions of the producer result slice returns the tiled producer loop dimensions.

Example:

%res = linalg.fill(%cst, %input)
scf.for %i
  scf.for %j
    %slice = tensor.extract_slice %res[%i, %j]

getTiledProducerLoops(res, [0, 1]) returns the loop indices [0, 1].

Definition at line 65 of file FusionOnTensors.cpp.

References mlir::AffineMap::getDimPosition(), mlir::AffineMap::getNumResults(), mlir::OpResult::getOwner(), mlir::OpResult::getResultNumber(), mlir::AffineMap::getSubMap(), and mlir::AffineMap::isProjectedPermutation().

Referenced by mlir::linalg::TileLoopNest::fuseProducer().

getTiledSliceDims()

static SmallVector<int64_t> getTiledSliceDims ( OpOperand *  consumerOperand, ArrayRef< int64_t >  tiledLoopDims  )

static

Returns the tiled slice dimensions given the tiled consumer loop dimensions.

The slice defines a hyper rectangular iteration space and fusing the producer is always possible. However, depending on the consumer indexing map, not all slice elements may be consumed and the tiles may overlap. In these cases, fusion introduces redundant computation.

Definition at line 37 of file FusionOnTensors.cpp.

References mlir::detail::enumerate(), mlir::detail::IROperandBase::getOwner(), and mlir::AffineMap::getResults().

Referenced by mlir::linalg::TileLoopNest::fuseProducer().