MLIR
17.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.
Namespaces | |
mlir | |
This header declares functions that assit transformations in the MemRef dialect. | |
Macros | |
#define | GEN_PASS_DEF_LINALGFOLDUNITEXTENTDIMS |
#define | DEBUG_TYPE "linalg-drop-unit-dims" |
Functions | |
static ArrayAttr | replaceUnitDims (DenseSet< unsigned > &unitDims, ArrayRef< AffineMap > indexingMaps, MLIRContext *context) |
Implements a pass that canonicalizes the uses of unit-extent dimensions for broadcasting. More... | |
static void | replaceUnitDimIndexOps (GenericOp genericOp, const DenseSet< unsigned > &unitDims, PatternRewriter &rewriter) |
Update the index accesses of linalg operations having index semantics. More... | |
static std::optional< UnitExtentReplacementInfo > | replaceUnitExtents (GenericOp genericOp, OpOperand *opOperand, MLIRContext *context) |
Utility function for replacing operands/results to a linalg generic operation with unit-extent dimensions. More... | |
#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 |
Update the index accesses of linalg operations having index semantics.
Definition at line 162 of file DropUnitDims.cpp.
References mlir::RewriterBase::replaceOpWithNewOp(), and mlir::OpBuilder::setInsertionPoint().
|
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> } Given dims of the iteration space of a structured op that are known to be single trip count (unitDims
), return the indexing maps to use in the canonicalized op with these dims removed, given the original indexingMaps
.
Definition at line 113 of file DropUnitDims.cpp.
References mlir::concatAffineMaps(), mlir::get(), mlir::getAffineConstantExpr(), mlir::getAffineDimExpr(), mlir::getAffineSymbolExpr(), mlir::inversePermutation(), and mlir::simplifyAffineMap().
|
static |
Utility function for replacing operands/results to a linalg generic operation with unit-extent dimensions.
These can be replaced with an operand/result with the unit-extent dimension removed. This is only done if the indexing map used to access that dimension has a AffineConstantExpr of value 0. Given the type
of an result/operand of a Linalg op, and its indexMap
the utility function returns:
Definition at line 373 of file DropUnitDims.cpp.
References mlir::IROperand< DerivedT, IRValueT >::get(), mlir::AffineMap::get(), mlir::getAffineConstantExpr(), mlir::AffineMap::getNumDims(), mlir::AffineMap::getNumSymbols(), mlir::AffineMap::getResults(), and mlir::Value::getType().