MLIR  19.0.0git
Classes | Namespaces | Macros | Functions
DropUnitDims.cpp File Reference
#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/Tensor/Utils/Utils.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/ADT/SetVector.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "mlir/Dialect/Linalg/Passes.h.inc"

Go to the source code of this file.

Classes

struct  UnitExtentReplacementInfo
 Compute the modified metadata for an operands of operation whose unit dims are being dropped. More...
 

Namespaces

 mlir
 Include the generated interface declarations.
 

Macros

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

Functions

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. More...
 
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. More...
 
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. More...
 
static UnitExtentReplacementInfo dropUnitExtentFromOperandMetadata (MLIRContext *context, GenericOp genericOp, 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. More...
 
static void populateFoldUnitExtentDimsViaSlicesPatterns (RewritePatternSet &patterns, ControlDropUnitDims &options)
 

Macro Definition Documentation

◆ DEBUG_TYPE

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

Definition at line 40 of file DropUnitDims.cpp.

◆ GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMSPASS

#define GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMSPASS

Definition at line 36 of file DropUnitDims.cpp.

Function Documentation

◆ collapseValue()

static 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 282 of file DropUnitDims.cpp.

References mlir::OpBuilder::create(), mlir::linalg::ControlDropUnitDims::ExtractInsertSlice, mlir::get(), mlir::Value::getType(), mlir::linalg::ControlDropUnitDims::ReassociativeReshape, and mlir::succeeded().

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

◆ dropUnitExtentFromOperandMetadata()

static UnitExtentReplacementInfo dropUnitExtentFromOperandMetadata ( MLIRContext context,
GenericOp  genericOp,
OpOperand opOperand,
llvm::SmallDenseMap< unsigned, unsigned > &  oldDimsToNewDimsMap,
ArrayRef< AffineExpr dimReplacements 
)
static

◆ expandValue()

static 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 256 of file DropUnitDims.cpp.

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

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

◆ populateFoldUnitExtentDimsViaReshapesPatterns()

static void populateFoldUnitExtentDimsViaReshapesPatterns ( RewritePatternSet patterns,
ControlDropUnitDims options 
)
static

◆ populateFoldUnitExtentDimsViaSlicesPatterns()

static void populateFoldUnitExtentDimsViaSlicesPatterns ( RewritePatternSet patterns,
ControlDropUnitDims options 
)
static

◆ replaceUnitDimIndexOps()

static 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,

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

#trait = { args_in = 2, args_out = 1, 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 = { args_in = 2, args_out = 1, 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 232 of file DropUnitDims.cpp.

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

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