MLIR  15.0.0git
Functions
ElementwiseOpFusion.cpp File Reference
#include <utility>
#include "PassDetail.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/SparseTensor/IR/SparseTensor.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+ Include dependency graph for ElementwiseOpFusion.cpp:

Go to the source code of this file.

Functions

static AffineMap getIndexingMapOfProducerOperandsInCoordinatesOfFusedOp (OpOperand *producerOpOperand, AffineMap producerResultIndexMap, AffineMap fusedConsumerArgIndexMap)
 Append to fusedOpIndexingMapAttrs the indexing maps for the operands of the producer to use in the fused operation given the indexing map of the result of the producer in the consumer. More...
 
static bool areElementwiseOpsFusable (GenericOp producer, GenericOp consumer, OpOperand *consumerOpOperand)
 Conditions for elementwise fusion of generic operations. More...
 
static void generateFusedElementwiseOpRegion (PatternRewriter &rewriter, GenericOp fusedOp, AffineMap consumerToProducerLoopsMap, OpOperand *consumerOpOperand, unsigned nloops)
 Generate the region of the fused tensor operation. More...
 
static Optional< SmallVector< Value > > fuseElementwiseOpsImpl (GenericOp producer, OpOperand *consumerOpOperand, const ControlFusionFn &controlFn, PatternRewriter &rewriter)
 
static Optional< SmallVector< Value > > fuseElementwiseOps (PatternRewriter &rewriter, OpOperand *consumerOpOperand, GenericOp producer, const ControlFusionFn &controlFn)
 
static bool isFusableWithReshapeByDimExpansion (GenericOp genericOp, OpOperand *fusableOpOperand)
 Conditions for folding a generic operation with a reshape op by expanding the iteration space dimensionality for tensor operations. More...
 
static LogicalResult isGenericOpExpandable (GenericOp genericOp, const ExpansionInfo &expansionInfo, PatternRewriter &rewriter)
 Epanding the body of a linalg operation requires adaptations of the accessed loop indices. More...
 
static AffineMap getIndexingMapInExpandedOp (OpBuilder &builder, AffineMap indexingMap, const ExpansionInfo &expansionInfo)
 Return the indexing map to use in the expanded op for a given the indexingMap of the original operation. More...
 
static RankedTensorType getExpandedType (RankedTensorType originalType, AffineMap indexingMap, const ExpansionInfo &expansionInfo)
 Return the type of the operand/result to use in the expanded op given the type in the original op. More...
 
static SmallVector< ReassociationIndicesgetReassociationForExpansion (AffineMap indexingMap, const ExpansionInfo &expansionInfo)
 Returns the reassociation maps to use in the tensor.expand_shape operation to convert the operands of the original operation to operands of the expanded operation. More...
 
static void updateExpandedGenericOpRegion (PatternRewriter &rewriter, Location loc, Region &fusedRegion, const ExpansionInfo &expansionInfo)
 Update the body of an expanded linalg operation having index semantics. More...
 
static Optional< SmallVector< Value > > fuseWithReshapeByExpansion (GenericOp genericOp, Operation *reshapeOp, OpOperand *fusableOpOperand, PatternRewriter &rewriter)
 Implements the fusion of a tensor.collapse_shape or a tensor.expand_shape op and a generic op as explained in isFusableWithReshapeByExpansion. More...
 
static ReassociationIndices getDomainReassociation (AffineMap indexingMap, ReassociationIndicesRef rangeReassociation)
 For a given list of indices in the range of the indexingMap that are folded, return the indices of the corresponding domain. More...
 
static bool isDimSequencePreserved (AffineMap indexingMap, ReassociationIndicesRef dimSequence)
 For a given dimSequence, check if the sequence is conserved in the indexingMap. More...
 
static SmallVector< ReassociationIndicesgetCollapsableIterationSpaceDims (GenericOp genericOp, OpOperand *fusableOperand, ArrayRef< ReassociationIndices > reassociation)
 
static SmallVector< StringRef > getCollapsedOpIteratorTypes (ArrayRef< Attribute > iteratorTypes, const CollapsingInfo &collapsingInfo)
 Get the iterator types for the collapsed operation given the original iterator types and collapsed dimensions. More...
 
static AffineMap getCollapsedOpIndexingMap (AffineMap indexingMap, const CollapsingInfo &collapsingInfo)
 Compute the indexing map in the collapsed op that corresponds to the given indexingMap of the original operation. More...
 
static SmallVector< ReassociationIndicesgetOperandReassociation (AffineMap indexingMap, const CollapsingInfo &collapsingInfo)
 Return the reassociation indices to use to collapse the operand when the iteration space of a generic op is collapsed. More...
 
static Value getCollapsedOpOperand (Location loc, GenericOp genericOp, OpOperand *opOperand, const CollapsingInfo &collapsingInfo, OpBuilder &builder)
 Get the new value to use for a given OpOperand in the collapsed operation. More...
 
void generateCollapsedIndexingRegion (Location loc, Block *block, const CollapsingInfo &collapsingInfo, ValueRange loopRange, PatternRewriter &rewriter)
 Modify the linalg.index operations in the original generic op, to its value in the collapsed operation. More...
 
static FailureOr< SmallVector< Value > > collapseGenericOpIterationDims (GenericOp genericOp, ArrayRef< ReassociationIndices > foldedIterationDims, OpOperand *fusableOpOperand, PatternRewriter &rewriter)
 Implementation of fusion with reshape operation by collapsing dimensions. More...
 

Function Documentation

◆ areElementwiseOpsFusable()

static bool areElementwiseOpsFusable ( GenericOp  producer,
GenericOp  consumer,
OpOperand consumerOpOperand 
)
static

◆ collapseGenericOpIterationDims()

static FailureOr<SmallVector<Value> > collapseGenericOpIterationDims ( GenericOp  genericOp,
ArrayRef< ReassociationIndices foldedIterationDims,
OpOperand fusableOpOperand,
PatternRewriter rewriter 
)
static

Implementation of fusion with reshape operation by collapsing dimensions.

Definition at line 1330 of file ElementwiseOpFusion.cpp.

References mlir::Region::begin(), mlir::Block::begin(), mlir::RewriterBase::cancelRootUpdate(), mlir::Type::cast(), mlir::RewriterBase::cloneRegionBefore(), mlir::concatAffineMaps(), mlir::OpBuilder::create(), mlir::OpBuilder::createOrFold(), mlir::Attribute::dyn_cast(), mlir::Type::dyn_cast(), mlir::detail::enumerate(), mlir::failed(), mlir::failure(), mlir::RewriterBase::finalizeRootUpdate(), mlir::Block::front(), generateCollapsedIndexingRegion(), mlir::Builder::getAffineMapArrayAttr(), mlir::Block::getArgument(), mlir::Block::getArguments(), getCollapsableIterationSpaceDims(), getCollapsedOpIndexingMap(), getCollapsedOpIteratorTypes(), getCollapsedOpOperand(), mlir::Value::getDefiningOp(), mlir::Builder::getFusedLoc(), mlir::Value::getLoc(), mlir::Operation::getLoc(), getOperandReassociation(), mlir::sparse_tensor::getSparseTensorEncoding(), mlir::DenseElementsAttr::getSplatValue(), mlir::Attribute::getType(), mlir::Value::getType(), mlir::DenseElementsAttr::getType(), mlir::inversePermutation(), mlir::DenseElementsAttr::isSplat(), mlir::m_One(), mlir::m_Zero(), mlir::BlockAndValueMapping::map(), mlir::matchPattern(), mlir::RewriterBase::mergeBlocks(), mlir::RewriterBase::notifyMatchFailure(), mlir::Range::offset, mlir::RewriterBase::replaceOp(), mlir::OpBuilder::setInsertionPoint(), mlir::Operation::setOperand(), mlir::Range::size, mlir::RewriterBase::startRootUpdate(), mlir::Range::stride, and mlir::success().

◆ fuseElementwiseOps()

static Optional<SmallVector<Value> > fuseElementwiseOps ( PatternRewriter rewriter,
OpOperand consumerOpOperand,
GenericOp  producer,
const ControlFusionFn controlFn 
)
static

◆ fuseElementwiseOpsImpl()

static Optional<SmallVector<Value> > fuseElementwiseOpsImpl ( GenericOp  producer,
OpOperand consumerOpOperand,
const ControlFusionFn controlFn,
PatternRewriter rewriter 
)
static

◆ fuseWithReshapeByExpansion()

static Optional<SmallVector<Value> > fuseWithReshapeByExpansion ( GenericOp  genericOp,
Operation reshapeOp,
OpOperand fusableOpOperand,
PatternRewriter rewriter 
)
static

◆ generateCollapsedIndexingRegion()

void generateCollapsedIndexingRegion ( Location  loc,
Block block,
const CollapsingInfo &  collapsingInfo,
ValueRange  loopRange,
PatternRewriter rewriter 
)

Modify the linalg.index operations in the original generic op, to its value in the collapsed operation.

Definition at line 1291 of file ElementwiseOpFusion.cpp.

References mlir::OpBuilder::create(), mlir::detail::enumerate(), mlir::Block::getOps(), mlir::RewriterBase::replaceOp(), and mlir::OpBuilder::setInsertionPointToStart().

Referenced by collapseGenericOpIterationDims().

◆ generateFusedElementwiseOpRegion()

static void generateFusedElementwiseOpRegion ( PatternRewriter rewriter,
GenericOp  fusedOp,
AffineMap  consumerToProducerLoopsMap,
OpOperand consumerOpOperand,
unsigned  nloops 
)
static

◆ getCollapsableIterationSpaceDims()

static SmallVector<ReassociationIndices> getCollapsableIterationSpaceDims ( GenericOp  genericOp,
OpOperand fusableOperand,
ArrayRef< ReassociationIndices reassociation 
)
static

◆ getCollapsedOpIndexingMap()

static AffineMap getCollapsedOpIndexingMap ( AffineMap  indexingMap,
const CollapsingInfo &  collapsingInfo 
)
static

Compute the indexing map in the collapsed op that corresponds to the given indexingMap of the original operation.

Definition at line 1217 of file ElementwiseOpFusion.cpp.

References mlir::AffineMap::get(), mlir::getAffineDimExpr(), mlir::AffineMap::getContext(), mlir::AffineMap::getResults(), and mlir::AffineMap::isProjectedPermutation().

Referenced by collapseGenericOpIterationDims().

◆ getCollapsedOpIteratorTypes()

static SmallVector<StringRef> getCollapsedOpIteratorTypes ( ArrayRef< Attribute iteratorTypes,
const CollapsingInfo &  collapsingInfo 
)
static

Get the iterator types for the collapsed operation given the original iterator types and collapsed dimensions.

Definition at line 1198 of file ElementwiseOpFusion.cpp.

Referenced by collapseGenericOpIterationDims().

◆ getCollapsedOpOperand()

static Value getCollapsedOpOperand ( Location  loc,
GenericOp  genericOp,
OpOperand opOperand,
const CollapsingInfo &  collapsingInfo,
OpBuilder builder 
)
static

Get the new value to use for a given OpOperand in the collapsed operation.

Definition at line 1269 of file ElementwiseOpFusion.cpp.

References mlir::OpBuilder::create(), mlir::IROperand< DerivedT, IRValueT >::get(), mlir::AffineMap::getNumResults(), getOperandReassociation(), and mlir::Operation::getResult().

Referenced by collapseGenericOpIterationDims().

◆ getDomainReassociation()

static ReassociationIndices getDomainReassociation ( AffineMap  indexingMap,
ReassociationIndicesRef  rangeReassociation 
)
static

For a given list of indices in the range of the indexingMap that are folded, return the indices of the corresponding domain.

Return llvm::None on failure. Ensures that all the elements of the returned reassociation are distinct.

Definition at line 898 of file ElementwiseOpFusion.cpp.

References mlir::AffineMap::getResults(), and mlir::AffineMap::isProjectedPermutation().

Referenced by getCollapsableIterationSpaceDims().

◆ getExpandedType()

static RankedTensorType getExpandedType ( RankedTensorType  originalType,
AffineMap  indexingMap,
const ExpansionInfo &  expansionInfo 
)
static

Return the type of the operand/result to use in the expanded op given the type in the original op.

Definition at line 603 of file ElementwiseOpFusion.cpp.

References mlir::AffineMap::getResults().

Referenced by fuseWithReshapeByExpansion().

◆ getIndexingMapInExpandedOp()

static AffineMap getIndexingMapInExpandedOp ( OpBuilder builder,
AffineMap  indexingMap,
const ExpansionInfo &  expansionInfo 
)
static

Return the indexing map to use in the expanded op for a given the indexingMap of the original operation.

Definition at line 585 of file ElementwiseOpFusion.cpp.

References mlir::AffineMap::get(), mlir::Builder::getAffineDimExpr(), mlir::Builder::getContext(), mlir::AffineMap::getNumSymbols(), and mlir::AffineMap::getResults().

Referenced by fuseWithReshapeByExpansion().

◆ getIndexingMapOfProducerOperandsInCoordinatesOfFusedOp()

static AffineMap getIndexingMapOfProducerOperandsInCoordinatesOfFusedOp ( OpOperand producerOpOperand,
AffineMap  producerResultIndexMap,
AffineMap  fusedConsumerArgIndexMap 
)
static

Append to fusedOpIndexingMapAttrs the indexing maps for the operands of the producer to use in the fused operation given the indexing map of the result of the producer in the consumer.

Definition at line 38 of file ElementwiseOpFusion.cpp.

References mlir::AffineMap::compose(), mlir::detail::IROperandBase::getOwner(), and mlir::inversePermutation().

Referenced by areElementwiseOpsFusable(), and fuseElementwiseOpsImpl().

◆ getOperandReassociation()

static SmallVector<ReassociationIndices> getOperandReassociation ( AffineMap  indexingMap,
const CollapsingInfo &  collapsingInfo 
)
static

Return the reassociation indices to use to collapse the operand when the iteration space of a generic op is collapsed.

Definition at line 1242 of file ElementwiseOpFusion.cpp.

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

Referenced by collapseGenericOpIterationDims(), and getCollapsedOpOperand().

◆ getReassociationForExpansion()

static SmallVector<ReassociationIndices> getReassociationForExpansion ( AffineMap  indexingMap,
const ExpansionInfo &  expansionInfo 
)
static

Returns the reassociation maps to use in the tensor.expand_shape operation to convert the operands of the original operation to operands of the expanded operation.

The same method is used to compute the tensor.collapse_shape used to collapse the result of the expanded op to get the value that can replace all uses of the results of the original op.

Definition at line 622 of file ElementwiseOpFusion.cpp.

References mlir::AffineMap::getResults().

Referenced by fuseWithReshapeByExpansion().

◆ isDimSequencePreserved()

static bool isDimSequencePreserved ( AffineMap  indexingMap,
ReassociationIndicesRef  dimSequence 
)
static

For a given dimSequence, check if the sequence is conserved in the indexingMap.

indexingMap is expected to be a projected permutation. Non-existence of the sequence returns true as well.

Definition at line 917 of file ElementwiseOpFusion.cpp.

References mlir::detail::enumerate(), mlir::AffineMap::getNumResults(), mlir::AffineMap::getResult(), mlir::AffineMap::getResults(), and mlir::AffineMap::isProjectedPermutation().

Referenced by getCollapsableIterationSpaceDims().

◆ isFusableWithReshapeByDimExpansion()

static bool isFusableWithReshapeByDimExpansion ( GenericOp  genericOp,
OpOperand fusableOpOperand 
)
static

Conditions for folding a generic operation with a reshape op by expanding the iteration space dimensionality for tensor operations.

These are preconditions assumed by foldReshapeByDimExpansion which implements the following fusion pattern.

Consider

c = linalg.generic ins(a, b : memref<?x?x?xf32>, memref<?x?xf32>) indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0, d2)>, affine_map<(d0, d1, d2) -> (d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d2, d1)>] d = tensor.expand_shape c [[0, 1], [2], [3, 4, 5]] : tensor<?x?x?xf32> into tensor<?x?x?x?x?x?xf32>

The reshape can be folded into the genericOp if its loop dimensionality is increased to match the result (operand) of the tensor.expand_shape. The indexing_map of the fused tensor in the genericOp and the reassociation map helps compute the indexing maps of the modified op. For the above example, based on the reassociation map it can be concluded that

  • The loop used to access the first dimension of the fused tensor is split into two.
  • The loop used to access the second dimension of the fused tensor is kept as is.
  • The loop used to access the third dimension of the fused tensor is split into three.

i.e. (e0, e1, e2, e3, e4) is the domain of the indexing map of the modified op, then

d0 -> e0, e1 d1 -> e2, e3, e4 d2 -> e5

substituting this, the generic op can be rewritten as

d = linalg.generic ins(%0, %1 : ) indexing_maps = [affine_map<(e0, e1, e2, e3, e4, e5) -> (e2, e3, e4, e0, e1, e5)>, affine_map<(e0, e1, e2, e3, e4, e5) -> (e2, e3, e4, e5)>, affine_map<(e0, e1, e2, e3, e4, e5) -> (e0, e1, e5, e2, e3, e4)>]

Since operands to the linalg generic are now 5D, reshapes can be introduced to make it consistent

%0 = tensor.expand_shape a [[0, 1, 2], [3, 4], [5]] : tensor<?x?x?xf32> into tensor<?x?x?x?x?x?xf32> %1 = tensor.expand_shape b [[0, 1, 2], [3]] : tensor<?x?x?xf32> into tensor<?x?x?x?xf32>

The added reshapes are again expanding patterns, so they will get fused with its producers if possible.

Definition at line 453 of file ElementwiseOpFusion.cpp.

References mlir::detail::enumerate(), mlir::failure(), mlir::AffineMap::getDimPosition(), mlir::AffineMap::getNumDims(), mlir::AffineMap::getNumResults(), mlir::getParallelIteratorTypeName(), mlir::AffineMap::getResults(), and mlir::success().

Referenced by fuseWithReshapeByExpansion().

◆ isGenericOpExpandable()

static LogicalResult isGenericOpExpandable ( GenericOp  genericOp,
const ExpansionInfo &  expansionInfo,
PatternRewriter rewriter 
)
static

Epanding the body of a linalg operation requires adaptations of the accessed loop indices.

Specifically, access of indices in the original operation need to be replaced with linearizations of indices in the expanded op. That requires the shape of the expanded dimensions to be static (at least all but the most significant). For now check that these are all statically sized. Note that this could be extended to handle dynamic case, but the implementation below uses affine.apply which seems to have issues when the shapes are not static.

Definition at line 563 of file ElementwiseOpFusion.cpp.

References mlir::RewriterBase::notifyMatchFailure(), and mlir::success().

Referenced by fuseWithReshapeByExpansion().

◆ updateExpandedGenericOpRegion()

static void updateExpandedGenericOpRegion ( PatternRewriter rewriter,
Location  loc,
Region fusedRegion,
const ExpansionInfo &  expansionInfo 
)
static

Update the body of an expanded linalg operation having index semantics.

The indices of the original operation need to be recovered by linearizing the indices of the correspoding dimensions of the expanded operation. For now it is assumed that the shapes of the expanded operation needed for linearization are static.

Definition at line 642 of file ElementwiseOpFusion.cpp.

References mlir::bindDims(), mlir::OpBuilder::create(), mlir::Region::front(), mlir::Builder::getContext(), mlir::Block::getOps(), mlir::RewriterBase::replaceOp(), and mlir::OpBuilder::setInsertionPointAfter().

Referenced by fuseWithReshapeByExpansion().