#include "mlir/Dialect/Linalg/Passes.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/Transforms/FoldUtils.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/Support/Debug.h"
#include "mlir/Dialect/Linalg/Passes.h.inc"

Classes
class	mlir::impl::LinalgFoldUnitExtentDimsPassBase< DerivedT >
struct	UnitExtentReplacementInfo
	Compute the modified metadata for an operands of operation whose unit dims are being dropped. More...

Namespaces
namespace	mlir
	Include the generated interface declarations.
namespace	mlir::impl
	Attribute collections provide a dictionary-like interface.

Macros
#define	GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMSPASS
#define	DEBUG_TYPE "linalg-drop-unit-dims"

Functions
std::unique_ptr<::mlir::Pass >	mlir::impl::createLinalgFoldUnitExtentDimsPass ()
std::unique_ptr<::mlir::Pass >	mlir::impl::createLinalgFoldUnitExtentDimsPass (LinalgFoldUnitExtentDimsPassOptions options)
std::unique_ptr<::mlir::Pass >	mlir::createLinalgFoldUnitExtentDimsPass ()
std::unique_ptr<::mlir::Pass >	mlir::createLinalgFoldUnitExtentDimsPass (LinalgFoldUnitExtentDimsPassOptions options)
static void	replaceUnitDimIndexOps (GenericOp genericOp, const llvm::SmallDenseSet< unsigned > &unitDims, RewriterBase &rewriter)
	Implements a pass that canonicalizes the uses of unit-extent dimensions for broadcasting.
static Value	expandValue (RewriterBase &rewriter, Location loc, Value result, Value origDest, ArrayRef< ReassociationIndices > reassociation, ControlDropUnitDims::RankReductionStrategy rankReductionStrategy)
	Expand the given value so that the type matches the type of origDest.
static Value	collapseValue (RewriterBase &rewriter, Location loc, Value operand, ArrayRef< int64_t > targetShape, ArrayRef< ReassociationIndices > reassociation, ControlDropUnitDims::RankReductionStrategy rankReductionStrategy)
	Collapse the given value so that the type matches the type of origOutput.
static UnitExtentReplacementInfo	dropUnitExtentFromOperandMetadata (MLIRContext context, IndexingMapOpInterface op, OpOperand opOperand, llvm::SmallDenseMap< unsigned, unsigned > &oldDimsToNewDimsMap, ArrayRef< AffineExpr > dimReplacements)
static void	populateFoldUnitExtentDimsViaReshapesPatterns (RewritePatternSet &patterns, ControlDropUnitDims &options)
	Patterns that are used to canonicalize the use of unit-extent dims for broadcasting.
static void	populateFoldUnitExtentDimsViaSlicesPatterns (RewritePatternSet &patterns, ControlDropUnitDims &options)

Macro Definition Documentation

◆ DEBUG_TYPE

#define DEBUG_TYPE "linalg-drop-unit-dims"

Definition at line 38 of file DropUnitDims.cpp.

◆ GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMSPASS

#define GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMSPASS

Definition at line 34 of file DropUnitDims.cpp.

Function Documentation

◆ collapseValue()

Value collapseValue	(	RewriterBase &	rewriter,
		Location	loc,
		Value	operand,
		ArrayRef< int64_t >	targetShape,
		ArrayRef< ReassociationIndices >	reassociation,
		ControlDropUnitDims::RankReductionStrategy	rankReductionStrategy )

static

Collapse the given value so that the type matches the type of origOutput.

The reassociation is used when rankReductionStrategy is set to RankReductionStrategy::ReassociativeReshape.

Definition at line 278 of file DropUnitDims.cpp.

References mlir::linalg::ControlDropUnitDims::ExtractInsertSlice, mlir::Value::getType(), and mlir::linalg::ControlDropUnitDims::ReassociativeReshape.

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

◆ dropUnitExtentFromOperandMetadata()

UnitExtentReplacementInfo dropUnitExtentFromOperandMetadata	(	MLIRContext *	context,
		IndexingMapOpInterface	op,
		OpOperand *	opOperand,
		llvm::SmallDenseMap< unsigned, unsigned > &	oldDimsToNewDimsMap,
		ArrayRef< AffineExpr >	dimReplacements )

static

Definition at line 334 of file DropUnitDims.cpp.

References mlir::AffineMap::get(), mlir::AffineMap::getNumSymbols(), mlir::AffineMap::getResults(), UnitExtentReplacementInfo::indexMap, UnitExtentReplacementInfo::reassociation, and UnitExtentReplacementInfo::targetShape.

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

◆ expandValue()

Value expandValue	(	RewriterBase &	rewriter,
		Location	loc,
		Value	result,
		Value	origDest,
		ArrayRef< ReassociationIndices >	reassociation,
		ControlDropUnitDims::RankReductionStrategy	rankReductionStrategy )

static

Expand the given value so that the type matches the type of origDest.

The reassociation is used when rankReductionStrategy is set to RankReductionStrategy::ReassociativeReshape.

Definition at line 251 of file DropUnitDims.cpp.

References mlir::OpBuilder::createOrFold(), mlir::linalg::ControlDropUnitDims::ExtractInsertSlice, mlir::Builder::getIndexAttr(), mlir::tensor::getMixedSizes(), mlir::Value::getType(), mlir::linalg::ControlDropUnitDims::ReassociativeReshape, and result.

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

◆ populateFoldUnitExtentDimsViaReshapesPatterns()

void populateFoldUnitExtentDimsViaReshapesPatterns	(	RewritePatternSet &	patterns,
		ControlDropUnitDims &	options )

static

Patterns that are used to canonicalize the use of unit-extent dims for broadcasting.

Definition at line 804 of file DropUnitDims.cpp.

References options, mlir::patterns, mlir::tensor::populateFoldTensorEmptyPatterns(), mlir::memref::populateResolveRankedShapedTypeResultDimsPatterns(), and mlir::memref::populateResolveShapedTypeResultDimsPatterns().

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

◆ populateFoldUnitExtentDimsViaSlicesPatterns()

void populateFoldUnitExtentDimsViaSlicesPatterns	(	RewritePatternSet &	patterns,
		ControlDropUnitDims &	options )

static

Definition at line 824 of file DropUnitDims.cpp.

References options, mlir::patterns, mlir::tensor::populateFoldTensorEmptyPatterns(), mlir::memref::populateResolveRankedShapedTypeResultDimsPatterns(), and mlir::memref::populateResolveShapedTypeResultDimsPatterns().

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

◆ replaceUnitDimIndexOps()

void replaceUnitDimIndexOps	(	GenericOp	genericOp,
		const llvm::SmallDenseSet< unsigned > &	unitDims,
		RewriterBase &	rewriter )

static

Implements a pass that canonicalizes the uses of unit-extent dimensions for broadcasting.

For example,

#accesses = [
  affine_map<(d0, d1) -> (0, d1)>,
  affine_map<(d0, d1) -> (d0, 0)>,
  affine_map<(d0, d1) -> (d0, d1)>
]
 
#trait = {
  indexing_maps = #accesses,
  iterator_types = ["parallel", "parallel"],
  library_call = "some_external_fn"
}
 
func @broadcast_test(%arg0 : tensor<5xf32>, %arg1 : tensor<5xf32>) ->
tensor<5x5xf32>
{
  %0 = linalg.tensor_reshape %arg0 [affine_map<(d0, d1) -> (d0, d1)>] :
       tensor<5xf32> into tensor<1x5xf32>
  %1 = linalg.tensor_reshape %arg1 [affine_map<(d0, d1) -> (d0, d1)>] :
       tensor<5xf32> into tensor<5x1xf32>
  %2 = linalg.generic #trait %0, %1 {
       ^bb0(%arg2: f32, %arg3: f32):
         %3 = arith.addf %arg2, %arg3 : f32
         linalg.yield %3 : f32
       } : tensor<1x5xf32>, tensor<5x1xf32> -> tensor<5x5xf32>
  return %2 : tensor<5x5xf32>
}
 
would canonicalize to

mlir #accesses = [ affine_map<(d0, d1) -> (d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0, d1)> ]

#trait = { indexing_maps = #accesses, iterator_types = ["parallel", "parallel"], library_call = "some_external_fn" }

func @broadcast_test(arg0 : tensor<5xf32>, arg1 : tensor<5xf32>) -> tensor<5x5xf32> { %0 = linalg.generic #trait arg0, arg1 { ^bb0(arg2: f32, arg3: f32): %3 = arith.addf arg2, arg3 : f32 linalg.yield %3 : f32 } : tensor<5xf32>, tensor<5xf32> -> tensor<5x5xf32> return %0 : tensor<5x5xf32> } Update the index accesses of linalg operations having index semantics.

Definition at line 227 of file DropUnitDims.cpp.

References mlir::RewriterBase::replaceOpWithNewOp(), and mlir::OpBuilder::setInsertionPoint().

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

Classes

Namespaces

Macros

Functions

Macro Definition Documentation

◆ DEBUG_TYPE

◆ GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMSPASS

Function Documentation

◆ collapseValue()

◆ dropUnitExtentFromOperandMetadata()

◆ expandValue()

◆ populateFoldUnitExtentDimsViaReshapesPatterns()

◆ populateFoldUnitExtentDimsViaSlicesPatterns()

◆ replaceUnitDimIndexOps()