MLIR
20.0.0git
|
#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/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/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) |
#define DEBUG_TYPE "linalg-drop-unit-dims" |
Definition at line 41 of file DropUnitDims.cpp.
#define GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMSPASS |
Definition at line 37 of file DropUnitDims.cpp.
|
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 280 of file DropUnitDims.cpp.
References mlir::OpBuilder::create(), mlir::linalg::ControlDropUnitDims::ExtractInsertSlice, mlir::get(), mlir::Value::getType(), and mlir::linalg::ControlDropUnitDims::ReassociativeReshape.
Referenced by mlir::linalg::dropUnitDims().
|
static |
Definition at line 336 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().
|
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 253 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().
|
static |
Patterns that are used to canonicalize the use of unit-extent dims for broadcasting.
Definition at line 764 of file DropUnitDims.cpp.
References mlir::RewritePatternSet::add(), mlir::RewritePatternSet::getContext(), options, mlir::tensor::populateFoldTensorEmptyPatterns(), mlir::memref::populateResolveRankedShapedTypeResultDimsPatterns(), and mlir::memref::populateResolveShapedTypeResultDimsPatterns().
Referenced by mlir::linalg::populateFoldUnitExtentDimsPatterns().
|
static |
Definition at line 784 of file DropUnitDims.cpp.
References mlir::RewritePatternSet::add(), mlir::RewritePatternSet::getContext(), options, mlir::tensor::populateFoldTensorEmptyPatterns(), mlir::memref::populateResolveRankedShapedTypeResultDimsPatterns(), and mlir::memref::populateResolveShapedTypeResultDimsPatterns().
Referenced by mlir::linalg::populateFoldUnitExtentDimsPatterns().
|
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 = { 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 229 of file DropUnitDims.cpp.
References mlir::RewriterBase::replaceOpWithNewOp(), and mlir::OpBuilder::setInsertionPoint().
Referenced by mlir::linalg::dropUnitDims().