|
MLIR
18.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/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_LINALGFOLDUNITEXTENTDIMS |
| #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 40 of file DropUnitDims.cpp.
| #define GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMS |
Definition at line 36 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 284 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().
|
static |
Definition at line 340 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 258 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 636 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 655 of file DropUnitDims.cpp.
References mlir::RewritePatternSet::add(), mlir::linalg::ControlDropUnitDims::ExtractInsertSlice, 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 = { 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 234 of file DropUnitDims.cpp.
References mlir::RewriterBase::replaceOpWithNewOp(), and mlir::OpBuilder::setInsertionPoint().
Referenced by mlir::linalg::dropUnitDims().
1.9.1