#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 "mlir/Transforms/WalkPatternRewriteDriver.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 UnitExtentReplacementInfo	dropUnitExtentFromOperandMetadata (MLIRContext context, IndexingMapOpInterface op, OpOperand opOperand, llvm::SmallDenseMap< unsigned, unsigned > &oldDimsToNewDimsMap, ArrayRef< AffineExpr > dimReplacements)

Macro Definition Documentation

◆ DEBUG_TYPE

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

Definition at line 39 of file DropUnitDims.cpp.

◆ GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMSPASS

#define GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMSPASS

Definition at line 35 of file DropUnitDims.cpp.

Function Documentation

◆ dropUnitExtentFromOperandMetadata()

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

static

Definition at line 345 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().

◆ 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 228 of file DropUnitDims.cpp.

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

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

Classes

Namespaces