MLIR  15.0.0git
Classes | Macros | Functions
Fusion.cpp File Reference
#include "PassDetail.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
#include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.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/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Dominance.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/RegionUtils.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/ScopeExit.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include <set>
+ Include dependency graph for Fusion.cpp:

Go to the source code of this file.

Classes

struct  ShapeDimension
 Implements a simple high-level fusion pass on linalg structured operations. More...
 

Macros

#define DEBUG_TYPE   "linalg-fusion"
 

Functions

static ShapeDimension getShapeDefiningLoopRange (LinalgOp op, unsigned loopDepth, bool fromSubViewOpOnly=false)
 
static SmallVector< ValuegetTiledOperands (LinalgOp producer)
 
static LinalgOp fuse (OpBuilder &b, LinalgOp producer, const DenseMap< unsigned, Range > &fusedLoopsAndRanges)
 Fuses the producer by cloning the producer. More...
 
static Range getRangeFromOperandShape (OpBuilder &b, Location loc, Value shapedOperand, unsigned dim)
 Get the loop range for a dimension dim based on the shapedOperand. More...
 
static LinalgOp fuse (OpBuilder &b, LinalgOp producerOp, AffineMap producerMap, OpOperand &consumerOpOperand)
 Fuses the producer into the loop immediately enclosing the consumer. More...
 
static bool isStructurallyFusableProducer (LinalgOp producer, Value consumedView, LinalgOp consumer)
 
static FailureOr< LinalgDependenceGraph::LinalgDependenceGraphElemfindFusableProducer (OpOperand &consumerOpOperand, const LinalgDependenceGraph &dependenceGraph)
 For consumer with buffer semantics, find the Linalg operation on buffers that is the last writer of consumerOpOperand. More...
 
static void getProducerOfTensor (Value tensor, OpResult &opResult)
 Walk back use-def chain through scf::For yields. More...
 
static AffineMap pruneReductionDimsFromMap (ArrayRef< Attribute > iteratorTypes, AffineMap map)
 Prune all dimensions that are of reduction iterator type from map. More...
 
static FailureOr< AffineMapgetConsumerLoopToProducerLoopMap (LinalgDependenceGraph::LinalgDependenceGraphElem dependence)
 Returns the mapping from iterations in the consumer that write to the same location as the iterations in the producer. More...
 
static bool doesTransposeAccess (AffineMap map, const std::set< unsigned > &fusableLoops)
 Given a projected permutation map, returns true if the map changes the order in which the fused loop dimension appear. More...
 
static std::set< unsignedcollectFusableLoops (ArrayRef< LinalgOp > ops, const FusableOpDependencesTy &fusableDependences)
 Returns the positions of the loop in op that can be tiled based on the operations that are to be fused with it. More...
 
static FailureOr< TiledLinalgOptileRootOperation (OpBuilder &b, LinalgOp op, ArrayRef< Value > tileSizeVector, const LinalgTilingOptions &options, const std::set< unsigned > &fusedLoops)
 Tile the fused loops in the root operation, by setting the tile sizes for all other loops to zero (those will be tiled later). More...
 
static SmallVector< LinalgOp, 1 > fuseOperations (OpBuilder &b, LinalgOp rootOp, TiledLinalgOp tiledLinalgOp, ArrayRef< LinalgOp > fusionCandidates, const FusableOpDependencesTy &fusableDependences, const std::set< unsigned > &fusedLoops)
 Fuse the operations in fusionCandidates with tiledOp. More...
 
static FailureOr< TiledAndFusedLinalgOpstileAndFuseLinalgOpsImpl (OpBuilder &b, ArrayRef< LinalgOp > ops, const LinalgDependenceGraph &dependenceGraph, const LinalgTilingOptions &tilingOptions)
 

Macro Definition Documentation

◆ DEBUG_TYPE

#define DEBUG_TYPE   "linalg-fusion"

Definition at line 36 of file Fusion.cpp.

Function Documentation

◆ collectFusableLoops()

static std::set<unsigned> collectFusableLoops ( ArrayRef< LinalgOp >  ops,
const FusableOpDependencesTy fusableDependences 
)
static

Returns the positions of the loop in op that can be tiled based on the operations that are to be fused with it.

For example, in a

linalg.matmul ins(a, b : ...) outs(c : ...)

if the producer of a needs to be fused with this op, only the i loop of the matmul can be tiled while fusing. If producer of a, and b are to be fused, then no loops can be tiled while fusing. The conditions used are:

  1. Only parallel loops can be used for tile + fuse. Find the number of common outer parallel loops between the op and its producers being fused.
  2. Of the parallel loops only some can be fused. Only those loops can be fused such where the fusable loops iteration space only touches one tile of the fused operation. This is because the producer (which is writing the fused subview) has update semantics.

Since an inverse computation is needed, we need to consider the projection of the producerIndexMap w.r.t the parallel loops. The actual fusable loops are the dimensions of the consumerLoopToProducerLoop map that correspond to parallel loops and appear in the result of the map

Example 1: linalg.fill(cst, c) linalg.matmul ins(a, b) outs(c) Number of parallel loops : 2 producerIndexMap = affine_map<(i, j) ->(i , j)> consumerIndexMap = affine_map<(i, j, k) -> (i, j)> consumerLoopToProducerLoop = affine_map<(i, j, k) -> (i, j)> Fused dimensions : i, j

Example 2: linalg.matmul ins(a, b) outs(c) linalg.generic {indexing_maps = [affine_map<(i, j) -> (j, i)>, ... iterator_types = ["parallel", "parallel"]} ins(c) ...

Number of parallel loops = 2: producerIndexMap (projected to parallel loops) = affine_map<(i, j) -> (i, j)> consumerLoopToProducerLoop2 = affine_map<(i, j) -> (j, i)> Fused dimensions : i, j

Example 3: memref.copy(s, b) linalg.matmul ins(a, b) outs(c)

Number of parallel loops = 2 produceIndexMap : affine_map<(i, j) -> (i, j)> consumerLoopToProduceLoops = affine_map<(i, j, k) -> (k, j)> submap with only parallel loops = affine_map<(i, j) -> (j)> Fused dimensions : j

Definition at line 593 of file Fusion.cpp.

References mlir::Attribute::cast(), doesTransposeAccess(), getConsumerLoopToProducerLoopMap(), mlir::getParallelIteratorTypeName(), and min().

Referenced by tileAndFuseLinalgOpsImpl().

◆ doesTransposeAccess()

static bool doesTransposeAccess ( AffineMap  map,
const std::set< unsigned > &  fusableLoops 
)
static

Given a projected permutation map, returns true if the map changes the order in which the fused loop dimension appear.

Definition at line 523 of file Fusion.cpp.

References mlir::AffineMap::getResults().

Referenced by collectFusableLoops().

◆ findFusableProducer()

static FailureOr<LinalgDependenceGraph::LinalgDependenceGraphElem> findFusableProducer ( OpOperand consumerOpOperand,
const LinalgDependenceGraph dependenceGraph 
)
static

For consumer with buffer semantics, find the Linalg operation on buffers that is the last writer of consumerOpOperand.

For now the fusable dependence is returned as an instance of the dependenceGraph.

Definition at line 277 of file Fusion.cpp.

References mlir::failure(), mlir::IROperand< DerivedT, IRValueT >::get(), mlir::linalg::LinalgDependenceGraph::getDependencesInto(), mlir::linalg::LinalgDependenceGraph::getDependenceTypeStr(), mlir::OpOperand::getOperandNumber(), mlir::detail::IROperandBase::getOwner(), and mlir::linalg::isFusableInto().

Referenced by mlir::linalg::findAllFusableDependences(), and mlir::linalg::fuseProducerOfBuffer().

◆ fuse() [1/2]

static LinalgOp fuse ( OpBuilder b,
LinalgOp  producer,
const DenseMap< unsigned, Range > &  fusedLoopsAndRanges 
)
static

Fuses the producer by cloning the producer.

The fusedLoopsAndRanges provides the loop range information for the fused loops. The rest are obtained from the producer itself, since they are not tiled + fused.

omitPartialTileCheck=

Definition at line 114 of file Fusion.cpp.

References mlir::linalg::addTileLoopIvsToIndexOpResults(), mlir::Operation::clone(), mlir::OpBuilder::create(), mlir::linalg::createOrFoldDimOp(), getShapeDefiningLoopRange(), getTiledOperands(), mlir::linalg::makeTiledShapes(), and mlir::Range::offset.

Referenced by fuse(), fuseOperations(), mlir::linalg::fuseProducerOfBuffer(), and mlir::linalg::fuseProducerOfTensor().

◆ fuse() [2/2]

static LinalgOp fuse ( OpBuilder b,
LinalgOp  producerOp,
AffineMap  producerMap,
OpOperand consumerOpOperand 
)
static

Fuses the producer into the loop immediately enclosing the consumer.

This is achieved by "recomputing" the producer at the time it is needed just before the consumer.

Definition at line 193 of file Fusion.cpp.

References mlir::Value::cast(), mlir::detail::enumerate(), fuse(), mlir::IROperand< DerivedT, IRValueT >::get(), mlir::Operation::getLoc(), mlir::detail::IROperandBase::getOwner(), getRangeFromOperandShape(), and mlir::AffineMap::getResults().

◆ fuseOperations()

static SmallVector<LinalgOp, 1> fuseOperations ( OpBuilder b,
LinalgOp  rootOp,
TiledLinalgOp  tiledLinalgOp,
ArrayRef< LinalgOp >  fusionCandidates,
const FusableOpDependencesTy fusableDependences,
const std::set< unsigned > &  fusedLoops 
)
static

Fuse the operations in fusionCandidates with tiledOp.

Latter is expected to be a tiled operation such that it is valid to fuse all operations in fusionCandidates, i.e. move the operation within the inter-tile loops of tiledOp.

Definition at line 741 of file Fusion.cpp.

References ShapeDimension::dimension, mlir::detail::enumerate(), fuse(), getRangeFromOperandShape(), getShapeDefiningLoopRange(), mlir::Value::getType(), mlir::Type::isa(), mlir::linalg::TiledLinalgOp::loops, mlir::linalg::TiledLinalgOp::op, resultIndex(), mlir::OpBuilder::setInsertionPoint(), and ShapeDimension::shape.

Referenced by tileAndFuseLinalgOpsImpl().

◆ getConsumerLoopToProducerLoopMap()

static FailureOr<AffineMap> getConsumerLoopToProducerLoopMap ( LinalgDependenceGraph::LinalgDependenceGraphElem  dependence)
static

Returns the mapping from iterations in the consumer that write to the same location as the iterations in the producer.

To do so use

  • indexing map of the fused view in the consumer : consumerIndexMap
  • indexing map of the fused view in the producer : producerIndexMap consumerLoopToProducerLoop = inverse(producerIndexMap).compose(consumerIndexMap)

Definition at line 481 of file Fusion.cpp.

References mlir::failure(), mlir::linalg::LinalgDependenceGraph::LinalgDependenceGraphElem::getDependentOp(), mlir::linalg::LinalgDependenceGraph::LinalgDependenceGraphElem::getDependentOpViewIndexingMap(), mlir::linalg::LinalgDependenceGraph::LinalgDependenceGraphElem::getIndexingOpViewIndexingMap(), mlir::AffineMap::getNumResults(), mlir::inversePermutation(), mlir::AffineMap::isPermutation(), mlir::AffineMap::print(), and pruneReductionDimsFromMap().

Referenced by collectFusableLoops().

◆ getProducerOfTensor()

static void getProducerOfTensor ( Value  tensor,
OpResult opResult 
)
static

Walk back use-def chain through scf::For yields.

Sets producer and outputIndex if it finds a producer LinalgOp

Definition at line 379 of file Fusion.cpp.

References mlir::Value::cast(), mlir::Value::dyn_cast(), mlir::Value::getDefiningOp(), mlir::Value::getType(), and mlir::Type::isa().

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

◆ getRangeFromOperandShape()

static Range getRangeFromOperandShape ( OpBuilder b,
Location  loc,
Value  shapedOperand,
unsigned  dim 
)
static

Get the loop range for a dimension dim based on the shapedOperand.

It is expected to be defined by a subview op or an extract_slice op.

Definition at line 180 of file Fusion.cpp.

References mlir::Value::getDefiningOp().

Referenced by fuse(), and fuseOperations().

◆ getShapeDefiningLoopRange()

static ShapeDimension getShapeDefiningLoopRange ( LinalgOp  op,
unsigned  loopDepth,
bool  fromSubViewOpOnly = false 
)
static

◆ getTiledOperands()

static SmallVector<Value> getTiledOperands ( LinalgOp  producer)
static

Definition at line 107 of file Fusion.cpp.

Referenced by fuse().

◆ isStructurallyFusableProducer()

static bool isStructurallyFusableProducer ( LinalgOp  producer,
Value  consumedView,
LinalgOp  consumer 
)
static

Definition at line 208 of file Fusion.cpp.

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

◆ pruneReductionDimsFromMap()

static AffineMap pruneReductionDimsFromMap ( ArrayRef< Attribute iteratorTypes,
AffineMap  map 
)
static

Prune all dimensions that are of reduction iterator type from map.

Definition at line 465 of file Fusion.cpp.

References mlir::detail::enumerate(), mlir::getProjectedMap(), and mlir::isParallelIterator().

Referenced by getConsumerLoopToProducerLoopMap().

◆ tileAndFuseLinalgOpsImpl()

static FailureOr<TiledAndFusedLinalgOps> tileAndFuseLinalgOpsImpl ( OpBuilder b,
ArrayRef< LinalgOp >  ops,
const LinalgDependenceGraph dependenceGraph,
const LinalgTilingOptions tilingOptions 
)
static

◆ tileRootOperation()

static FailureOr<TiledLinalgOp> tileRootOperation ( OpBuilder b,
LinalgOp  op,
ArrayRef< Value tileSizeVector,
const LinalgTilingOptions options,
const std::set< unsigned > &  fusedLoops 
)
static

Tile the fused loops in the root operation, by setting the tile sizes for all other loops to zero (those will be tiled later).

Definition at line 721 of file Fusion.cpp.

References mlir::OpBuilder::create(), options, mlir::linalg::LinalgTilingOptions::setTileSizes(), and mlir::linalg::tileLinalgOp().

Referenced by tileAndFuseLinalgOpsImpl().